From 342d76f653fe33cdec7c39e6bd7dc95c091d3bbc Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 15 Oct 2026 08:42:15 +0200 Subject: [PATCH 1/9] Rewrite data, model and training core around one default path - data: one input config for CSV/h5ad/SCE, labels attached at load time, splits drawn at setup (stratified, patient-grouped), data.unit (region|tile|polygon), one ego-graph sampler subclassing PyG's ShaDow sampler, fingerprinted processed graphs, fixed PCA serialisation - models: GINE encoder, projected gated attention, logit-space pooling; one manual-optimisation training path with gradient clipping after AMP unscaling, probe head fitted in setup, Cox/regression objectives over a hyperbatch, modules rebuilt from their checkpoints alone - configs: model/task/regime groups; finetune is the default regime; regimes renamed (finetune, finetune_reg, attention_only, mlp_head, mlp_head_frozen, pretrain) - CLIs: train, nested-cv (refit without validation), extract, predict Tests, niche analysis and docs follow in later commits. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NqExFvaq8RhDYBj6RShCMk --- src/grass_mil/_paths.py | 29 - src/grass_mil/batching.py | 166 -- src/grass_mil/callbacks.py | 110 -- src/grass_mil/configs/__init__.py | 2 +- .../configs/augmentation/bgrl_paper.yaml | 17 - .../configs/augmentation/uniform.yaml | 6 - src/grass_mil/configs/callbacks/default.yaml | 22 - .../configs/callbacks/early_stopping.yaml | 17 +- .../configs/callbacks/mil_selection.yaml | 27 - .../configs/callbacks/model_checkpoint.yaml | 23 +- .../configs/callbacks/model_summary.yaml | 4 +- src/grass_mil/configs/callbacks/none.yaml | 0 .../callbacks/pretrain_checkpoint.yaml | 10 +- src/grass_mil/configs/callbacks/refit.yaml | 18 - .../configs/callbacks/rich_progress_bar.yaml | 2 - .../configs/callbacks/selection.yaml | 26 + .../configs/data/sampler/ego_graph.yaml | 18 + .../configs/data/sampler/ego_graph_mil.yaml | 36 - .../data/sampler/ego_graph_pretrain.yaml | 40 +- .../configs/data/sampler/identity.yaml | 2 - .../configs/data/sampler/shadow_custom.yaml | 33 - .../configs/data/sampler/shadow_native.yaml | 33 - src/grass_mil/configs/data/spatial_omics.yaml | 151 +- src/grass_mil/configs/debug/default.yaml | 35 - src/grass_mil/configs/debug/fdr.yaml | 9 - src/grass_mil/configs/debug/limit.yaml | 12 - src/grass_mil/configs/debug/overfit.yaml | 13 - src/grass_mil/configs/debug/profiler.yaml | 12 - src/grass_mil/configs/eval.yaml | 29 - .../configs/experiment/example_cohort.yaml | 66 +- .../experiment/grass_mil_pretrain_a3.yaml | 21 - src/grass_mil/configs/extract.yaml | 31 +- src/grass_mil/configs/extras/default.yaml | 8 - src/grass_mil/configs/hydra/default.yaml | 5 - .../inference/aggregation/default.yaml | 15 - .../configs/inference/embeddings/default.yaml | 5 - .../inference/interpretability/default.yaml | 40 - .../configs/inference/metrics/default.yaml | 5 - src/grass_mil/configs/inference/predict.yaml | 42 - .../clustering/agglomerative.yaml | 6 - .../interpretability/clustering/dbscan.yaml | 10 - .../interpretability/clustering/hdbscan.yaml | 18 - .../interpretability/clustering/kmeans.yaml | 9 - .../interpretability/clustering/optics.yaml | 13 - .../interpretability/clustering/spectral.yaml | 8 - .../interpretability/plugins/default.yaml | 72 - .../interpretability/plugins/full.yaml | 24 - .../interpretability/reduction/pca.yaml | 5 - .../interpretability/reduction/tsne.yaml | 8 - .../interpretability/reduction/umap.yaml | 8 - .../configs/interpretability/report.yaml | 41 - src/grass_mil/configs/local/.gitkeep | 0 src/grass_mil/configs/logger/aim.yaml | 28 - src/grass_mil/configs/logger/comet.yaml | 12 - src/grass_mil/configs/logger/csv.yaml | 9 +- .../configs/logger/many_loggers.yaml | 9 - src/grass_mil/configs/logger/mlflow.yaml | 12 - src/grass_mil/configs/logger/neptune.yaml | 9 - src/grass_mil/configs/logger/tensorboard.yaml | 10 +- src/grass_mil/configs/logger/wandb.yaml | 16 - src/grass_mil/configs/loocv.yaml | 14 - .../configs/model/attention/gated.yaml | 7 - .../model/attention/gated_projected.yaml | 10 - src/grass_mil/configs/model/bgrl.yaml | 15 + src/grass_mil/configs/model/bgrl_a3.yaml | 28 - src/grass_mil/configs/model/bgrl_module.yaml | 22 - src/grass_mil/configs/model/encoder/gat.yaml | 18 - src/grass_mil/configs/model/encoder/gcn.yaml | 18 - src/grass_mil/configs/model/encoder/gin.yaml | 20 - src/grass_mil/configs/model/encoder/gine.yaml | 36 +- .../configs/model/encoder/gine_length.yaml | 45 - .../configs/model/encoder/graphsage.yaml | 18 - src/grass_mil/configs/model/grass_mil.yaml | 54 +- .../configs/model/supervised_module.yaml | 40 - src/grass_mil/configs/nested_cv.yaml | 57 +- src/grass_mil/configs/optim/adamw.yaml | 4 +- src/grass_mil/configs/paths/default.yaml | 24 +- src/grass_mil/configs/predict.yaml | 21 + src/grass_mil/configs/probe.yaml | 29 - .../configs/regime/attention_only.yaml | 16 + src/grass_mil/configs/regime/finetune.yaml | 39 + .../configs/regime/finetune_reg.yaml | 19 + .../configs/regime/linear_ftall.yaml | 28 - .../configs/regime/linear_ftall_reg.yaml | 21 - .../configs/regime/linear_probehead.yaml | 27 - src/grass_mil/configs/regime/mil.yaml | 65 - src/grass_mil/configs/regime/mil_frozen.yaml | 17 - src/grass_mil/configs/regime/mlp_head.yaml | 15 + .../configs/regime/mlp_head_frozen.yaml | 13 + src/grass_mil/configs/regime/pretrain.yaml | 39 + src/grass_mil/configs/regime/pretrain_a3.yaml | 22 - src/grass_mil/configs/scheduler/cosine.yaml | 4 + .../configs/scheduler/cosine_epoch.yaml | 3 - .../configs/scheduler/cosine_mil_step.yaml | 6 - .../configs/scheduler/cosine_step.yaml | 4 - src/grass_mil/configs/scores.yaml | 29 - .../configs/task/classification.yaml | 13 + src/grass_mil/configs/task/finetune_mean.yaml | 16 - src/grass_mil/configs/task/finetune_mil.yaml | 74 - .../configs/task/finetune_survival.yaml | 43 - src/grass_mil/configs/task/pretrain_bgrl.yaml | 30 - .../configs/task/pretrain_bgrl_a3.yaml | 53 - src/grass_mil/configs/task/regression.yaml | 9 + src/grass_mil/configs/task/survival.yaml | 8 + src/grass_mil/configs/train.yaml | 71 +- src/grass_mil/configs/trainer/cpu.yaml | 3 - src/grass_mil/configs/trainer/ddp.yaml | 9 - src/grass_mil/configs/trainer/ddp_sim.yaml | 7 - src/grass_mil/configs/trainer/default.yaml | 20 +- src/grass_mil/configs/trainer/gpu.yaml | 2 + src/grass_mil/configs/trainer/gpu_amp.yaml | 6 - src/grass_mil/configs/trainer/mps.yaml | 5 - src/grass_mil/contracts.py | 138 -- src/grass_mil/data/__init__.py | 15 +- src/grass_mil/data/components/__init__.py | 43 - src/grass_mil/data/components/datasets.py | 86 - src/grass_mil/data/components/ego_radius.py | 180 -- .../data/components/feature_reducers.py | 101 -- .../data/components/graph_builders.py | 189 --- src/grass_mil/data/components/loaders.py | 382 ----- src/grass_mil/data/components/merge.py | 79 - src/grass_mil/data/components/patching.py | 91 -- src/grass_mil/data/components/precompute.py | 1181 -------------- src/grass_mil/data/components/samplers.py | 1060 ------------ .../data/components/seed_sampling.py | 277 ---- .../data/components/spatial_types.py | 19 - src/grass_mil/data/components/splits.py | 197 --- src/grass_mil/data/components/transforms.py | 78 - src/grass_mil/data/datamodule.py | 309 ++++ src/grass_mil/data/ego_radius.py | 95 ++ src/grass_mil/data/graphs.py | 89 + src/grass_mil/data/interior.py | 152 ++ src/grass_mil/data/labels.py | 109 ++ src/grass_mil/data/loaders.py | 243 +++ src/grass_mil/data/precompute.py | 301 ++++ src/grass_mil/data/reducers.py | 83 + src/grass_mil/data/samplers.py | 291 ++++ .../data/spatial_omics_datamodule.py | 355 ---- src/grass_mil/data/splits.py | 125 ++ src/grass_mil/data/units.py | 64 + src/grass_mil/eval.py | 73 - src/grass_mil/extract.py | 159 +- src/grass_mil/inference/__init__.py | 18 +- src/grass_mil/inference/aggregation.py | 427 ----- src/grass_mil/inference/cell_cache.py | 339 ++-- src/grass_mil/inference/collectors.py | 483 ------ .../inference/interpretability_export.py | 638 -------- src/grass_mil/inference/io.py | 107 -- src/grass_mil/inference/metrics.py | 273 ---- src/grass_mil/inference/model_loading.py | 69 - src/grass_mil/inference/predict.py | 529 ------ src/grass_mil/inference/predictions.py | 92 ++ src/grass_mil/inference/schemas.py | 72 - src/grass_mil/loocv.py | 285 ---- src/grass_mil/models/__init__.py | 54 +- src/grass_mil/models/augmentations.py | 183 +++ src/grass_mil/models/bgrl_module.py | 266 --- src/grass_mil/models/checkpoints.py | 26 + src/grass_mil/models/components/__init__.py | 45 - src/grass_mil/models/components/attention.py | 134 -- src/grass_mil/models/components/backbones.py | 267 --- src/grass_mil/models/components/embeddings.py | 280 ---- src/grass_mil/models/components/factory.py | 105 -- src/grass_mil/models/components/heads.py | 34 - src/grass_mil/models/components/losses.py | 145 -- src/grass_mil/models/components/pooling.py | 44 - src/grass_mil/models/components/ssl.py | 67 - src/grass_mil/models/encoder.py | 177 ++ src/grass_mil/models/mil.py | 143 ++ src/grass_mil/models/objectives.py | 162 ++ src/grass_mil/models/optimization.py | 54 + src/grass_mil/models/pretraining.py | 172 ++ src/grass_mil/models/probe.py | 91 ++ src/grass_mil/models/supervised.py | 334 ++++ src/grass_mil/models/supervised_module.py | 1448 ----------------- src/grass_mil/models/training/__init__.py | 80 - .../models/training/augmentations.py | 381 ----- src/grass_mil/models/training/bagging.py | 272 ---- src/grass_mil/models/training/builders.py | 212 --- .../models/training/checkpoint_init.py | 136 -- .../models/training/class_weights.py | 67 - src/grass_mil/models/training/loss_utils.py | 51 - src/grass_mil/models/training/optimization.py | 168 -- src/grass_mil/models/training/probe.py | 275 ---- src/grass_mil/models/training/ssl_runtime.py | 22 - src/grass_mil/nested_cv.py | 619 ++----- src/grass_mil/predict.py | 110 ++ src/grass_mil/probe.py | 133 -- src/grass_mil/runtime.py | 82 + src/grass_mil/train.py | 229 +-- src/grass_mil/utils/__init__.py | 19 - src/grass_mil/utils/instantiators.py | 56 - src/grass_mil/utils/logging_utils.py | 57 - src/grass_mil/utils/pylogger.py | 51 - src/grass_mil/utils/rich_utils.py | 101 -- src/grass_mil/utils/utils.py | 119 -- 196 files changed, 4353 insertions(+), 15697 deletions(-) delete mode 100644 src/grass_mil/_paths.py delete mode 100644 src/grass_mil/batching.py delete mode 100644 src/grass_mil/callbacks.py delete mode 100644 src/grass_mil/configs/augmentation/bgrl_paper.yaml delete mode 100644 src/grass_mil/configs/augmentation/uniform.yaml delete mode 100644 src/grass_mil/configs/callbacks/default.yaml delete mode 100644 src/grass_mil/configs/callbacks/mil_selection.yaml delete mode 100644 src/grass_mil/configs/callbacks/none.yaml delete mode 100644 src/grass_mil/configs/callbacks/refit.yaml create mode 100644 src/grass_mil/configs/callbacks/selection.yaml create mode 100644 src/grass_mil/configs/data/sampler/ego_graph.yaml delete mode 100644 src/grass_mil/configs/data/sampler/ego_graph_mil.yaml delete mode 100644 src/grass_mil/configs/data/sampler/identity.yaml delete mode 100644 src/grass_mil/configs/data/sampler/shadow_custom.yaml delete mode 100644 src/grass_mil/configs/data/sampler/shadow_native.yaml delete mode 100644 src/grass_mil/configs/debug/default.yaml delete mode 100644 src/grass_mil/configs/debug/fdr.yaml delete mode 100644 src/grass_mil/configs/debug/limit.yaml delete mode 100644 src/grass_mil/configs/debug/overfit.yaml delete mode 100644 src/grass_mil/configs/debug/profiler.yaml delete mode 100644 src/grass_mil/configs/eval.yaml delete mode 100644 src/grass_mil/configs/experiment/grass_mil_pretrain_a3.yaml delete mode 100644 src/grass_mil/configs/extras/default.yaml delete mode 100644 src/grass_mil/configs/inference/aggregation/default.yaml delete mode 100644 src/grass_mil/configs/inference/embeddings/default.yaml delete mode 100644 src/grass_mil/configs/inference/interpretability/default.yaml delete mode 100644 src/grass_mil/configs/inference/metrics/default.yaml delete mode 100644 src/grass_mil/configs/inference/predict.yaml delete mode 100644 src/grass_mil/configs/interpretability/clustering/agglomerative.yaml delete mode 100644 src/grass_mil/configs/interpretability/clustering/dbscan.yaml delete mode 100644 src/grass_mil/configs/interpretability/clustering/hdbscan.yaml delete mode 100644 src/grass_mil/configs/interpretability/clustering/kmeans.yaml delete mode 100644 src/grass_mil/configs/interpretability/clustering/optics.yaml delete mode 100644 src/grass_mil/configs/interpretability/clustering/spectral.yaml delete mode 100644 src/grass_mil/configs/interpretability/plugins/default.yaml delete mode 100644 src/grass_mil/configs/interpretability/plugins/full.yaml delete mode 100644 src/grass_mil/configs/interpretability/reduction/pca.yaml delete mode 100644 src/grass_mil/configs/interpretability/reduction/tsne.yaml delete mode 100644 src/grass_mil/configs/interpretability/reduction/umap.yaml delete mode 100644 src/grass_mil/configs/interpretability/report.yaml delete mode 100644 src/grass_mil/configs/local/.gitkeep delete mode 100644 src/grass_mil/configs/logger/aim.yaml delete mode 100644 src/grass_mil/configs/logger/comet.yaml delete mode 100644 src/grass_mil/configs/logger/many_loggers.yaml delete mode 100644 src/grass_mil/configs/logger/mlflow.yaml delete mode 100644 src/grass_mil/configs/logger/neptune.yaml delete mode 100644 src/grass_mil/configs/logger/wandb.yaml delete mode 100644 src/grass_mil/configs/loocv.yaml delete mode 100644 src/grass_mil/configs/model/attention/gated.yaml delete mode 100644 src/grass_mil/configs/model/attention/gated_projected.yaml create mode 100644 src/grass_mil/configs/model/bgrl.yaml delete mode 100644 src/grass_mil/configs/model/bgrl_a3.yaml delete mode 100644 src/grass_mil/configs/model/bgrl_module.yaml delete mode 100644 src/grass_mil/configs/model/encoder/gat.yaml delete mode 100644 src/grass_mil/configs/model/encoder/gcn.yaml delete mode 100644 src/grass_mil/configs/model/encoder/gin.yaml delete mode 100644 src/grass_mil/configs/model/encoder/gine_length.yaml delete mode 100644 src/grass_mil/configs/model/encoder/graphsage.yaml delete mode 100644 src/grass_mil/configs/model/supervised_module.yaml create mode 100644 src/grass_mil/configs/predict.yaml delete mode 100644 src/grass_mil/configs/probe.yaml create mode 100644 src/grass_mil/configs/regime/attention_only.yaml create mode 100644 src/grass_mil/configs/regime/finetune.yaml create mode 100644 src/grass_mil/configs/regime/finetune_reg.yaml delete mode 100644 src/grass_mil/configs/regime/linear_ftall.yaml delete mode 100644 src/grass_mil/configs/regime/linear_ftall_reg.yaml delete mode 100644 src/grass_mil/configs/regime/linear_probehead.yaml delete mode 100644 src/grass_mil/configs/regime/mil.yaml delete mode 100644 src/grass_mil/configs/regime/mil_frozen.yaml create mode 100644 src/grass_mil/configs/regime/mlp_head.yaml create mode 100644 src/grass_mil/configs/regime/mlp_head_frozen.yaml create mode 100644 src/grass_mil/configs/regime/pretrain.yaml delete mode 100644 src/grass_mil/configs/regime/pretrain_a3.yaml create mode 100644 src/grass_mil/configs/scheduler/cosine.yaml delete mode 100644 src/grass_mil/configs/scheduler/cosine_epoch.yaml delete mode 100644 src/grass_mil/configs/scheduler/cosine_mil_step.yaml delete mode 100644 src/grass_mil/configs/scheduler/cosine_step.yaml delete mode 100644 src/grass_mil/configs/scores.yaml create mode 100644 src/grass_mil/configs/task/classification.yaml delete mode 100644 src/grass_mil/configs/task/finetune_mean.yaml delete mode 100644 src/grass_mil/configs/task/finetune_mil.yaml delete mode 100644 src/grass_mil/configs/task/finetune_survival.yaml delete mode 100644 src/grass_mil/configs/task/pretrain_bgrl.yaml delete mode 100644 src/grass_mil/configs/task/pretrain_bgrl_a3.yaml create mode 100644 src/grass_mil/configs/task/regression.yaml create mode 100644 src/grass_mil/configs/task/survival.yaml delete mode 100644 src/grass_mil/configs/trainer/ddp.yaml delete mode 100644 src/grass_mil/configs/trainer/ddp_sim.yaml delete mode 100644 src/grass_mil/configs/trainer/gpu_amp.yaml delete mode 100644 src/grass_mil/configs/trainer/mps.yaml delete mode 100644 src/grass_mil/contracts.py delete mode 100644 src/grass_mil/data/components/__init__.py delete mode 100644 src/grass_mil/data/components/datasets.py delete mode 100644 src/grass_mil/data/components/ego_radius.py delete mode 100644 src/grass_mil/data/components/feature_reducers.py delete mode 100644 src/grass_mil/data/components/graph_builders.py delete mode 100644 src/grass_mil/data/components/loaders.py delete mode 100644 src/grass_mil/data/components/merge.py delete mode 100644 src/grass_mil/data/components/patching.py delete mode 100644 src/grass_mil/data/components/precompute.py delete mode 100644 src/grass_mil/data/components/samplers.py delete mode 100644 src/grass_mil/data/components/seed_sampling.py delete mode 100644 src/grass_mil/data/components/spatial_types.py delete mode 100644 src/grass_mil/data/components/splits.py delete mode 100644 src/grass_mil/data/components/transforms.py create mode 100644 src/grass_mil/data/datamodule.py create mode 100644 src/grass_mil/data/ego_radius.py create mode 100644 src/grass_mil/data/graphs.py create mode 100644 src/grass_mil/data/interior.py create mode 100644 src/grass_mil/data/labels.py create mode 100644 src/grass_mil/data/loaders.py create mode 100644 src/grass_mil/data/precompute.py create mode 100644 src/grass_mil/data/reducers.py create mode 100644 src/grass_mil/data/samplers.py delete mode 100644 src/grass_mil/data/spatial_omics_datamodule.py create mode 100644 src/grass_mil/data/splits.py create mode 100644 src/grass_mil/data/units.py delete mode 100644 src/grass_mil/eval.py delete mode 100644 src/grass_mil/inference/aggregation.py delete mode 100644 src/grass_mil/inference/collectors.py delete mode 100644 src/grass_mil/inference/interpretability_export.py delete mode 100644 src/grass_mil/inference/io.py delete mode 100644 src/grass_mil/inference/metrics.py delete mode 100644 src/grass_mil/inference/model_loading.py delete mode 100644 src/grass_mil/inference/predict.py create mode 100644 src/grass_mil/inference/predictions.py delete mode 100644 src/grass_mil/inference/schemas.py delete mode 100644 src/grass_mil/loocv.py create mode 100644 src/grass_mil/models/augmentations.py delete mode 100644 src/grass_mil/models/bgrl_module.py create mode 100644 src/grass_mil/models/checkpoints.py delete mode 100644 src/grass_mil/models/components/__init__.py delete mode 100644 src/grass_mil/models/components/attention.py delete mode 100644 src/grass_mil/models/components/backbones.py delete mode 100644 src/grass_mil/models/components/embeddings.py delete mode 100644 src/grass_mil/models/components/factory.py delete mode 100644 src/grass_mil/models/components/heads.py delete mode 100644 src/grass_mil/models/components/losses.py delete mode 100644 src/grass_mil/models/components/pooling.py delete mode 100644 src/grass_mil/models/components/ssl.py create mode 100644 src/grass_mil/models/encoder.py create mode 100644 src/grass_mil/models/mil.py create mode 100644 src/grass_mil/models/objectives.py create mode 100644 src/grass_mil/models/optimization.py create mode 100644 src/grass_mil/models/pretraining.py create mode 100644 src/grass_mil/models/probe.py create mode 100644 src/grass_mil/models/supervised.py delete mode 100644 src/grass_mil/models/supervised_module.py delete mode 100644 src/grass_mil/models/training/__init__.py delete mode 100644 src/grass_mil/models/training/augmentations.py delete mode 100644 src/grass_mil/models/training/bagging.py delete mode 100644 src/grass_mil/models/training/builders.py delete mode 100644 src/grass_mil/models/training/checkpoint_init.py delete mode 100644 src/grass_mil/models/training/class_weights.py delete mode 100644 src/grass_mil/models/training/loss_utils.py delete mode 100644 src/grass_mil/models/training/optimization.py delete mode 100644 src/grass_mil/models/training/probe.py delete mode 100644 src/grass_mil/models/training/ssl_runtime.py create mode 100644 src/grass_mil/predict.py delete mode 100644 src/grass_mil/probe.py create mode 100644 src/grass_mil/runtime.py delete mode 100644 src/grass_mil/utils/__init__.py delete mode 100644 src/grass_mil/utils/instantiators.py delete mode 100644 src/grass_mil/utils/logging_utils.py delete mode 100644 src/grass_mil/utils/pylogger.py delete mode 100644 src/grass_mil/utils/rich_utils.py delete mode 100644 src/grass_mil/utils/utils.py diff --git a/src/grass_mil/_paths.py b/src/grass_mil/_paths.py deleted file mode 100644 index 51ba02a..0000000 --- a/src/grass_mil/_paths.py +++ /dev/null @@ -1,29 +0,0 @@ -"""Project-root resolution shared by the CLI entrypoints. - -In a source checkout, ``rootutils`` walks up from this file to the -``.project-root`` marker and exports ``PROJECT_ROOT``, which -``configs/paths/default.yaml`` interpolates. An installed wheel has no such -marker above it, so resolution is best-effort: when it fails, the paths config -falls back to the Hydra runtime working directory. -""" - -from __future__ import annotations - -from typing import Optional - -import rootutils - -__all__ = ["setup_project_root"] - - -def setup_project_root() -> Optional[str]: - """Export ``PROJECT_ROOT`` when a ``.project-root`` marker is present. - - Returns: - The resolved project root, or ``None`` when running from an installed - package where no marker exists. - """ - try: - return str(rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)) - except (FileNotFoundError, RuntimeError): - return None diff --git a/src/grass_mil/batching.py b/src/grass_mil/batching.py deleted file mode 100644 index eecf383..0000000 --- a/src/grass_mil/batching.py +++ /dev/null @@ -1,166 +0,0 @@ -"""Splitting batched graphs back into their per-subgraph parts. - -A leaf module: it depends only on torch and PyG, so both the data layer (which -builds these batches) and the model layer (which un-builds them at inference) -can use it without importing each other. - -`Batch.to_data_list()` only works for batches assembled by -`Batch.from_data_list()`, which records the `_slice_dict`/`_inc_dict` -bookkeeping it needs. The ShaDow samplers construct their `Batch` directly and -delimit subgraphs with a `ptr` tensor instead, so `to_data_list()` raises on -them. `split_batch_by_ptr` reconstructs the pieces from `ptr`. -""" - -from __future__ import annotations - -from typing import Any, Callable, List, Optional, Sequence - -import torch -from torch_geometric.data import Batch, Data - -# Handled explicitly when rebuilding a subgraph, so they must not be copied -# across by the generic attribute loop below. -_STRUCTURAL_KEYS = { - "x", - "edge_index", - "edge_attr", - "batch", - "ptr", - "root_n_id", - "n_id", - "adj_t", - "num_nodes", -} - - -def _select_sequence_value( - *, key: str, value: Sequence[Any], subgraph_index: int, num_subgraphs: int -) -> Any: - """Pick this subgraph's entry from a per-subgraph sequence. - - Sequences that are not per-subgraph (schema name lists, for instance) are - passed through whole rather than sliced by position. - """ - if len(value) == num_subgraphs: - return [value[subgraph_index]] - return value - - -def split_batch_by_ptr(batch: Batch) -> List[Data]: - """Reconstruct the per-subgraph `Data` objects delimited by `batch.ptr`. - - Node-level tensors are sliced to the subgraph's node range, edges are - restricted to it and relabelled, and `root_n_id` is rebased onto the - subgraph's local indexing. - """ - from torch_geometric.utils import subgraph - - if getattr(batch, "ptr", None) is None: - raise ValueError("split_batch_by_ptr requires a batch carrying `ptr`.") - - num_subgraphs = len(batch.ptr) - 1 - num_nodes = int(batch.x.size(0)) - has_edge_attr = "edge_attr" in batch - num_edges = int(batch.edge_index.size(1)) - - out: List[Data] = [] - for i in range(num_subgraphs): - node_start = int(batch.ptr[i]) - node_end = int(batch.ptr[i + 1]) - node_ids = torch.arange(node_start, node_end, device=batch.ptr.device) - - sub_edge_index, sub_edge_attr = subgraph( - subset=node_ids, - edge_index=batch.edge_index, - edge_attr=batch.edge_attr if has_edge_attr else None, - relabel_nodes=True, - num_nodes=num_nodes, - ) - - root_n_id = batch.root_n_id[i : i + 1] if "root_n_id" in batch else None - # Rebase batch-space root indices onto the subgraph. Global ids (node - # ids in the source graph, flagged by `root_n_id_is_global`) must be - # left alone: subtracting a node offset from one silently turns it into - # an unrelated node, which is how the root stopped being findable in its - # own subgraph. - root_is_global = False - flags = getattr(batch, "root_n_id_is_global", None) - if flags is not None: - flags = torch.as_tensor(flags).view(-1) - root_is_global = bool(flags[0]) if flags.numel() else False - if root_n_id is not None and root_n_id.numel() == 1 and not root_is_global: - root_value = int(root_n_id.item()) - if node_start <= root_value < node_end: - root_n_id = root_n_id - node_start - - sub_data = Data( - x=batch.x[node_ids], - edge_index=sub_edge_index, - edge_attr=sub_edge_attr, - root_n_id=root_n_id, - ) - if hasattr(batch, "n_id") and isinstance(batch.n_id, torch.Tensor): - sub_data.n_id = batch.n_id[node_ids] - - for key, val in batch: - if key in _STRUCTURAL_KEYS: - continue - if isinstance(val, torch.Tensor): - if val.dim() > 0 and val.size(0) == num_subgraphs: - setattr(sub_data, key, val[i : i + 1]) - elif val.dim() == 0: - setattr(sub_data, key, val) - elif val.dim() > 0 and val.size(0) == num_nodes: - # Node-level tensors must be sliced down to this subgraph. - # They are NOT redundant with `x`: `categorical_codes` - # carries the cell-type codes the encoder embeds, and `pos` - # carries coordinates. Dropping them here silently removed - # cell type from every sampled subgraph. - setattr(sub_data, key, val[node_ids]) - elif val.dim() > 0 and has_edge_attr and val.size(0) == num_edges: - continue - else: - setattr(sub_data, key, val) - elif isinstance(val, (list, tuple)): - setattr( - sub_data, - key, - _select_sequence_value( - key=key, - value=val, - subgraph_index=i, - num_subgraphs=num_subgraphs, - ), - ) - else: - setattr(sub_data, key, val) - out.append(sub_data) - return out - - -def extract_subgraphs(batch: Any) -> Optional[List[Data]]: - """Best available split of `batch` into per-subgraph `Data` objects. - - Prefers PyG's own `to_data_list()`, and falls back to `ptr` for batches the - ShaDow samplers built directly. Returns None only when the object carries - neither route, which means it is not a batched graph at all. - - Deliberately does not swallow errors: a genuine failure to split raises, - rather than surfacing several steps downstream as a missing payload field. - """ - if hasattr(batch, "to_data_list"): - try: - return list(batch.to_data_list()) - except (RuntimeError, KeyError, AttributeError): - # Batch was not assembled by `Batch.from_data_list()`. The ShaDow - # samplers construct it directly and track subgraph boundaries in - # `ptr`, so fall through to that. - pass - if getattr(batch, "ptr", None) is not None: - return split_batch_by_ptr(batch) - return None - - -def apply_per_subgraph(batch: Batch, transform: Callable[[Data], Data]) -> Batch: - """Apply `transform` to each subgraph and re-collate.""" - return Batch.from_data_list([transform(d) for d in split_batch_by_ptr(batch)]) diff --git a/src/grass_mil/callbacks.py b/src/grass_mil/callbacks.py deleted file mode 100644 index 430f9e9..0000000 --- a/src/grass_mil/callbacks.py +++ /dev/null @@ -1,110 +0,0 @@ -"""Lightweight run-artefact callbacks. - -``EpochMetricsWriter`` appends one JSON line per epoch with every logged -``val/*`` and ``train/*`` metric to ``metrics_per_epoch.jsonl`` in the run's -output directory. The nested cross-validation driver reads it to pick the best -epoch (``argmax val/auroc``) of each inner fold. -""" - -from __future__ import annotations - -import json -import math -from pathlib import Path -from typing import Any, Dict, Optional - -import lightning as L -import torch - -__all__ = ["EpochMetricsWriter", "write_test_predictions"] - - -def _to_float(value: Any) -> Optional[float]: - try: - if isinstance(value, torch.Tensor): - value = value.detach().float().cpu().item() - value = float(value) - except (TypeError, ValueError): - return None - if math.isnan(value): - return None - return value - - -class EpochMetricsWriter(L.Callback): - def __init__(self, output_dir: str, filename: str = "metrics_per_epoch.jsonl") -> None: - super().__init__() - self.output_dir = Path(output_dir) - self.filename = filename - self._path: Optional[Path] = None - - @property - def path(self) -> Path: - if self._path is None: - self.output_dir.mkdir(parents=True, exist_ok=True) - self._path = self.output_dir / self.filename - return self._path - - def on_train_start(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None: - if trainer.is_global_zero and self.path.exists(): - self.path.unlink() - - def on_validation_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None: - # ``on_validation_end`` runs after the module's ``on_validation_epoch_end`` - # (where the region-level metrics are logged), unlike the callback's own - # ``on_validation_epoch_end``, which would see the previous epoch's values. - if trainer.sanity_checking or not trainer.is_global_zero: - return - row: Dict[str, Any] = { - "epoch": int(trainer.current_epoch), - "step": int(trainer.global_step), - } - for key, value in trainer.callback_metrics.items(): - if key.startswith(("val/", "train/")) and not key.endswith("_step"): - v = _to_float(value) - if v is not None: - row[key] = v - with self.path.open("a") as fh: - fh.write(json.dumps(row) + "\n") - - -def write_test_predictions(pl_module: L.LightningModule, output_dir: str) -> Optional[Path]: - """Write region-level test predictions (``test_predictions.csv``) if present.""" - preds = getattr(pl_module, "test_predictions", None) - if not preds: - return None - import pandas as pd - - logits = preds["logits"] - targets = preds["targets"] - rows = [] - n = int(logits.shape[0]) - region_ids = preds.get("region_ids") or [None] * n - sample_ids = preds.get("sample_ids") or [None] * n - for i in range(n): - row: Dict[str, Any] = { - "bag_id": preds["bag_ids"][i], - "region_id": region_ids[i] if i < len(region_ids) else None, - "sample_id": sample_ids[i] if i < len(sample_ids) else None, - } - lg = logits[i].view(-1) - for c in range(int(lg.numel())): - row[f"logit_{c}"] = float(lg[c]) - if preds.get("target_type") == "categorical": - probs = torch.softmax(lg.float(), dim=0) - for c in range(int(lg.numel())): - row[f"prob_class_{c}"] = float(probs[c]) - row["label"] = int(targets[i].view(-1)[0]) - elif preds.get("target_type") == "binary": - row["prob_class_1"] = float(torch.sigmoid(lg.float())[0]) - row["label"] = int(targets[i].view(-1)[0]) - else: - tv = targets[i].view(-1) - for c in range(int(tv.numel())): - row[f"target_{c}"] = float(tv[c]) - rows.append(row) - out = Path(output_dir) - out.mkdir(parents=True, exist_ok=True) - path = out / "test_predictions.csv" - pd.DataFrame(rows).to_csv(path, index=False) - return path diff --git a/src/grass_mil/configs/__init__.py b/src/grass_mil/configs/__init__.py index 56bf7f4..1684139 100644 --- a/src/grass_mil/configs/__init__.py +++ b/src/grass_mil/configs/__init__.py @@ -1 +1 @@ -# this file is needed here to include configs when building project as a package +"""Hydra configuration tree shipped with the package (``config_module="grass_mil.configs"``).""" diff --git a/src/grass_mil/configs/augmentation/bgrl_paper.yaml b/src/grass_mil/configs/augmentation/bgrl_paper.yaml deleted file mode 100644 index 5874df8..0000000 --- a/src/grass_mil/configs/augmentation/bgrl_paper.yaml +++ /dev/null @@ -1,17 +0,0 @@ -# Degree-importance-weighted node/edge dropping plus Gaussian noise on the -# continuous cell-size feature. This is the production BGRL view generator. -# -# Node importance is the mean-max normalised log-degree, so a node is dropped -# with probability min((1 - importance) * mu, p_lambda): high-degree nodes are -# preserved, low-degree ones are masked. Dropped nodes have their features -# zeroed and their cell type set to the reserved "unassigned" code, which -# requires model.encoder.categorical_embedding.reserve_unassigned=true. -mode: importance -mu: 0.25 -p_lambda: 0.45 - -# Noise is applied only to the named continuous columns. Resolve by NAME: -# a hardcoded index points at cell size in one cohort and at something else in -# another, which silently feeds the encoder pure noise. -feature_noise_std: 0.2 -feature_noise_columns: ["SIZE"] diff --git a/src/grass_mil/configs/augmentation/uniform.yaml b/src/grass_mil/configs/augmentation/uniform.yaml deleted file mode 100644 index 2645df7..0000000 --- a/src/grass_mil/configs/augmentation/uniform.yaml +++ /dev/null @@ -1,6 +0,0 @@ -# Uniform edge dropping and feature masking, independent of graph structure. -# Cohorts with no continuous node features (cell type only) cannot use the -# Gaussian-noise step, so this is the applicable preset there. -mode: uniform -drop_edge_p: 0.1 -drop_feat_p: 0.1 diff --git a/src/grass_mil/configs/callbacks/default.yaml b/src/grass_mil/configs/callbacks/default.yaml deleted file mode 100644 index 521d6e4..0000000 --- a/src/grass_mil/configs/callbacks/default.yaml +++ /dev/null @@ -1,22 +0,0 @@ -defaults: - - model_checkpoint - - early_stopping - - model_summary - - rich_progress_bar - - _self_ - -model_checkpoint: - dirpath: ${paths.output_dir}/checkpoints - filename: "epoch_{epoch:03d}" - monitor: "val/loss" - mode: "min" - save_last: True - auto_insert_metric_name: False - -early_stopping: - monitor: "val/loss" - patience: 100 - mode: "min" - -model_summary: - max_depth: -1 diff --git a/src/grass_mil/configs/callbacks/early_stopping.yaml b/src/grass_mil/configs/callbacks/early_stopping.yaml index c826c8d..d548e91 100644 --- a/src/grass_mil/configs/callbacks/early_stopping.yaml +++ b/src/grass_mil/configs/callbacks/early_stopping.yaml @@ -1,15 +1,6 @@ -# https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.EarlyStopping.html - early_stopping: _target_: lightning.pytorch.callbacks.EarlyStopping - monitor: ??? # quantity to be monitored, must be specified !!! - min_delta: 0. # minimum change in the monitored quantity to qualify as an improvement - patience: 3 # number of checks with no improvement after which training will be stopped - verbose: False # verbosity mode - mode: "min" # "max" means higher metric value is better, can be also "min" - strict: True # whether to crash the training if monitor is not found in the validation metrics - check_finite: True # when set True, stops training when the monitor becomes NaN or infinite - stopping_threshold: null # stop training immediately once the monitored quantity reaches this threshold - divergence_threshold: null # stop training as soon as the monitored quantity becomes worse than this threshold - check_on_train_epoch_end: null # whether to run early stopping at the end of the training epoch - # log_rank_zero_only: False # this keyword argument isn't available in stable version + monitor: ??? + mode: min + patience: 3 + strict: true diff --git a/src/grass_mil/configs/callbacks/mil_selection.yaml b/src/grass_mil/configs/callbacks/mil_selection.yaml deleted file mode 100644 index 661a338..0000000 --- a/src/grass_mil/configs/callbacks/mil_selection.yaml +++ /dev/null @@ -1,27 +0,0 @@ -# Model selection for MIL fine-tuning (preprint): keep the epoch with the best -# region-level validation AUROC, stop early on the region-level validation loss -# (patience 8 after a minimum of 10 epochs, set via trainer.min_epochs). -defaults: - - model_checkpoint - - early_stopping - - model_summary - - rich_progress_bar - - _self_ - -model_checkpoint: - dirpath: ${paths.output_dir}/checkpoints - filename: "epoch_{epoch:03d}" - monitor: "val/auroc" - mode: "max" - save_last: True - save_top_k: 1 - auto_insert_metric_name: False - -early_stopping: - monitor: "val/loss_region" - patience: 8 - mode: "min" - strict: False - -model_summary: - max_depth: -1 diff --git a/src/grass_mil/configs/callbacks/model_checkpoint.yaml b/src/grass_mil/configs/callbacks/model_checkpoint.yaml index bf946e8..13c923b 100644 --- a/src/grass_mil/configs/callbacks/model_checkpoint.yaml +++ b/src/grass_mil/configs/callbacks/model_checkpoint.yaml @@ -1,17 +1,10 @@ -# https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html - model_checkpoint: _target_: lightning.pytorch.callbacks.ModelCheckpoint - dirpath: null # directory to save the model file - filename: null # checkpoint filename - monitor: null # name of the logged metric which determines when model is improving - verbose: False # verbosity mode - save_last: null # additionally always save an exact copy of the last checkpoint to a file last.ckpt - save_top_k: 1 # save k best models (determined by above metric) - mode: "min" # "max" means higher metric value is better, can be also "min" - auto_insert_metric_name: True # when True, the checkpoints filenames will contain the metric name - save_weights_only: False # if True, then only the model’s weights will be saved - every_n_train_steps: null # number of training steps between checkpoints - train_time_interval: null # checkpoints are monitored at the specified time interval - every_n_epochs: null # number of epochs between checkpoints - save_on_train_epoch_end: null # whether to run checkpointing at the end of the training epoch or the end of validation + dirpath: null + filename: null + monitor: null + mode: min + save_last: null + save_top_k: 1 + auto_insert_metric_name: true + every_n_train_steps: null diff --git a/src/grass_mil/configs/callbacks/model_summary.yaml b/src/grass_mil/configs/callbacks/model_summary.yaml index b75981d..9ddc006 100644 --- a/src/grass_mil/configs/callbacks/model_summary.yaml +++ b/src/grass_mil/configs/callbacks/model_summary.yaml @@ -1,5 +1,3 @@ -# https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.RichModelSummary.html - model_summary: _target_: lightning.pytorch.callbacks.RichModelSummary - max_depth: 1 # the maximum depth of layer nesting that the summary will include + max_depth: 1 diff --git a/src/grass_mil/configs/callbacks/none.yaml b/src/grass_mil/configs/callbacks/none.yaml deleted file mode 100644 index e69de29..0000000 diff --git a/src/grass_mil/configs/callbacks/pretrain_checkpoint.yaml b/src/grass_mil/configs/callbacks/pretrain_checkpoint.yaml index 9c831eb..739d32b 100644 --- a/src/grass_mil/configs/callbacks/pretrain_checkpoint.yaml +++ b/src/grass_mil/configs/callbacks/pretrain_checkpoint.yaml @@ -1,5 +1,4 @@ -# Pretraining: a checkpoint every N optimizer steps (preprint: 5000), all kept, -# plus last.ckpt. No validation metric is involved. +# Pretraining: keep a checkpoint every N optimizer steps, plus last.ckpt. defaults: - model_checkpoint - model_summary @@ -10,10 +9,7 @@ model_checkpoint: dirpath: ${paths.output_dir}/checkpoints filename: "step_{step:07d}" monitor: null - save_last: True + save_last: true save_top_k: -1 every_n_train_steps: 5000 - auto_insert_metric_name: False - -model_summary: - max_depth: -1 + auto_insert_metric_name: false diff --git a/src/grass_mil/configs/callbacks/refit.yaml b/src/grass_mil/configs/callbacks/refit.yaml deleted file mode 100644 index 9c0c44f..0000000 --- a/src/grass_mil/configs/callbacks/refit.yaml +++ /dev/null @@ -1,18 +0,0 @@ -# Refit on all development folds for a fixed number of epochs: no early -# stopping, no metric-based selection; the final-epoch weights are the model. -defaults: - - model_checkpoint - - model_summary - - rich_progress_bar - - _self_ - -model_checkpoint: - dirpath: ${paths.output_dir}/checkpoints - filename: "last_epoch" - monitor: null - save_last: True - save_top_k: 0 - auto_insert_metric_name: False - -model_summary: - max_depth: -1 diff --git a/src/grass_mil/configs/callbacks/rich_progress_bar.yaml b/src/grass_mil/configs/callbacks/rich_progress_bar.yaml index de6f1cc..9ffb16b 100644 --- a/src/grass_mil/configs/callbacks/rich_progress_bar.yaml +++ b/src/grass_mil/configs/callbacks/rich_progress_bar.yaml @@ -1,4 +1,2 @@ -# https://lightning.ai/docs/pytorch/latest/api/lightning.pytorch.callbacks.RichProgressBar.html - rich_progress_bar: _target_: lightning.pytorch.callbacks.RichProgressBar diff --git a/src/grass_mil/configs/callbacks/selection.yaml b/src/grass_mil/configs/callbacks/selection.yaml new file mode 100644 index 0000000..9d4c65f --- /dev/null +++ b/src/grass_mil/configs/callbacks/selection.yaml @@ -0,0 +1,26 @@ +# Keep the epoch with the best validation metric of the task and stop early when +# it has not improved for `task.selection.patience` epochs (after +# trainer.min_epochs). +defaults: + - model_checkpoint + - early_stopping + - model_summary + - rich_progress_bar + - _self_ + +model_checkpoint: + dirpath: ${paths.output_dir}/checkpoints + filename: "epoch_{epoch:03d}" + monitor: ${task.selection.metric} + mode: ${task.selection.mode} + save_last: true + save_top_k: 1 + auto_insert_metric_name: false + +early_stopping: + monitor: ${task.selection.metric} + mode: ${task.selection.mode} + patience: ${task.selection.patience} + strict: false + # NaN (e.g. AUROC of a single-class validation set) counts as no improvement. + check_finite: false diff --git a/src/grass_mil/configs/data/sampler/ego_graph.yaml b/src/grass_mil/configs/data/sampler/ego_graph.yaml new file mode 100644 index 0000000..fcfe3cd --- /dev/null +++ b/src/grass_mil/configs/data/sampler/ego_graph.yaml @@ -0,0 +1,18 @@ +# One bag per unit per step: roots_per_unit ego-graphs of depth `depth`. +depth: 4 # keep equal to model.encoder.num_layers +num_neighbors: -1 # -1 = full k-hop neighbourhood +roots_per_unit: 256 +batch_size: 64 # ego-graphs per batch when scoring +merge_units: false +root_weighting: null # null | inverse | sqrt_inverse (by cell-type frequency) +interior_roots: + enabled: false + n_hops: null # null = depth + hull_tolerance: 1.0 + keep_largest_component: true + cache_dir: null +ego_radius: + enabled: false + radius_per_hop: 75.0 # radius = radius_per_hop * depth + radius_offset + radius_offset: 55.0 + keep_root_component: true diff --git a/src/grass_mil/configs/data/sampler/ego_graph_mil.yaml b/src/grass_mil/configs/data/sampler/ego_graph_mil.yaml deleted file mode 100644 index b02deae..0000000 --- a/src/grass_mil/configs/data/sampler/ego_graph_mil.yaml +++ /dev/null @@ -1,36 +0,0 @@ -# Ego-graph bags for MIL fine-tuning (preprint defaults). -# Training: `subsample_count` random roots per region per epoch (uniform, not -# cell-type weighted), one batch per region (`unit_as_bag`), full k-hop -# ego-graphs (depth 4, unlimited neighbours, no radius cap). -# Validation: a fixed 1024-root subset per region (`val_sampler`, below). -# Test/predict: every root once (derived automatically from this config). -name: shadow_custom -kwargs: {} -runtime: - enabled: true - depth: 4 - num_neighbors: -1 # -1 = all neighbours at every hop - subgraph_batch_size: 64 # upper bound per chunk; ignored when unit_as_bag - replace: false - shuffle_subgraphs: true - proportional_root_sampling: false - property_name: cell_type - weight_mode: inverse - min_weight: 1.0e-6 - subsample_fraction: 1.0 - subsample_count: 256 - subsample_seed: null - fixed_per_unit_seed: null - unit_as_bag: true - merge_units: false - interior_seeds: - enabled: false - n_hops: null - seed_tolerance_multiplier: 1.0 - keep_largest_component: true - cache_dir: null - ego_radius: - enabled: false - radius_per_hop: 75.0 - radius_offset: 55.0 - keep_root_component: true diff --git a/src/grass_mil/configs/data/sampler/ego_graph_pretrain.yaml b/src/grass_mil/configs/data/sampler/ego_graph_pretrain.yaml index 1f79f87..fb82682 100644 --- a/src/grass_mil/configs/data/sampler/ego_graph_pretrain.yaml +++ b/src/grass_mil/configs/data/sampler/ego_graph_pretrain.yaml @@ -1,33 +1,7 @@ -# Ego-graph batches for BGRL pretraining (preprint "A3" defaults): roots drawn -# uniformly from the union of ALL regions (`merge_units`), 64 full k-hop -# ego-graphs per step (depth 4, unlimited neighbours, no radius cap). -name: shadow_custom -kwargs: {} -runtime: - enabled: true - depth: 4 - num_neighbors: -1 - subgraph_batch_size: 64 - replace: false - shuffle_subgraphs: true - proportional_root_sampling: false - property_name: cell_type - weight_mode: inverse - min_weight: 1.0e-6 - subsample_fraction: 1.0 - subsample_count: null - subsample_seed: null - fixed_per_unit_seed: null - unit_as_bag: false - merge_units: true - interior_seeds: - enabled: false - n_hops: null - seed_tolerance_multiplier: 1.0 - keep_largest_component: true - cache_dir: null - ego_radius: - enabled: false - radius_per_hop: 75.0 - radius_offset: 55.0 - keep_root_component: true +# Pretraining: roots of all units pooled, batch_size ego-graphs per step. +defaults: + - ego_graph + - _self_ + +roots_per_unit: null +merge_units: true diff --git a/src/grass_mil/configs/data/sampler/identity.yaml b/src/grass_mil/configs/data/sampler/identity.yaml deleted file mode 100644 index 3366849..0000000 --- a/src/grass_mil/configs/data/sampler/identity.yaml +++ /dev/null @@ -1,2 +0,0 @@ -name: identity -kwargs: {} diff --git a/src/grass_mil/configs/data/sampler/shadow_custom.yaml b/src/grass_mil/configs/data/sampler/shadow_custom.yaml deleted file mode 100644 index 8176031..0000000 --- a/src/grass_mil/configs/data/sampler/shadow_custom.yaml +++ /dev/null @@ -1,33 +0,0 @@ -name: shadow_custom -kwargs: {} -runtime: - enabled: true - depth: 2 - num_neighbors: 8 - subgraph_batch_size: 32 - replace: false - shuffle_subgraphs: true - proportional_root_sampling: true - property_name: cell_type - weight_mode: inverse - min_weight: 1.0e-6 - # Restrict ego-graph seeds to the tissue interior. Ego-graphs rooted at the - # section edge are truncated by the boundary rather than by biology, so they - # are not comparable to interior ones. Excludes only the n_hops graph-hop rim - # around the convex hull; internal holes are not treated as boundary. - interior_seeds: - enabled: false - # null uses `depth`, so the rim matches the ego-graph radius. - n_hops: null - seed_tolerance_multiplier: 1.0 - keep_largest_component: true - # Directory for the on-disk index cache; null disables caching. - cache_dir: null - # Cap the ego-graph's physical extent: r_k = radius_per_hop * depth + offset. - # Distances are in datamodule coordinate units (um after coord_scale_um), NOT - # pixels -- convert the cohort's pixel values before setting these. - ego_radius: - enabled: false - radius_per_hop: 75.0 - radius_offset: 55.0 - keep_root_component: true diff --git a/src/grass_mil/configs/data/sampler/shadow_native.yaml b/src/grass_mil/configs/data/sampler/shadow_native.yaml deleted file mode 100644 index 97971fb..0000000 --- a/src/grass_mil/configs/data/sampler/shadow_native.yaml +++ /dev/null @@ -1,33 +0,0 @@ -name: shadow_native -kwargs: {} -runtime: - enabled: true - depth: 2 - num_neighbors: 8 - subgraph_batch_size: 32 - replace: false - shuffle_subgraphs: true - proportional_root_sampling: true - property_name: cell_type - weight_mode: inverse - min_weight: 1.0e-6 - # Restrict ego-graph seeds to the tissue interior. Ego-graphs rooted at the - # section edge are truncated by the boundary rather than by biology, so they - # are not comparable to interior ones. Excludes only the n_hops graph-hop rim - # around the convex hull; internal holes are not treated as boundary. - interior_seeds: - enabled: false - # null uses `depth`, so the rim matches the ego-graph radius. - n_hops: null - seed_tolerance_multiplier: 1.0 - keep_largest_component: true - # Directory for the on-disk index cache; null disables caching. - cache_dir: null - # Cap the ego-graph's physical extent: r_k = radius_per_hop * depth + offset. - # Distances are in datamodule coordinate units (um after coord_scale_um), NOT - # pixels -- convert the cohort's pixel values before setting these. - ego_radius: - enabled: false - radius_per_hop: 75.0 - radius_offset: 55.0 - keep_root_component: true diff --git a/src/grass_mil/configs/data/spatial_omics.yaml b/src/grass_mil/configs/data/spatial_omics.yaml index 9cee887..663e937 100644 --- a/src/grass_mil/configs/data/spatial_omics.yaml +++ b/src/grass_mil/configs/data/spatial_omics.yaml @@ -1,109 +1,66 @@ +# A cohort of cell tables listed in a manifest, cut into graph units, with +# labels attached per unit. See docs/data_preparation.md. defaults: - - sampler: shadow_custom + - sampler: ego_graph - _self_ -_target_: grass_mil.data.spatial_omics_datamodule.SpatialOmicsDataModule +_target_: grass_mil.data.SpatialOmicsDataModule -data_dir: ${paths.data_dir} -raw_manifest_path: ${paths.data_dir}/raw/manifest.csv -processed_dir: ${paths.data_dir}/processed/spatial_omics - -batch_size: 1 -num_workers: 0 -pin_memory: False - -coord_scale_um: 1.0 -sample_unit: tile # tile | full -reducer_scope: sample # sample | dataset -keep_raw_molecular: false +# CSV with sample_id, input_path and optionally region_id, input_type, +# polygons_path (paths relative to the manifest). +manifest_path: ${paths.data_dir}/manifest.csv +# Graphs are built here once and reused while the settings below are unchanged. +processed_dir: ${paths.data_dir}/processed force_precompute: false -min_cells: 10 -use_molecular_features: true - -categorical_features: - include_labels: ["cell_type"] - -split: - train_val_test_split: [0.8, 0.1, 0.1] - split_by: sample # sample | region | patch - loocv: - enabled: false - fold_unit: region # region | sample - holdout_id: null - validation_strategy: heldout_fold_items # heldout_fold_items | patches_from_train_items - val_ratio: 0.1 # exactly 0.0 disables val; non-zero guarantees at least one val item - seed: 42 - # Explicit id lists (region ids by default, or sample ids with id_kind: sample) - # re-partition an existing precompute without re-processing. Each entry is a - # list or a path to a .json / .txt file; ids not listed are dropped. This is - # what `grass-mil-nested-cv` and `grass-mil-probe` use. - explicit: - train: null - val: null - test: null - id_kind: region # region | sample | patch (every tile its own unit) - -manifest: - sample_id: sample_id - input_path: input_path - input_type: input_type - region_id: region_id - polygons_path: polygons_path - -csv: - sep: "," - coord_columns: ["x", "y"] - cell_id_column: cell_id - categorical_label_columns: [] - molecular_columns: null -h5ad: - coord_source: obsm # obs | obsm - coord_key: spatial - coord_columns: ["x", "y"] - cell_id_column: null - categorical_label_columns: [] - molecular_layer: null - molecular_features: null +# What one bag is: a whole region, a square tile (data.tiling), or a polygon of +# the region's polygons_path. +unit: region # region | tile | polygon +coord_scale_um: 1.0 # coordinates are multiplied by this +min_cells: 10 # units with fewer cells are dropped +tiling: + size: 200.0 + stride: 200.0 -sce: - assay_name: null - coord_source: colData # colData | reducedDims - coord_key: null - coord_columns: ["x", "y"] +input: + coord_columns: [x, y] + obsm_key: spatial # h5ad: coordinates in obsm[obsm_key]; null = obs coord_columns + reduced_dim: null # SCE: coordinates in reducedDim; null = colData coord_columns cell_id_column: null - categorical_label_columns: [] - transpose_assay: true - molecular_features: null + cell_type_column: cell_type # null = no cell types (features only) + features: all # continuous node features: all | [names] | null (cell type only) + h5ad_layer: null + sce_assay: null + sep: "," -graph_builder: - name: delaunay - kwargs: - edge_features: ["distance", "neighbor"] - neighbor_cutoff_um: 20.0 +graph: + edge_features: [log_distance, neighbor] + neighbor_cutoff: 20.0 # contact distance, in coordinate units + log_distance_bounds: [2.0, 5.0] # ln-length mapped to [0, 1] feature_reducer: - name: identity - fit_mode: train_only # train_only | global - kwargs: {} + name: identity # identity | pca + n_components: 32 + fit_on: all # all | train | sample + state_from: null # processed_index.json of a cohort whose PCA basis to reuse + +labels: + file: null # CSV with one row per labelled id + id_column: region_id # unit, region or sample ids + columns: ${oc.select:task.target_columns,null} + patient_column: null # grouping for splits; never a target +# Skip units without a label (supervised training); pretraining uses all units. +labelled_only: true -# `sampler` comes from the data/sampler group (shadow_custom by default; -# `data/sampler=ego_graph_mil` is the preprint MIL setting). Override single -# keys as `data.sampler.runtime.=...`. -val_sampler: null -# Test/predict loader; null derives an all-roots, unweighted variant of `sampler`. -test_sampler: null - -tiling: - tile_size_um: 200.0 - stride_um: 200.0 - min_cells: 10 - -graph_labels: - label_file: null - id_column: id - tasks: null - scope: region - patient_column: null # grouping column; attached as data.patient_id, never a target - -transforms: [] +split: + ratios: [0.8, 0.1, 0.1] # random train/val/test split, grouped by patient + seed: 0 + # Explicit ids (lists or .json/.txt files) replace the random split. + train: null + val: null + test: null + +# Validation scores a fixed subset of roots per unit; test scores every cell. +val_roots_per_unit: 1024 +num_workers: 0 +pin_memory: false diff --git a/src/grass_mil/configs/debug/default.yaml b/src/grass_mil/configs/debug/default.yaml deleted file mode 100644 index 1886902..0000000 --- a/src/grass_mil/configs/debug/default.yaml +++ /dev/null @@ -1,35 +0,0 @@ -# @package _global_ - -# default debugging setup, runs 1 full epoch -# other debugging configs can inherit from this one - -# overwrite task name so debugging logs are stored in separate folder -task_name: "debug" - -# disable callbacks and loggers during debugging -callbacks: null -logger: null - -extras: - ignore_warnings: False - enforce_tags: False - -# sets level of all command line loggers to 'DEBUG' -# https://hydra.cc/docs/tutorials/basic/running_your_app/logging/ -hydra: - job_logging: - root: - level: DEBUG - - # use this to also set hydra loggers to 'DEBUG' - # verbose: True - -trainer: - max_epochs: 1 - accelerator: cpu # debuggers don't like gpus - devices: 1 # debuggers don't like multiprocessing - detect_anomaly: true # raise exception if NaN or +/-inf is detected in any tensor - -data: - num_workers: 0 # debuggers don't like multiprocessing - pin_memory: False # disable gpu memory pin diff --git a/src/grass_mil/configs/debug/fdr.yaml b/src/grass_mil/configs/debug/fdr.yaml deleted file mode 100644 index 7f2d34f..0000000 --- a/src/grass_mil/configs/debug/fdr.yaml +++ /dev/null @@ -1,9 +0,0 @@ -# @package _global_ - -# runs 1 train, 1 validation and 1 test step - -defaults: - - default - -trainer: - fast_dev_run: true diff --git a/src/grass_mil/configs/debug/limit.yaml b/src/grass_mil/configs/debug/limit.yaml deleted file mode 100644 index 514d77f..0000000 --- a/src/grass_mil/configs/debug/limit.yaml +++ /dev/null @@ -1,12 +0,0 @@ -# @package _global_ - -# uses only 1% of the training data and 5% of validation/test data - -defaults: - - default - -trainer: - max_epochs: 3 - limit_train_batches: 0.01 - limit_val_batches: 0.05 - limit_test_batches: 0.05 diff --git a/src/grass_mil/configs/debug/overfit.yaml b/src/grass_mil/configs/debug/overfit.yaml deleted file mode 100644 index 9906586..0000000 --- a/src/grass_mil/configs/debug/overfit.yaml +++ /dev/null @@ -1,13 +0,0 @@ -# @package _global_ - -# overfits to 3 batches - -defaults: - - default - -trainer: - max_epochs: 20 - overfit_batches: 3 - -# model ckpt and early stopping need to be disabled during overfitting -callbacks: null diff --git a/src/grass_mil/configs/debug/profiler.yaml b/src/grass_mil/configs/debug/profiler.yaml deleted file mode 100644 index 2bd7da8..0000000 --- a/src/grass_mil/configs/debug/profiler.yaml +++ /dev/null @@ -1,12 +0,0 @@ -# @package _global_ - -# runs with execution time profiling - -defaults: - - default - -trainer: - max_epochs: 1 - profiler: "simple" - # profiler: "advanced" - # profiler: "pytorch" diff --git a/src/grass_mil/configs/eval.yaml b/src/grass_mil/configs/eval.yaml deleted file mode 100644 index 5b34bd5..0000000 --- a/src/grass_mil/configs/eval.yaml +++ /dev/null @@ -1,29 +0,0 @@ -# @package _global_ - -defaults: - - _self_ - - data: spatial_omics # choose datamodule with `test_dataloader()` for evaluation - - model: supervised_module - - optim: adamw - - scheduler: cosine_epoch - # `task` is merged after `model`, `optim` and `scheduler` so that a task - # preset can override them. finetune_mil sets the per-class head and - # attention widths, and pretrain_bgrl overrides the scheduler; both are - # silently discarded if `task` is listed first. - - task: finetune_mean - - logger: null - - trainer: default - - paths: default - - extras: default - - hydra: default - - callbacks: none - # Regime presets compose here too (`regime=mil`), so the evaluated model - # matches the checkpoint. - - regime: null - -task_name: "eval" - -tags: ["dev"] - -# passing checkpoint path is necessary for evaluation -ckpt_path: ??? diff --git a/src/grass_mil/configs/experiment/example_cohort.yaml b/src/grass_mil/configs/experiment/example_cohort.yaml index 324549b..43f70f3 100644 --- a/src/grass_mil/configs/experiment/example_cohort.yaml +++ b/src/grass_mil/configs/experiment/example_cohort.yaml @@ -1,53 +1,33 @@ # @package _global_ -# Template for a cohort of whole regions (cores, ROIs, slides) given as CSV cell -# tables. Copy this file, rename it (e.g. my_cohort.yaml), fill in the paths and -# column names, then select it on every command with `experiment=my_cohort`: +# Template describing one cohort. Copy it to my_cohort.yaml, fill in paths and +# column names, and pass experiment=my_cohort to every command: # -# grass-mil-train experiment=my_cohort regime=pretrain_a3 paths.output_dir=runs/pretrain_a3 -# grass-mil-nested-cv experiment=my_cohort regime=linear_ftall_reg model.init_from_ckpt=... -# grass-mil-extract experiment=my_cohort regime=linear_ftall_reg extract.nested_cv_summary=... -# -# Keys listed here are the ones that describe a dataset. Every other setting -# (encoder, ego-graphs, regimes, clustering) keeps its preset default and can -# still be overridden on the command line. Values are examples. +# grass-mil-train experiment=my_cohort regime=pretrain +# grass-mil-nested-cv experiment=my_cohort model.init_from_ckpt=/checkpoints/last.ckpt +# grass-mil-extract experiment=my_cohort extract.nested_cv_summary=/nested_cv_summary.json paths: - data_dir: /path/to/my_cohort # holds raw/ and processed/ + data_dir: /path/to/my_cohort data: - # --- inputs ------------------------------------------------------------- - raw_manifest_path: ${paths.data_dir}/raw/manifest.csv # sample_id,input_path,input_type,region_id - processed_dir: ${paths.data_dir}/processed # graphs are written here once - sample_unit: full # full = one graph per region; tile = cut regions into tiles (see docs) - coord_scale_um: 1.0 # multiply coordinates by this to get the unit used below - min_cells: 10 # regions with fewer cells are dropped - csv: - sep: "," + manifest_path: ${paths.data_dir}/manifest.csv # sample_id,input_path[,region_id] + processed_dir: ${paths.data_dir}/processed + unit: region + coord_scale_um: 1.0 + input: coord_columns: [x, y] cell_id_column: cell_id - categorical_label_columns: [cell_type] # the cell-type column (any name) - molecular_columns: [SIZE] # continuous node features; [] for cell type only - categorical_features: - include_labels: [cell_type] # must match the column above - # --- graph -------------------------------------------------------------- - graph_builder: - name: delaunay - kwargs: - edge_features: [log_distance, neighbor] # GRASS-MIL default - neighbor_cutoff_um: 55.0 # contact distance, in the coordinate unit after coord_scale_um - # --- labels and grouping ------------------------------------------------- - graph_labels: - label_file: ${paths.data_dir}/raw/labels.csv # REGION_ID,STAS,PATIENT - id_column: REGION_ID - tasks: [STAS] # target column(s); 0/1 or category names - scope: region - patient_column: PATIENT # grouping; never a target - split: - split_by: region # persisted random split used only when no explicit split is given - train_val_test_split: [0.8, 0.1, 0.1] + cell_type_column: cell_type + features: [area] # continuous node features; null for cell type only + graph: + neighbor_cutoff: 20.0 # contact distance in coordinate units + labels: + file: ${paths.data_dir}/labels.csv # region_id,label,patient_id + id_column: region_id + patient_column: patient_id -nested_cv: - folds_file: ${paths.data_dir}/raw/folds.json # JSON list of lists of region ids, patient-grouped +task: + target_columns: [label] -task_name: my_cohort -tags: [my_cohort] +nested_cv: + folds_file: ${paths.data_dir}/folds.json # JSON list of lists of region ids diff --git a/src/grass_mil/configs/experiment/grass_mil_pretrain_a3.yaml b/src/grass_mil/configs/experiment/grass_mil_pretrain_a3.yaml deleted file mode 100644 index b158991..0000000 --- a/src/grass_mil/configs/experiment/grass_mil_pretrain_a3.yaml +++ /dev/null @@ -1,21 +0,0 @@ -# @package _global_ -# Same recipe as `regime=pretrain_a3`, for runs without a cohort experiment file: -# grass-mil-train experiment=grass_mil_pretrain_a3 data.raw_manifest_path=... data.processed_dir=... -# With a cohort file use `experiment=my_cohort regime=pretrain_a3` (one experiment -# file per command). Keep the two files identical. -defaults: - - override /model: bgrl_a3 - - override /task: pretrain_bgrl_a3 - - override /callbacks: pretrain_checkpoint - - override /data/sampler: ego_graph_pretrain - -regime_name: pretrain_a3 -task_name: pretrain_a3 -tags: ["grass_mil", "pretrain", "a3"] -test: false - -data: - split: - train_val_test_split: [1.0, 0.0, 0.0] - split_by: region - val_sampler: null diff --git a/src/grass_mil/configs/extract.yaml b/src/grass_mil/configs/extract.yaml index 65b515c..d038783 100644 --- a/src/grass_mil/configs/extract.yaml +++ b/src/grass_mil/configs/extract.yaml @@ -1,33 +1,24 @@ -# @package _global_ -# grass-mil-extract: per-cell cache (embeddings, instance logits, attention) of -# one or more trained models over ALL cells of every region. +# grass-mil-extract: per-cell outputs (embeddings, attention space, logits) of +# trained models over every cell of every unit. # -# grass-mil-extract regime=linear_ftall_reg data=... \ -# extract.nested_cv_summary=/nested_cv_summary.json -# grass-mil-extract regime=mil data=... +extract.models.final=/path/to/last.ckpt +# grass-mil-extract experiment=my_cohort extract.nested_cv_summary=/nested_cv_summary.json +# grass-mil-extract experiment=my_cohort +extract.models.final=/checkpoints/last.ckpt # -# The model is composed from the same `regime`/`model` keys as the training run -# so the checkpoint fits; the ego-graph settings come from `data.sampler`. +# Models are rebuilt from their checkpoints, including the ego-graph settings +# they were trained with; `data` only needs to point at the same cohort. defaults: - train - _self_ -task_name: "extract" -train: false -test: false +task_name: extract extract: - # Either a nested-CV summary (one model per outer fold, held-out regions - # recorded) or an explicit {name: checkpoint} mapping. + # One model per outer fold of a cross-validation (held-out units recorded) ... nested_cv_summary: null + # ... and/or explicit {name: checkpoint} models. models: {} - # Optional held-out region lists per explicit model: +extract.held_out.=[...] - held_out: {} output_dir: ${paths.output_dir}/cells - regions: null # optional list / json file of unit ids to extract (null = all) - # auto = patch when the precompute holds several tiles per region, else region. - unit_id: auto # auto | region | sample | patch + units: null # subset of unit ids (list or file); null = every unit batch_size: 512 # ego-graphs per forward pass - autocast: true # fp16 autocast on GPU - device: auto + device: auto # auto | cpu | cuda resume: true diff --git a/src/grass_mil/configs/extras/default.yaml b/src/grass_mil/configs/extras/default.yaml deleted file mode 100644 index b9c6b62..0000000 --- a/src/grass_mil/configs/extras/default.yaml +++ /dev/null @@ -1,8 +0,0 @@ -# disable python warnings if they annoy you -ignore_warnings: False - -# ask user for tags if none are provided in the config -enforce_tags: True - -# pretty print config tree at the start of the run using Rich library -print_config: True diff --git a/src/grass_mil/configs/hydra/default.yaml b/src/grass_mil/configs/hydra/default.yaml index a61e9b3..2e35558 100644 --- a/src/grass_mil/configs/hydra/default.yaml +++ b/src/grass_mil/configs/hydra/default.yaml @@ -1,11 +1,7 @@ -# https://hydra.cc/docs/configure_hydra/intro/ - -# enable color logging defaults: - override hydra_logging: colorlog - override job_logging: colorlog -# output directory, generated dynamically on each run run: dir: ${paths.log_dir}/${task_name}/runs/${now:%Y-%m-%d}_${now:%H-%M-%S} sweep: @@ -15,5 +11,4 @@ sweep: job_logging: handlers: file: - # Incorporates fix from https://github.com/facebookresearch/hydra/pull/2242 filename: ${hydra.runtime.output_dir}/${task_name}.log diff --git a/src/grass_mil/configs/inference/aggregation/default.yaml b/src/grass_mil/configs/inference/aggregation/default.yaml deleted file mode 100644 index e26132e..0000000 --- a/src/grass_mil/configs/inference/aggregation/default.yaml +++ /dev/null @@ -1,15 +0,0 @@ -enabled: true -# mean | max | none | attention_weighted | mean_probability -mode: mean -# patch | region | sample | patient -# `patient` pools the regions of a patient: with mode=mean_probability the -# patient score is the mean of its regions' probabilities (preprint default); -# with mode=attention_weighted all cells of all regions form one merged bag. -bag_scope: region -subsample_fraction: 1.0 # in (0, 1]; when <1.0, predict path applies pre-forward root subsampling for efficiency -subsample_seed: null # optional seed for deterministic subsampling -# Region -> patient mapping used by bag_scope=patient (a label CSV). -patient_map: - label_file: null - id_column: REGION_ID - patient_column: PATIENT diff --git a/src/grass_mil/configs/inference/embeddings/default.yaml b/src/grass_mil/configs/inference/embeddings/default.yaml deleted file mode 100644 index 6cc6ea7..0000000 --- a/src/grass_mil/configs/inference/embeddings/default.yaml +++ /dev/null @@ -1,5 +0,0 @@ -enabled: true -extract_node: false -save: true -filename: graph_embeddings.csv -node_filename: node_embeddings.csv diff --git a/src/grass_mil/configs/inference/interpretability/default.yaml b/src/grass_mil/configs/inference/interpretability/default.yaml deleted file mode 100644 index 4dd9d37..0000000 --- a/src/grass_mil/configs/inference/interpretability/default.yaml +++ /dev/null @@ -1,40 +0,0 @@ -enabled: false -instance_filename: instance_table.csv -spatial_filename: spatial_table.csv -id_column: instance_id -bag_id_column: bag_id -composition_label: cell_type -# Single categorical label per instance, needed by `filtration_curves` and -# `per_niche_cell_type_enrichment`. An instance is a neighbourhood, not a cell, -# so this is the dominant cell type of the neighbourhood (argmax of the -# composition vector). Set to null to omit the column. -cell_type_column: cell_type -composition_prefix: comp_ -embedding_prefix: inst_emb -score_column: score -# How to export `score` from `instance_logits`. -# - sigmoid: sigmoid(logit[score_logit_index]) (legacy/default) -# - identity: raw logit[score_logit_index] (useful for survival hazard-style outputs) -# - softmax: softmax(logits)[score_logit_index] (class probability) -score_mode: sigmoid -# Which logit column to project into `score` (0-based). -score_logit_index: 0 -attention_column: attention -# Which attention channel becomes the canonical `attention` column when the -# head emits one channel per class. null selects the last channel (the -# positive class for binary tasks). Per-class columns are always exported -# as attention_c{j} alongside the raw attention_logit_c{j}. -attention_class_index: null -require_composition: true - -spatial: - enabled: false - undirected: true - -# Cell-level export: one row per node of every sampled ego-graph, plus the -# intra-subgraph edges with their distances. Requires a shadow sampler, since -# only then does the payload carry per-instance subgraphs. -cells: - enabled: false - cell_filename: cell_table.csv - cell_edge_filename: cell_edges.csv diff --git a/src/grass_mil/configs/inference/metrics/default.yaml b/src/grass_mil/configs/inference/metrics/default.yaml deleted file mode 100644 index ee6b00f..0000000 --- a/src/grass_mil/configs/inference/metrics/default.yaml +++ /dev/null @@ -1,5 +0,0 @@ -enabled: true -per_group: true - -categorical: - threshold: 0.5 diff --git a/src/grass_mil/configs/inference/predict.yaml b/src/grass_mil/configs/inference/predict.yaml deleted file mode 100644 index c20ba78..0000000 --- a/src/grass_mil/configs/inference/predict.yaml +++ /dev/null @@ -1,42 +0,0 @@ -# @package _global_ - -defaults: - - _self_ - - /data: spatial_omics - - /model: supervised_module - - /optim: adamw - - /scheduler: cosine_epoch - # `task` is merged after `model`, `optim` and `scheduler` so a task preset - # can override them. finetune_mil sets the per-class head and attention - # widths; listed first, they are silently discarded and the model is built - # narrower than the checkpoint it is asked to load. - - /task: finetune_mean - - /trainer: default - - /paths: default - - /extras: default - - /hydra: default - - metrics: default - - aggregation: default - - embeddings: default - - interpretability: default - # Training-regime presets compose here too, so `regime=mil` rebuilds the - # exact model (head width, attention, sampler depth) the checkpoint was trained with. - - /callbacks: none - - /regime: null - -task_name: "predict" -tags: ["dev"] - -ckpt_path: ??? - -predict: - save_predictions: true - save_metrics: true - include_targets: true - include_attention: true - output_subdir: predictions - -output: - predictions_filename: predictions.csv - metrics_filename: metrics.json - summary_filename: inference_summary.json diff --git a/src/grass_mil/configs/interpretability/clustering/agglomerative.yaml b/src/grass_mil/configs/interpretability/clustering/agglomerative.yaml deleted file mode 100644 index 24a27ca..0000000 --- a/src/grass_mil/configs/interpretability/clustering/agglomerative.yaml +++ /dev/null @@ -1,6 +0,0 @@ -enabled: true -method: agglomerative -params: - n_clusters: 8 - linkage: ward - metric: euclidean diff --git a/src/grass_mil/configs/interpretability/clustering/dbscan.yaml b/src/grass_mil/configs/interpretability/clustering/dbscan.yaml deleted file mode 100644 index 7f028e6..0000000 --- a/src/grass_mil/configs/interpretability/clustering/dbscan.yaml +++ /dev/null @@ -1,10 +0,0 @@ -enabled: true -method: dbscan -params: - eps: 0.5 - min_samples: 5 - metric: cosine - algorithm: auto - leaf_size: 30 - p: null - n_jobs: null diff --git a/src/grass_mil/configs/interpretability/clustering/hdbscan.yaml b/src/grass_mil/configs/interpretability/clustering/hdbscan.yaml deleted file mode 100644 index 9bbfcf3..0000000 --- a/src/grass_mil/configs/interpretability/clustering/hdbscan.yaml +++ /dev/null @@ -1,18 +0,0 @@ -enabled: true -method: hdbscan -cluster_on_pca: true -pca_components: 10 -params: - # `auto` = 0.5% of the instances being clustered. A fixed value does not - # transfer between cohorts: 100 is a sensible floor at ~20k instances and - # swallows every niche at 2k. Set an integer to override. - min_cluster_size: auto - min_samples: 1 - cluster_selection_method: eom - cluster_selection_epsilon: 0.0 - alpha: 1.0 - algorithm: auto - leaf_size: 40 - allow_single_cluster: false - n_jobs: null - metric: cosine diff --git a/src/grass_mil/configs/interpretability/clustering/kmeans.yaml b/src/grass_mil/configs/interpretability/clustering/kmeans.yaml deleted file mode 100644 index a04c88d..0000000 --- a/src/grass_mil/configs/interpretability/clustering/kmeans.yaml +++ /dev/null @@ -1,9 +0,0 @@ -enabled: true -method: kmeans -params: - n_clusters: 8 - random_state: 42 - init: k-means++ - n_init: auto - max_iter: 300 - tol: 0.0001 diff --git a/src/grass_mil/configs/interpretability/clustering/optics.yaml b/src/grass_mil/configs/interpretability/clustering/optics.yaml deleted file mode 100644 index aa7e3ec..0000000 --- a/src/grass_mil/configs/interpretability/clustering/optics.yaml +++ /dev/null @@ -1,13 +0,0 @@ -enabled: true -method: optics -params: - min_samples: 5 - max_eps: .inf - metric: minkowski - p: 2 - cluster_method: xi - xi: 0.05 - min_cluster_size: null - algorithm: auto - leaf_size: 30 - n_jobs: null diff --git a/src/grass_mil/configs/interpretability/clustering/spectral.yaml b/src/grass_mil/configs/interpretability/clustering/spectral.yaml deleted file mode 100644 index fe5ec15..0000000 --- a/src/grass_mil/configs/interpretability/clustering/spectral.yaml +++ /dev/null @@ -1,8 +0,0 @@ -enabled: true -method: spectral -params: - n_clusters: 8 - random_state: 42 - affinity: nearest_neighbors - n_neighbors: 10 - assign_labels: kmeans diff --git a/src/grass_mil/configs/interpretability/plugins/default.yaml b/src/grass_mil/configs/interpretability/plugins/default.yaml deleted file mode 100644 index bd83d7e..0000000 --- a/src/grass_mil/configs/interpretability/plugins/default.yaml +++ /dev/null @@ -1,72 +0,0 @@ -enabled: - - niche_profiles -params: - niche_profiles: - variance_estimator: unbiased - margin_attribution: - # Per-class attention columns, emitted when the head has one attention - # channel per class. Falls back to the single `attention` column. - attention_prefix: attention_c - logit_prefix: logit_ - # Head output bias per class, subtracted so the margin reflects the - # instance-driven part of the decision rather than the head's prior. - # Point logit_bias_checkpoint at the trained checkpoint and it is read from - # there; logit_bias only exists to override that by hand. - logit_bias_checkpoint: null - head_prefix: graph_head - logit_bias: null - n_bootstrap: 200 - random_state: 0 - # Classes to summarise. null summarises the positive class for binary - # heads (the class-0 margin is its exact negation) and every class for - # multi-class heads, each with a one-vs-rest margin. - focus_classes: null - niche_agreement: - active_name: active - # Columns holding partitions from other representation spaces, e.g. joined - # from a previous run's niche_labels.csv. - label_columns: [] - metrics: [ari, ami, nmi] - tissue_graph: - sample_column: sample_id - sample_value: null - id_column: instance_id - x_column: center_x - y_column: center_y - show_edges: true - node_size: 5.0 - edge_width: 0.5 - edge_opacity: 0.25 - reverse_y: true - # --- cell-level analyses (need interpretability.cells.enabled at predict) --- - cell_type_enrichment_per_niche: - cell_type_column: cell_type - niche_column: niche_label - # n_perms=0 uses the analytical null; >0 gives permutation z-scores. - n_perms: 0 - enrichment_mode: zscore - min_cells: 10 - skip_background: false - differential_vs_background: true - reference_niche: -1 - cell_type_moran_per_niche: - cell_type_column: cell_type - niche_column: niche_label - n_perms: 100 - min_nodes: 3 - niche_label_moran: - niche_column: niche_label - n_perms: 100 - cell_filtration_curves: - cell_type_column: cell_type - niche_column: niche_label - threshold_start: 0.0 - threshold_stop: 55.0 - threshold_count: 500 - cell_type_ripley_per_niche: - cell_type_column: cell_type - niche_column: niche_label - x_column: x - y_column: y - n_radii: 50 - min_count: 5 diff --git a/src/grass_mil/configs/interpretability/plugins/full.yaml b/src/grass_mil/configs/interpretability/plugins/full.yaml deleted file mode 100644 index 1ac0ca3..0000000 --- a/src/grass_mil/configs/interpretability/plugins/full.yaml +++ /dev/null @@ -1,24 +0,0 @@ -# Every builtin analysis enabled. -# -# The tier-2 analyses are cell-level: they measure the cells inside the sampled -# ego-graphs, which needs `cell_table`/`cell_edge_table` from -# `grass-mil-predict interpretability.cells.enabled=true`. `niche_label_moran` -# is the global counterpart and needs only `spatial_table`. -# -# Slower than the default preset: the per-niche statistics scale with niche size -# and permutation count. - -defaults: - - default - - _self_ - -enabled: - - niche_profiles - - margin_attribution - # cell-level, within each niche - - cell_type_enrichment_per_niche - - cell_type_moran_per_niche - - cell_filtration_curves - - cell_type_ripley_per_niche - # global, over niches - - niche_label_moran diff --git a/src/grass_mil/configs/interpretability/reduction/pca.yaml b/src/grass_mil/configs/interpretability/reduction/pca.yaml deleted file mode 100644 index caeb9b2..0000000 --- a/src/grass_mil/configs/interpretability/reduction/pca.yaml +++ /dev/null @@ -1,5 +0,0 @@ -enabled: true -method: pca -params: - n_components: 2 - random_state: 42 diff --git a/src/grass_mil/configs/interpretability/reduction/tsne.yaml b/src/grass_mil/configs/interpretability/reduction/tsne.yaml deleted file mode 100644 index 01e9544..0000000 --- a/src/grass_mil/configs/interpretability/reduction/tsne.yaml +++ /dev/null @@ -1,8 +0,0 @@ -enabled: true -method: tsne -params: - n_components: 2 - random_state: 42 - perplexity: 30.0 - learning_rate: auto - init: pca diff --git a/src/grass_mil/configs/interpretability/reduction/umap.yaml b/src/grass_mil/configs/interpretability/reduction/umap.yaml deleted file mode 100644 index e29300d..0000000 --- a/src/grass_mil/configs/interpretability/reduction/umap.yaml +++ /dev/null @@ -1,8 +0,0 @@ -enabled: true -method: umap -params: - n_components: 2 - random_state: null - n_neighbors: 50 - min_dist: 0.1 - metric: euclidean diff --git a/src/grass_mil/configs/interpretability/report.yaml b/src/grass_mil/configs/interpretability/report.yaml deleted file mode 100644 index 9e6a5ae..0000000 --- a/src/grass_mil/configs/interpretability/report.yaml +++ /dev/null @@ -1,41 +0,0 @@ -# @package _global_ - -defaults: - - _self_ - - /paths: default - - /extras: default - - reduction: umap - - clustering: hdbscan - - plugins: default - -task_name: "interpretability_report" -tags: ["dev"] - -output_dir: ${paths.output_dir}/interpretability_report -composition_prefix: comp_ -variance_estimator: unbiased -embedding_prefixes: - - inst_emb_ - - emb_ - - graph_emb_ - -data: - instance_table: ??? - bag_table: null - spatial_table: null - # Cell-level tables from `grass-mil-predict interpretability.cells.enabled=true`. - # Required by the analyses that measure cells rather than neighbourhoods. - cell_table: null - cell_edge_table: null - id_column: instance_id - bag_id_column: bag_id - cell_type_column: cell_type - condition_column: null - -report: - html: - enabled: true - pdf: - enabled: true - engine: playwright - snapshot_dpi_scale: 2.0 diff --git a/src/grass_mil/configs/local/.gitkeep b/src/grass_mil/configs/local/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/src/grass_mil/configs/logger/aim.yaml b/src/grass_mil/configs/logger/aim.yaml deleted file mode 100644 index 8f9f6ad..0000000 --- a/src/grass_mil/configs/logger/aim.yaml +++ /dev/null @@ -1,28 +0,0 @@ -# https://aimstack.io/ - -# example usage in lightning module: -# https://github.com/aimhubio/aim/blob/main/examples/pytorch_lightning_track.py - -# open the Aim UI with the following command (run in the folder containing the `.aim` folder): -# `aim up` - -aim: - _target_: aim.pytorch_lightning.AimLogger - repo: ${paths.root_dir} # .aim folder will be created here - # repo: "aim://ip_address:port" # can instead provide IP address pointing to Aim remote tracking server which manages the repo, see https://aimstack.readthedocs.io/en/latest/using/remote_tracking.html# - - # aim allows to group runs under experiment name - experiment: null # any string, set to "default" if not specified - - train_metric_prefix: "train/" - val_metric_prefix: "val/" - test_metric_prefix: "test/" - - # sets the tracking interval in seconds for system usage metrics (CPU, GPU, memory, etc.) - system_tracking_interval: 10 # set to null to disable system metrics tracking - - # enable/disable logging of system params such as installed packages, git info, env vars, etc. - log_system_params: true - - # enable/disable tracking console logs (default value is true) - capture_terminal_logs: false # set to false to avoid infinite console log loop issue https://github.com/aimhubio/aim/issues/2550 diff --git a/src/grass_mil/configs/logger/comet.yaml b/src/grass_mil/configs/logger/comet.yaml deleted file mode 100644 index f366bb0..0000000 --- a/src/grass_mil/configs/logger/comet.yaml +++ /dev/null @@ -1,12 +0,0 @@ -# https://www.comet.ml - -comet: - _target_: lightning.pytorch.loggers.comet.CometLogger - api_key: ${oc.env:COMET_API_TOKEN} # api key is loaded from environment variable - save_dir: "${paths.output_dir}" - project_name: "grass-mil" - rest_api_key: null - # experiment_name: "" - experiment_key: null # set to resume experiment - offline: False - prefix: "" diff --git a/src/grass_mil/configs/logger/csv.yaml b/src/grass_mil/configs/logger/csv.yaml index fa028e9..018a01f 100644 --- a/src/grass_mil/configs/logger/csv.yaml +++ b/src/grass_mil/configs/logger/csv.yaml @@ -1,7 +1,4 @@ -# csv logger built in lightning - csv: - _target_: lightning.pytorch.loggers.csv_logs.CSVLogger - save_dir: "${paths.output_dir}" - name: "csv/" - prefix: "" + _target_: lightning.pytorch.loggers.CSVLogger + save_dir: ${paths.output_dir} + name: csv diff --git a/src/grass_mil/configs/logger/many_loggers.yaml b/src/grass_mil/configs/logger/many_loggers.yaml deleted file mode 100644 index dd58680..0000000 --- a/src/grass_mil/configs/logger/many_loggers.yaml +++ /dev/null @@ -1,9 +0,0 @@ -# train with many loggers at once - -defaults: - # - comet - - csv - # - mlflow - # - neptune - - tensorboard - - wandb diff --git a/src/grass_mil/configs/logger/mlflow.yaml b/src/grass_mil/configs/logger/mlflow.yaml deleted file mode 100644 index f8fb7e6..0000000 --- a/src/grass_mil/configs/logger/mlflow.yaml +++ /dev/null @@ -1,12 +0,0 @@ -# https://mlflow.org - -mlflow: - _target_: lightning.pytorch.loggers.mlflow.MLFlowLogger - # experiment_name: "" - # run_name: "" - tracking_uri: ${paths.log_dir}/mlflow/mlruns # run `mlflow ui` command inside the `logs/mlflow/` dir to open the UI - tags: null - # save_dir: "./mlruns" - prefix: "" - artifact_location: null - # run_id: "" diff --git a/src/grass_mil/configs/logger/neptune.yaml b/src/grass_mil/configs/logger/neptune.yaml deleted file mode 100644 index 9c7881f..0000000 --- a/src/grass_mil/configs/logger/neptune.yaml +++ /dev/null @@ -1,9 +0,0 @@ -# https://neptune.ai - -neptune: - _target_: lightning.pytorch.loggers.neptune.NeptuneLogger - api_key: ${oc.env:NEPTUNE_API_TOKEN} # api key is loaded from environment variable - project: org/grass-mil - # name: "" - log_model_checkpoints: True - prefix: "" diff --git a/src/grass_mil/configs/logger/tensorboard.yaml b/src/grass_mil/configs/logger/tensorboard.yaml index 2bd31f6..b18f551 100644 --- a/src/grass_mil/configs/logger/tensorboard.yaml +++ b/src/grass_mil/configs/logger/tensorboard.yaml @@ -1,10 +1,4 @@ -# https://www.tensorflow.org/tensorboard/ - tensorboard: - _target_: lightning.pytorch.loggers.tensorboard.TensorBoardLogger - save_dir: "${paths.output_dir}/tensorboard/" + _target_: lightning.pytorch.loggers.TensorBoardLogger + save_dir: ${paths.output_dir}/tensorboard name: null - log_graph: False - default_hp_metric: True - prefix: "" - # version: "" diff --git a/src/grass_mil/configs/logger/wandb.yaml b/src/grass_mil/configs/logger/wandb.yaml deleted file mode 100644 index 095dd9a..0000000 --- a/src/grass_mil/configs/logger/wandb.yaml +++ /dev/null @@ -1,16 +0,0 @@ -# https://wandb.ai - -wandb: - _target_: lightning.pytorch.loggers.wandb.WandbLogger - # name: "" # name of the run (normally generated by wandb) - save_dir: "${paths.output_dir}" - offline: False - id: null # pass correct id to resume experiment! - anonymous: null # enable anonymous logging - project: "grass-mil" - log_model: False # upload lightning ckpts - prefix: "" # a string to put at the beginning of metric keys - # entity: "" # set to name of your wandb team - group: "" - tags: [] - job_type: "" diff --git a/src/grass_mil/configs/loocv.yaml b/src/grass_mil/configs/loocv.yaml deleted file mode 100644 index 288e942..0000000 --- a/src/grass_mil/configs/loocv.yaml +++ /dev/null @@ -1,14 +0,0 @@ -# @package _global_ - -defaults: - - train - - _self_ - -task_name: "loocv" - -loocv: - mode: single # single | all - fold_index: null # single mode selector; takes priority over data.split.loocv.holdout_id - per_fold_processed_dir: true - force_precompute_per_fold: true - summary_filename: loocv_summary.json diff --git a/src/grass_mil/configs/model/attention/gated.yaml b/src/grass_mil/configs/model/attention/gated.yaml deleted file mode 100644 index 695305f..0000000 --- a/src/grass_mil/configs/model/attention/gated.yaml +++ /dev/null @@ -1,7 +0,0 @@ -# Plain gated attention (Ilse et al. 2018), one attention channel per class. -attention_type: gated -input_dim: ${model.encoder.out_dim} -hidden_dim: 64 -dropout: false -n_classes: 1 -init: kaiming diff --git a/src/grass_mil/configs/model/attention/gated_projected.yaml b/src/grass_mil/configs/model/attention/gated_projected.yaml deleted file mode 100644 index ce1365d..0000000 --- a/src/grass_mil/configs/model/attention/gated_projected.yaml +++ /dev/null @@ -1,10 +0,0 @@ -# Projected gated attention (GRASS-MIL default): Linear + LeakyReLU projection, -# tanh/sigmoid gate, one attention channel per class. `n_classes` must equal -# graph_head.output_dim; the task preset sets both together. -attention_type: gated_projected -input_dim: ${model.encoder.out_dim} -projection_dim: 32 -hidden_dim: 8 -dropout: true -n_classes: 1 -init: kaiming diff --git a/src/grass_mil/configs/model/bgrl.yaml b/src/grass_mil/configs/model/bgrl.yaml new file mode 100644 index 0000000..268cf87 --- /dev/null +++ b/src/grass_mil/configs/model/bgrl.yaml @@ -0,0 +1,15 @@ +# BGRL pretraining of the ego-graph encoder (selected by regime=pretrain). +defaults: + - encoder: gine + - _self_ + +_target_: grass_mil.models.BGRLModule +_recursive_: false + +encoder: + dropout: 0.1 +pretrain: ${pretrain} +optim: ${optim} +scheduler: ${scheduler} +# Continue from an earlier encoder checkpoint; null = random initialisation. +init_from_ckpt: null diff --git a/src/grass_mil/configs/model/bgrl_a3.yaml b/src/grass_mil/configs/model/bgrl_a3.yaml deleted file mode 100644 index 30b179d..0000000 --- a/src/grass_mil/configs/model/bgrl_a3.yaml +++ /dev/null @@ -1,28 +0,0 @@ -defaults: - - encoder: gine_length - - _self_ - -# BGRL pretraining of the GRASS-MIL encoder (preprint "A3" recipe): the same -# GINE ego-graph encoder as model/grass_mil, encoder dropout 0.1 during -# pretraining only, predictor Linear(256->512) -> PReLU -> Linear(512->256). -# Use with `task=pretrain_bgrl_a3` (experiment=grass_mil_pretrain_a3 bundles -# model, task, sampler and callbacks). Every key can be overridden. - -_target_: grass_mil.models.bgrl_module.BGRLModule -_recursive_: false - -encoder: - dropout: 0.1 - -ssl: - method: bgrl - predictor: - hidden_size: 512 - -optim: ${oc.select:optim,null} -scheduler: ${oc.select:scheduler,null} -task: ${task} - -init_from_ckpt: null -init_strict: false -encoder_init_map: auto_bgrl_or_identity diff --git a/src/grass_mil/configs/model/bgrl_module.yaml b/src/grass_mil/configs/model/bgrl_module.yaml deleted file mode 100644 index 300c4d5..0000000 --- a/src/grass_mil/configs/model/bgrl_module.yaml +++ /dev/null @@ -1,22 +0,0 @@ -defaults: - - encoder: gin - - _self_ - -_target_: grass_mil.models.bgrl_module.BGRLModule -_recursive_: false - -# `encoder` comes from the model/encoder group (override with -# `model/encoder=gine_length`). - -ssl: - method: bgrl - predictor: - hidden_size: 512 - -optim: ${oc.select:optim,null} -scheduler: ${oc.select:scheduler,null} -task: ${task} - -init_from_ckpt: null -init_strict: false -encoder_init_map: auto_bgrl_or_identity diff --git a/src/grass_mil/configs/model/encoder/gat.yaml b/src/grass_mil/configs/model/encoder/gat.yaml deleted file mode 100644 index ca9dbda..0000000 --- a/src/grass_mil/configs/model/encoder/gat.yaml +++ /dev/null @@ -1,18 +0,0 @@ -input_dim: 0 -hidden_dim: 128 -out_dim: 128 -num_layers: 3 -dropout: 0.1 -conv_type: gat -norm: batchnorm -jk: last -act: relu -pooling: mean -gat_heads: 4 -use_edge_attr: false -edge_weight_index: 0 -edge_attr_dim: null -edge_feature_index: null -categorical_embedding: null -train_eps: true -input_init: xavier diff --git a/src/grass_mil/configs/model/encoder/gcn.yaml b/src/grass_mil/configs/model/encoder/gcn.yaml deleted file mode 100644 index 7ceb192..0000000 --- a/src/grass_mil/configs/model/encoder/gcn.yaml +++ /dev/null @@ -1,18 +0,0 @@ -input_dim: 0 -hidden_dim: 128 -out_dim: 128 -num_layers: 3 -dropout: 0.1 -conv_type: gcn -norm: batchnorm -jk: last -act: relu -pooling: mean -gat_heads: 4 -use_edge_attr: true -edge_weight_index: 0 -edge_attr_dim: null -edge_feature_index: null -categorical_embedding: null -train_eps: true -input_init: xavier diff --git a/src/grass_mil/configs/model/encoder/gin.yaml b/src/grass_mil/configs/model/encoder/gin.yaml deleted file mode 100644 index 946517e..0000000 --- a/src/grass_mil/configs/model/encoder/gin.yaml +++ /dev/null @@ -1,20 +0,0 @@ -# Small GIN encoder without edge features: a lightweight default for smoke runs. -# `input_dim: 0` lets the datamodule report the continuous node-feature width. -input_dim: 0 -hidden_dim: 128 -out_dim: 128 -num_layers: 3 -dropout: 0.1 -conv_type: gin -norm: batchnorm -jk: last -act: relu -pooling: mean -gat_heads: 4 -use_edge_attr: false -edge_weight_index: 0 -edge_attr_dim: null -edge_feature_index: null -categorical_embedding: null -train_eps: true -input_init: xavier diff --git a/src/grass_mil/configs/model/encoder/gine.yaml b/src/grass_mil/configs/model/encoder/gine.yaml index e4996c7..2204e46 100644 --- a/src/grass_mil/configs/model/encoder/gine.yaml +++ b/src/grass_mil/configs/model/encoder/gine.yaml @@ -1,18 +1,20 @@ -input_dim: 0 -hidden_dim: 128 -out_dim: 128 -num_layers: 3 -dropout: 0.1 -conv_type: gine -norm: batchnorm -jk: last -act: relu -pooling: mean -gat_heads: 4 -use_edge_attr: true -edge_weight_index: 0 -edge_attr_dim: 2 -edge_feature_index: null -categorical_embedding: null +# Ego-graph encoder: cell-type embedding + projected continuous features, +# GINE layers conditioned on one edge feature, max pooling over the ego-graph. +# Input widths (features, cell types, edge features) come from the data. +hidden_dim: 256 +out_dim: 256 +num_layers: 4 # keep equal to data.sampler.depth +dropout: 0.0 +norm: layernorm # layernorm | batchnorm | none +jk: last # last | concat | max | sum +act: relu # relu | gelu | leaky_relu +pooling: max # max | mean | sum +mlp_ratio: 2 # GINE MLP width = mlp_ratio * hidden_dim train_eps: true -input_init: xavier +# Column of data.graph.edge_features seen by the convolutions (0 = the first, +# log_distance by default); null = all columns. +edge_feature: 0 +input_init: xavier # xavier | default +# Cell type standing in for a masked node during augmentation; null reserves an +# extra embedding row instead. +unassigned_cell_type: null diff --git a/src/grass_mil/configs/model/encoder/gine_length.yaml b/src/grass_mil/configs/model/encoder/gine_length.yaml deleted file mode 100644 index 5b468c7..0000000 --- a/src/grass_mil/configs/model/encoder/gine_length.yaml +++ /dev/null @@ -1,45 +0,0 @@ -# GRASS-MIL encoder (preprint defaults): 4 GINE layers of width 256 conditioned -# on the bounded log edge length, layernorm, max pooling over the ego-graph, -# cell identity through a learned embedding summed with the projection of the -# continuous node features. Every value is a plain Hydra key and can be -# overridden (e.g. `model.encoder.num_layers=3 model.encoder.hidden_dim=128`). -# -# `input_dim: 0` lets the datamodule report the continuous node-feature width -# (0 for cell-type-only cohorts, 1 when cell size is present). Likewise -# `num_embeddings: null` and `column_index: null` are resolved from the -# processed dataset's label map and categorical slices. -input_dim: 0 -hidden_dim: 256 -out_dim: 256 -num_layers: 4 -dropout: 0.0 -conv_type: gine -norm: layernorm -jk: last -act: relu -pooling: max -gat_heads: 4 - -# Message passing reads exactly one edge column. With the graph builder's -# `edge_features: [log_distance, neighbor]` (preprint), index 0 is the bounded -# log edge length clip((ln d - 2)/3, 0, 1). -use_edge_attr: true -edge_weight_index: 0 -edge_attr_dim: 1 -edge_feature_index: 0 - -# Learnable GIN epsilon and xavier-initialised input embeddings, as in the preprint. -train_eps: true -input_init: xavier - -categorical_embedding: - label: cell_type - num_embeddings: null - column_index: null - # Reserve a trailing row so augmentations can mark a dropped node - # "unassigned" without resizing the table. Alternatively set - # `reserve_unassigned: false` and `unassigned_label: other` to reuse an - # existing category as the unassigned token. - reserve_unassigned: true - unassigned_label: null - unassigned_index: null diff --git a/src/grass_mil/configs/model/encoder/graphsage.yaml b/src/grass_mil/configs/model/encoder/graphsage.yaml deleted file mode 100644 index 12fcdc9..0000000 --- a/src/grass_mil/configs/model/encoder/graphsage.yaml +++ /dev/null @@ -1,18 +0,0 @@ -input_dim: 0 -hidden_dim: 128 -out_dim: 128 -num_layers: 3 -dropout: 0.1 -conv_type: graphsage -norm: batchnorm -jk: last -act: relu -pooling: mean -gat_heads: 4 -use_edge_attr: false -edge_weight_index: 0 -edge_attr_dim: null -edge_feature_index: null -categorical_embedding: null -train_eps: true -input_init: xavier diff --git a/src/grass_mil/configs/model/grass_mil.yaml b/src/grass_mil/configs/model/grass_mil.yaml index 1c61989..30085da 100644 --- a/src/grass_mil/configs/model/grass_mil.yaml +++ b/src/grass_mil/configs/model/grass_mil.yaml @@ -1,45 +1,25 @@ +# GRASS-MIL: ego-graph encoder, linear (or MLP) instance head and projected +# gated attention with one channel per class, pooled in logit space. defaults: - - encoder: gine_length - - attention: gated_projected + - encoder: gine - _self_ -# GRASS-MIL supervised module with the preprint defaults: BGRL-pretrainable GINE -# ego-graph encoder (model/encoder/gine_length), projected gated attention with -# one channel per class (model/attention/gated_projected), and an MLP instance -# head pooled in logit space (bag logit_c = sum_i A[i,c] * f_c(h_i)), so each -# prediction decomposes exactly into per-neighbourhood contributions. -# -# Use with `model=grass_mil task=finetune_mil` (and a `regime=...` preset). Any -# key can be overridden, e.g. `model.encoder.num_layers=3`, -# `model.graph_head.num_layers=1` (linear head for the probe regimes). - -_target_: grass_mil.models.supervised_module.SupervisedModule +_target_: grass_mil.models.SupervisedModule _recursive_: false -flags: - use_ssl: false - -graph_head: - input_dim: ${model.encoder.out_dim} - output_dim: 1 - hidden_dim: ${model.encoder.out_dim} - # 2 = MLP instance head (default); 1 = linear head for grass-mil-probe regimes. - num_layers: 2 +head: + num_layers: 1 # 1 = linear head; 2 = MLP + hidden_dim: null # MLP hidden width (null = embedding width) dropout: 0.0 - -ssl: - method: bgrl - predictor: - hidden_size: 512 - -loss: - loss_type: ${task.loss} - -optim: ${oc.select:optim,null} -scheduler: ${oc.select:scheduler,null} + act: leaky_relu +attention: + projection_dim: 32 + hidden_dim: 8 + dropout: 0.25 + init: kaiming # kaiming | default task: ${task} - +finetune: ${finetune} +optim: ${optim} +scheduler: ${scheduler} +# Pretrained encoder checkpoint (from regime=pretrain); null = random encoder. init_from_ckpt: null -init_strict: false -encoder_init_map: auto_bgrl_or_identity -freeze_encoder: false diff --git a/src/grass_mil/configs/model/supervised_module.yaml b/src/grass_mil/configs/model/supervised_module.yaml deleted file mode 100644 index 1f6d0e2..0000000 --- a/src/grass_mil/configs/model/supervised_module.yaml +++ /dev/null @@ -1,40 +0,0 @@ -defaults: - - encoder: gin - - attention: gated_projected - - _self_ - -_target_: grass_mil.models.supervised_module.SupervisedModule -_recursive_: false - -flags: - use_ssl: false - -# `encoder` and `attention` come from the model/encoder and model/attention -# groups (override with `model/encoder=gine_length`, `model/attention=gated`). -# The attention block is only used when task.aggregation == mil_attention. - -graph_head: - input_dim: ${model.encoder.out_dim} - output_dim: 1 - hidden_dim: ${model.encoder.out_dim} - # 2 = MLP instance head (default); 1 = linear head, required by the probe - # pre-fit regimes (grass-mil-probe folds a logistic regression into it). - num_layers: 2 - dropout: 0.0 - -ssl: - method: bgrl - predictor: - hidden_size: 512 - -loss: - loss_type: ${task.loss} - -optim: ${oc.select:optim,null} -scheduler: ${oc.select:scheduler,null} -task: ${task} - -init_from_ckpt: null -init_strict: false -encoder_init_map: auto_bgrl_or_identity -freeze_encoder: false diff --git a/src/grass_mil/configs/nested_cv.yaml b/src/grass_mil/configs/nested_cv.yaml index 9643f40..0729017 100644 --- a/src/grass_mil/configs/nested_cv.yaml +++ b/src/grass_mil/configs/nested_cv.yaml @@ -1,55 +1,24 @@ -# @package _global_ -# grass-mil-nested-cv: the preprint's nested, patient-grouped cross-validation. +# grass-mil-nested-cv: nested, patient-grouped cross-validation. # -# grass-mil-nested-cv regime=linear_ftall_reg data=... \ -# model.init_from_ckpt=/path/to/bgrl.ckpt \ -# nested_cv.folds_file=folds.json \ -# data.graph_labels.label_file=labels.csv data.graph_labels.id_column=REGION_ID \ -# data.graph_labels.tasks=[STAS] data.graph_labels.patient_column=PATIENT +# grass-mil-nested-cv experiment=my_cohort model.init_from_ckpt=/checkpoints/last.ckpt # -# For every outer fold: the remaining folds form the development set, which is -# split into `inner_k` stratified, patient-grouped inner folds; each inner run -# trains with the selected regime and records the epoch with the best -# `selection_metric`; the median of those epochs is the refit length; the model -# is refitted on the whole development set for that many epochs (no early -# stopping, last weights) and scored once on the held-out fold, on all cells. -# The linear_* regimes fit their probe head on each training split first -# (cached under `probe.cache_dir`, so a second regime reuses them). +# For every outer fold, the other folds form the development set. It is split +# into `inner_k` group-aware inner folds (stratified for classification); each +# inner run trains with the selected regime and records the epoch with the best +# `task.selection.metric`. The model is then refitted on the whole development +# set for the reduced number of epochs (no validation, last weights) and scored +# once on the held-out fold, on every cell. defaults: - train - _self_ -task_name: "nested_cv" +task_name: nested_cv nested_cv: - folds_file: ??? # JSON list of lists of unit ids (patient-grouped outer folds) - # What one unit (bag, label, fold member) is: a region (one graph per region, - # data.sample_unit=full), a sample, or a patch (tiles as units; sets - # task.bag_key=patch_id and data.graph_labels.scope=patch). - id_kind: region - label_file: ${data.graph_labels.label_file} - id_column: ${data.graph_labels.id_column} - label_column: null # null = task.target_columns[0] or data.graph_labels.tasks[0] - patient_column: ${data.graph_labels.patient_column} - outer_folds: null # e.g. [2] to run one outer fold; null = all + folds_file: ??? # JSON list of lists of unit (or region / sample) ids inner_k: 3 inner_seed: 0 - selection_metric: val/auroc - selection_mode: max # max | min - epoch_reduce: median # median | mean | max (over the inner folds' best epochs) + epoch_reduce: median # median | mean | max, over the inner best epochs min_refit_epochs: 1 - resume: true # skip inner runs / outer folds whose results already exist - probe: - # auto = fit a probe wherever the regime asks for one (task.head_init_path - # left unset, as in the linear_* regimes); true = always; false = never. - enabled: auto - cache_dir: ${paths.output_dir}/probes - roots_per_unit: 96 - seed: 0 - batch_size: 256 - C: 1.0 - class_weight: balanced - max_iter: 3000 - device: auto - patient_aggregation: mean # mean of the region probabilities per patient - summary_filename: nested_cv_summary.json + outer_folds: null # e.g. [2] to run one outer fold; null = all + resume: true # reuse finished inner runs and outer folds diff --git a/src/grass_mil/configs/optim/adamw.yaml b/src/grass_mil/configs/optim/adamw.yaml index cd34fc1..a313ab7 100644 --- a/src/grass_mil/configs/optim/adamw.yaml +++ b/src/grass_mil/configs/optim/adamw.yaml @@ -1,3 +1,3 @@ _target_: torch.optim.AdamW -lr: 0.001 -weight_decay: 0.01 +lr: 1.0e-4 +weight_decay: 1.0e-3 diff --git a/src/grass_mil/configs/paths/default.yaml b/src/grass_mil/configs/paths/default.yaml index 45ddc6d..b31d42b 100644 --- a/src/grass_mil/configs/paths/default.yaml +++ b/src/grass_mil/configs/paths/default.yaml @@ -1,20 +1,6 @@ -# path to root directory -# resolved from the PROJECT_ROOT environment variable when set (rootutils sets -# it from the .project-root marker in a source checkout), otherwise the Hydra -# runtime working directory, so an installed package works without setup. -# you can replace it with "." if you want the root to be the current working directory -root_dir: ${oc.env:PROJECT_ROOT,${hydra:runtime.cwd}} - -# path to data directory -data_dir: ${paths.root_dir}/data/ - -# path to logging directory -log_dir: ${paths.root_dir}/logs/ - -# path to output directory, created dynamically by hydra -# path generation pattern is specified in `configs/hydra/default.yaml` -# use it to store all files generated during the run, like ckpts and metrics +# Relative paths resolve against the directory the command is run from. +root_dir: ${hydra:runtime.cwd} +data_dir: ${paths.root_dir}/data +log_dir: ${paths.root_dir}/logs +# Created per run by Hydra (see hydra/default.yaml). output_dir: ${hydra:runtime.output_dir} - -# path to working directory -work_dir: ${hydra:runtime.cwd} diff --git a/src/grass_mil/configs/predict.yaml b/src/grass_mil/configs/predict.yaml new file mode 100644 index 0000000..48d5460 --- /dev/null +++ b/src/grass_mil/configs/predict.yaml @@ -0,0 +1,21 @@ +# grass-mil-predict: score a cohort with a trained checkpoint. +# +# grass-mil-predict experiment=new_cohort ckpt_path=/checkpoints/last.ckpt +# +# The model and its ego-graph settings come from the checkpoint; `data` points +# at the cohort to score (labels optional; with labels, metrics are reported). +defaults: + - train + - _self_ + +task_name: predict +ckpt_path: ??? + +predict: + units: null # subset of unit ids (list or file); null = every unit + # Patient scores (needs data.labels.patient_column): `mean` averages the + # unit predictions; `merged_bag` pools every cell of the patient's units in + # one attention bag. + patient_aggregation: mean # mean | merged_bag + batch_size: 512 + device: auto diff --git a/src/grass_mil/configs/probe.yaml b/src/grass_mil/configs/probe.yaml deleted file mode 100644 index cdaaffa..0000000 --- a/src/grass_mil/configs/probe.yaml +++ /dev/null @@ -1,29 +0,0 @@ -# @package _global_ -# grass-mil-probe: fit the linear probe head on a frozen encoder. -# -# grass-mil-probe regime=linear_probehead data=... \ -# model.init_from_ckpt=/path/to/bgrl.ckpt \ -# data.split.explicit.train=train_ids.json data.split.explicit.val=val_ids.json -# -# The encoder is taken from `model` (with `init_from_ckpt`) and never trained; -# the ego-graph settings come from `data.sampler.runtime` (depth, neighbours, -# radius cap). The result is written as `graph_head.net.0.{weight,bias}` and is -# consumed through `task.head_init_path` by the linear_* regimes. -defaults: - - train - - _self_ - -task_name: "probe" -train: false -test: false - -probe: - roots_per_unit: 96 # ego-graphs per region used to fit the probe (null = all roots) - seed: 0 # fixed per-region root subset - batch_size: 256 # ego-graphs per forward pass - C: 1.0 # inverse L2 strength of the logistic regression - class_weight: balanced # balanced | null - max_iter: 3000 - eval_splits: [val] # splits scored after the fit (instance and region-mean AUROC) - output_path: ${paths.output_dir}/probe_head.pt - device: auto # auto | cpu | cuda diff --git a/src/grass_mil/configs/regime/attention_only.yaml b/src/grass_mil/configs/regime/attention_only.yaml new file mode 100644 index 0000000..da0188d --- /dev/null +++ b/src/grass_mil/configs/regime/attention_only.yaml @@ -0,0 +1,16 @@ +# @package _global_ +# Frozen encoder and frozen probe head: only the attention is trained, starting +# from uniform attention (the mean-pooled probe). +defaults: + - finetune + +regime_name: attention_only + +finetune: + lr_scales: + encoder: 0.0 + head: 0.0 + freeze_head: true +data: + sampler: + roots_per_unit: 512 diff --git a/src/grass_mil/configs/regime/finetune.yaml b/src/grass_mil/configs/regime/finetune.yaml new file mode 100644 index 0000000..02c7849 --- /dev/null +++ b/src/grass_mil/configs/regime/finetune.yaml @@ -0,0 +1,39 @@ +# @package _global_ +# Default regime: the BGRL-pretrained encoder is fine-tuned end to end; the +# linear instance head starts from a logistic probe on the frozen embeddings +# and learns slowly; the attention learns at twice the base rate. +regime_name: finetune + +finetune: + lr_scales: + encoder: 1.0 + head: 0.01 + attention: 2.0 + head_weight_decay: 1.0e-3 # null = optim.weight_decay + head_init: probe # probe | random + freeze_head: false + hyperbatch_size: 12 # units per optimizer step + grad_clip_norm: 1.0 + warmup_steps: 0 + augmentation: + drop_edge_p: 0.0 + mask_node_p: 0.0 + seed: 0 + probe: + roots_per_unit: 96 + seed: 0 + C: 1.0 + class_weight: balanced + max_iter: 3000 + +model: + head: + num_layers: 1 + +data: + sampler: + roots_per_unit: 256 + +trainer: + max_epochs: 25 + min_epochs: 10 diff --git a/src/grass_mil/configs/regime/finetune_reg.yaml b/src/grass_mil/configs/regime/finetune_reg.yaml new file mode 100644 index 0000000..d12ebb3 --- /dev/null +++ b/src/grass_mil/configs/regime/finetune_reg.yaml @@ -0,0 +1,19 @@ +# @package _global_ +# finetune with a slower encoder and regularisation (weight decay, encoder +# dropout, edge drop and node masking). +defaults: + - finetune + +regime_name: finetune_reg + +optim: + weight_decay: 1.0e-2 +finetune: + lr_scales: + encoder: 0.3 + augmentation: + drop_edge_p: 0.05 + mask_node_p: 0.05 +model: + encoder: + dropout: 0.2 diff --git a/src/grass_mil/configs/regime/linear_ftall.yaml b/src/grass_mil/configs/regime/linear_ftall.yaml deleted file mode 100644 index bbea138..0000000 --- a/src/grass_mil/configs/regime/linear_ftall.yaml +++ /dev/null @@ -1,28 +0,0 @@ -# @package _global_ -# Preprint "ftALL": encoder fine-tuned end to end (lr x1.0), linear head -# initialised from the probe pre-fit and trained slowly (lr x0.01, wd 1e-3), -# attention at lr x2.0. -defaults: - - mil - -regime_name: linear_ftall - -task: - # "auto" = fitted per split by grass-mil-nested-cv; for a single grass-mil-train - # run pass the output of grass-mil-probe (task.head_init_path=/path/probe_head.pt). - head_init_path: auto - freeze_head: false - optimization: - encoder_lr_scale: 1.0 - head_lr_scale: 0.01 - attention_lr_scale: 2.0 - head_weight_decay: 1.0e-3 - -model: - graph_head: - num_layers: 1 - -data: - sampler: - runtime: - subsample_count: 256 diff --git a/src/grass_mil/configs/regime/linear_ftall_reg.yaml b/src/grass_mil/configs/regime/linear_ftall_reg.yaml deleted file mode 100644 index 9233540..0000000 --- a/src/grass_mil/configs/regime/linear_ftall_reg.yaml +++ /dev/null @@ -1,21 +0,0 @@ -# @package _global_ -# Preprint "ftALL, lr x0.3 + regularisation" (the STAS model of record): -# encoder lr x0.3, weight decay 1e-2, encoder dropout 0.2, edge/feature drop 0.05. -defaults: - - linear_ftall - -regime_name: linear_ftall_reg - -optim: - weight_decay: 1.0e-2 - -task: - optimization: - encoder_lr_scale: 0.3 - finetune_augmentation: - drop_edge_p: 0.05 - drop_feat_p: 0.05 - -model: - encoder: - dropout: 0.2 diff --git a/src/grass_mil/configs/regime/linear_probehead.yaml b/src/grass_mil/configs/regime/linear_probehead.yaml deleted file mode 100644 index 63887c8..0000000 --- a/src/grass_mil/configs/regime/linear_probehead.yaml +++ /dev/null @@ -1,27 +0,0 @@ -# @package _global_ -# Preprint "probehead": frozen encoder, LINEAR head pre-fitted by grass-mil-probe -# (task.head_init_path) and frozen, so the attention is the only trainable part. -# Its uniform-attention starting point is exactly the mean-pooled probe. -defaults: - - mil - -regime_name: linear_probehead - -task: - # "auto" = fitted per split by grass-mil-nested-cv; for a single grass-mil-train - # run pass the output of grass-mil-probe (task.head_init_path=/path/probe_head.pt). - head_init_path: auto - freeze_head: true - optimization: - encoder_lr_scale: 0.0 - head_lr_scale: 0.0 - attention_lr_scale: 2.0 - -model: - graph_head: - num_layers: 1 - -data: - sampler: - runtime: - subsample_count: 512 diff --git a/src/grass_mil/configs/regime/mil.yaml b/src/grass_mil/configs/regime/mil.yaml deleted file mode 100644 index fc2182d..0000000 --- a/src/grass_mil/configs/regime/mil.yaml +++ /dev/null @@ -1,65 +0,0 @@ -# @package _global_ -# GRASS-MIL default regime: BGRL-initialised encoder fine-tuned end to end, -# MLP instance head, projected gated attention, one region per bag. -# Every value below is a plain override and can be changed on the command line -# (e.g. `regime=mil task.region_accumulation.hyperbatch_size=8`). -defaults: - - override /model: grass_mil - - override /task: finetune_mil - - override /scheduler: cosine_mil_step - - override /callbacks: mil_selection - - override /data/sampler: ego_graph_mil - -regime_name: mil - -optim: - lr: 1.0e-4 - weight_decay: 1.0e-3 - -task: - optimization: - encoder_lr_scale: 1.0 - head_lr_scale: 1.0 - attention_lr_scale: 2.0 - head_weight_decay: null - finetune_augmentation: - drop_edge_p: 0.0 - drop_feat_p: 0.0 - region_accumulation: - hyperbatch_size: 12 - grad_clip_norm: 1.0 - -model: - encoder: - dropout: 0.0 - graph_head: - num_layers: 2 - -data: - sampler: - runtime: - subsample_count: 256 - # Validation on a FIXED 1024-root subset per region, so epoch-to-epoch metric - # differences reflect the model rather than the root sample. - val_sampler: - name: shadow_custom - kwargs: {} - runtime: - enabled: true - depth: ${data.sampler.runtime.depth} - num_neighbors: ${data.sampler.runtime.num_neighbors} - subgraph_batch_size: 64 - replace: false - shuffle_subgraphs: false - proportional_root_sampling: false - subsample_count: 1024 - fixed_per_unit_seed: 0 - unit_as_bag: false - merge_units: false - interior_seeds: ${data.sampler.runtime.interior_seeds} - ego_radius: ${data.sampler.runtime.ego_radius} - -trainer: - max_epochs: 25 - min_epochs: 10 - precision: 16-mixed diff --git a/src/grass_mil/configs/regime/mil_frozen.yaml b/src/grass_mil/configs/regime/mil_frozen.yaml deleted file mode 100644 index afde9f8..0000000 --- a/src/grass_mil/configs/regime/mil_frozen.yaml +++ /dev/null @@ -1,17 +0,0 @@ -# @package _global_ -# Frozen encoder (learning-rate scale 0, kept in eval mode); MLP head and -# attention trained on the pretrained embeddings. Larger bags fit in memory -# because no activations are stored for the encoder. -defaults: - - mil - -regime_name: mil_frozen - -task: - optimization: - encoder_lr_scale: 0.0 - -data: - sampler: - runtime: - subsample_count: 512 diff --git a/src/grass_mil/configs/regime/mlp_head.yaml b/src/grass_mil/configs/regime/mlp_head.yaml new file mode 100644 index 0000000..9741433 --- /dev/null +++ b/src/grass_mil/configs/regime/mlp_head.yaml @@ -0,0 +1,15 @@ +# @package _global_ +# Encoder fine-tuned end to end with a randomly initialised two-layer MLP head. +defaults: + - finetune + +regime_name: mlp_head + +finetune: + lr_scales: + head: 1.0 + head_weight_decay: null + head_init: random +model: + head: + num_layers: 2 diff --git a/src/grass_mil/configs/regime/mlp_head_frozen.yaml b/src/grass_mil/configs/regime/mlp_head_frozen.yaml new file mode 100644 index 0000000..d36a21c --- /dev/null +++ b/src/grass_mil/configs/regime/mlp_head_frozen.yaml @@ -0,0 +1,13 @@ +# @package _global_ +# Frozen encoder; MLP head and attention trained on the pretrained embeddings. +defaults: + - mlp_head + +regime_name: mlp_head_frozen + +finetune: + lr_scales: + encoder: 0.0 +data: + sampler: + roots_per_unit: 512 diff --git a/src/grass_mil/configs/regime/pretrain.yaml b/src/grass_mil/configs/regime/pretrain.yaml new file mode 100644 index 0000000..088662e --- /dev/null +++ b/src/grass_mil/configs/regime/pretrain.yaml @@ -0,0 +1,39 @@ +# @package _global_ +# Label-free BGRL pretraining of the encoder on every unit of the cohort: +# grass-mil-train experiment=my_cohort regime=pretrain trainer=gpu +defaults: + - override /model: bgrl + - override /callbacks: pretrain_checkpoint + - override /data/sampler: ego_graph_pretrain + +regime_name: pretrain +task_name: pretrain +test: false + +pretrain: + total_steps: 200000 + warmup_steps: 1000 + momentum_start: 0.99 + momentum_end: 1.0 + predictor_hidden_dim: 512 + seed: 0 + augmentation: + mu: 0.25 # node/edge drop scale + p_lambda: 0.45 # maximum drop probability + noise_std: 0.2 + noise_features: [] # continuous features to perturb, by name (e.g. [area]) + +optim: + lr: 5.0e-5 + weight_decay: 1.0e-5 + +data: + labelled_only: false + split: + ratios: [1.0, 0.0, 0.0] + +trainer: + max_steps: ${pretrain.total_steps} + max_epochs: -1 + min_epochs: null + limit_val_batches: 0 diff --git a/src/grass_mil/configs/regime/pretrain_a3.yaml b/src/grass_mil/configs/regime/pretrain_a3.yaml deleted file mode 100644 index a7b78ab..0000000 --- a/src/grass_mil/configs/regime/pretrain_a3.yaml +++ /dev/null @@ -1,22 +0,0 @@ -# @package _global_ -# GRASS-MIL encoder pretraining (preprint "A3"): BGRL on the ego-graphs of -# every region, label-free. Composes with a cohort file: -# grass-mil-train experiment=my_cohort regime=pretrain_a3 paths.output_dir=runs/pretrain_a3 trainer=gpu -# Every value is overridable, e.g. task.total_steps=50000 model.encoder.num_layers=3 -# data.sampler.runtime.depth=3 task.augmentation.feature_noise_columns=[]. -defaults: - - override /model: bgrl_a3 - - override /task: pretrain_bgrl_a3 - - override /callbacks: pretrain_checkpoint - - override /data/sampler: ego_graph_pretrain - -regime_name: pretrain_a3 -task_name: pretrain_a3 -tags: ["grass_mil", "pretrain", "a3"] -test: false - -data: - split: - train_val_test_split: [1.0, 0.0, 0.0] - split_by: region - val_sampler: null diff --git a/src/grass_mil/configs/scheduler/cosine.yaml b/src/grass_mil/configs/scheduler/cosine.yaml new file mode 100644 index 0000000..4632640 --- /dev/null +++ b/src/grass_mil/configs/scheduler/cosine.yaml @@ -0,0 +1,4 @@ +# Cosine annealing over every optimizer step of the run (after any warmup). +_target_: torch.optim.lr_scheduler.CosineAnnealingLR +T_max: auto +eta_min: 1.0e-7 diff --git a/src/grass_mil/configs/scheduler/cosine_epoch.yaml b/src/grass_mil/configs/scheduler/cosine_epoch.yaml deleted file mode 100644 index 96f1858..0000000 --- a/src/grass_mil/configs/scheduler/cosine_epoch.yaml +++ /dev/null @@ -1,3 +0,0 @@ -_target_: torch.optim.lr_scheduler.CosineAnnealingLR -T_max: 100 -eta_min: 1.0e-6 diff --git a/src/grass_mil/configs/scheduler/cosine_mil_step.yaml b/src/grass_mil/configs/scheduler/cosine_mil_step.yaml deleted file mode 100644 index cc5789d..0000000 --- a/src/grass_mil/configs/scheduler/cosine_mil_step.yaml +++ /dev/null @@ -1,6 +0,0 @@ -# Cosine annealing per optimizer step over the whole run (preprint): with -# task.scheduler_interval=step, `T_max: auto` resolves to -# max_epochs * ceil(n_train_regions / hyperbatch_size). -_target_: torch.optim.lr_scheduler.CosineAnnealingLR -T_max: auto -eta_min: 1.0e-7 diff --git a/src/grass_mil/configs/scheduler/cosine_step.yaml b/src/grass_mil/configs/scheduler/cosine_step.yaml deleted file mode 100644 index b1afa31..0000000 --- a/src/grass_mil/configs/scheduler/cosine_step.yaml +++ /dev/null @@ -1,4 +0,0 @@ -_target_: torch.optim.lr_scheduler.CosineAnnealingLR -# Step-wise scheduler intended for BGRL pretraining. -T_max: 10000 -eta_min: 1.0e-7 diff --git a/src/grass_mil/configs/scores.yaml b/src/grass_mil/configs/scores.yaml deleted file mode 100644 index ee18ab7..0000000 --- a/src/grass_mil/configs/scores.yaml +++ /dev/null @@ -1,29 +0,0 @@ -# grass-mil-scores: evaluation statistics for held-out scores. -# -# grass-mil-scores input=/oof_predictions.csv -# grass-mil-scores input=table.csv scores=[grass,clinical] strata=[subtype] -# -# Input: a CSV with one row per unit (region or patient) holding a label -# column, one or more score columns, optionally a fold column (for within-fold -# standardisation and per-fold AUROCs), a patient column (scores are averaged -# per patient first when `patient_level: true`) and strata columns. -defaults: - - _self_ - - paths: default - - hydra: default - -task_name: scores -input: ??? -output_dir: ${paths.output_dir} -label_column: label -scores: [prob_class_1] # score columns to evaluate -fold_column: outer # null = no folds -patient_column: patient_id # null = evaluate rows as they are -patient_level: true # average region scores per patient before everything else -standardise: z # z | rank | none (within fold, before pooling) -strata: [] # optional columns; AUROC reported within each level -n_boot: 2000 -seed: 0 -operating_rules: [0.2, 0.3, 0.4, 0.5, 0.6] # flag the top q within fold -prevalences: [0.1, 0.2, 0.3, 0.4, 0.5] # PPV/NPV at other prevalences (first score) -report_filename: report_scores.md diff --git a/src/grass_mil/configs/task/classification.yaml b/src/grass_mil/configs/task/classification.yaml new file mode 100644 index 0000000..69a7a7f --- /dev/null +++ b/src/grass_mil/configs/task/classification.yaml @@ -0,0 +1,13 @@ +# Classification of units (regions, tiles) into num_classes classes. +target_type: classification +# Label-file column(s) to predict; null = the first non-patient column. +target_columns: null +num_classes: 2 +# auto = inverse class frequency on the training units (mean 1); a list of +# per-class weights; or null. +class_weights: auto +# Epoch selection and early stopping (validation metrics are region-level). +selection: + metric: val/auroc + mode: max + patience: 8 diff --git a/src/grass_mil/configs/task/finetune_mean.yaml b/src/grass_mil/configs/task/finetune_mean.yaml deleted file mode 100644 index 6ea7c43..0000000 --- a/src/grass_mil/configs/task/finetune_mean.yaml +++ /dev/null @@ -1,16 +0,0 @@ -# @package _global_ - -task: - name: finetune_mean - aggregation: mean - bag_key: patch_id - bag_fallback_key: sample_id - max_instances_per_bag: 0 - instance_sampling: all - target_type: binary - target_columns: null - loss: categorical_bce - lr_warmup: - enabled: false - warmup_steps: 0 - start_factor: 0.1 diff --git a/src/grass_mil/configs/task/finetune_mil.yaml b/src/grass_mil/configs/task/finetune_mil.yaml deleted file mode 100644 index b351c60..0000000 --- a/src/grass_mil/configs/task/finetune_mil.yaml +++ /dev/null @@ -1,74 +0,0 @@ -# @package _global_ - -task: - name: finetune_mil - aggregation: mil_attention - # One bag per region: with `data.sampler.runtime.unit_as_bag: true` every - # batch holds all sampled ego-graphs of one region, so the attention softmax - # and the loss cover the whole region. - bag_key: region_id - bag_fallback_key: sample_id - max_instances_per_bag: 0 - instance_sampling: all - # The MIL head is per-class: the bag logit for class c is - # sum_i A[i,c] * l[i,c], which requires one logit and one attention channel - # per class. A single-logit BCE head cannot express the attribution identity, - # so this preset is categorical with C=2 for binary tasks. - target_type: categorical - target_columns: null - loss: categorical_ce - - # Class weights for the region cross-entropy: `auto_train` = inverse frequency - # on the TRAINING regions (mean 1, preprint); a list of floats; or null. - class_weights: auto_train - - # Per-component learning-rate multipliers on optim.lr. A scale of 0 freezes the - # component (no gradients, encoder kept in eval mode). `head_weight_decay` - # overrides optim.weight_decay for the head group (null = same as optimizer). - optimization: - encoder_lr_scale: 1.0 - head_lr_scale: 1.0 - attention_lr_scale: 2.0 - head_weight_decay: null - # Legacy absolute pair (set both or neither); ignored when scales are used. - backbone_lr: null - attention_lr: null - - # Training-time augmentation of each ego-graph batch (preprint regularised - # regime: 0.05 / 0.05). Edge drop keeps both directions together; feature drop - # is element-wise and also writes the unassigned cell-type token. - finetune_augmentation: - drop_edge_p: 0.0 - drop_feat_p: 0.0 - seed: 0 - - # Linear-head pre-fit (grass-mil-probe output) and head freezing, used by the - # linear_* regimes. Requires model.graph_head.num_layers=1. - head_init_path: null - freeze_head: false - - # Gradient-norm clipping applied inside the region-accumulation flush. - grad_clip_norm: 1.0 - # Scheduler stepping: `step` (per optimizer step, with scheduler=cosine_mil_step - # and T_max: auto this anneals over the whole run) or `epoch`. - scheduler_interval: step - - lr_warmup: - enabled: false - warmup_steps: 0 - start_factor: 0.1 - # Hyperbatch has one meaning only: - # number of regions accumulated before one optimizer/scheduler step. - region_accumulation: - enabled: true - mode: manual_region_buffer - hyperbatch_size: 12 - flush_on_epoch_end: true - -# Per-class attention and matching head width. Override both together when -# moving to a task with more than two classes. -model: - graph_head: - output_dim: 2 - attention: - n_classes: 2 diff --git a/src/grass_mil/configs/task/finetune_survival.yaml b/src/grass_mil/configs/task/finetune_survival.yaml deleted file mode 100644 index cc9f46b..0000000 --- a/src/grass_mil/configs/task/finetune_survival.yaml +++ /dev/null @@ -1,43 +0,0 @@ -# @package _global_ - -# Cox survival with MIL attention pooling. The head emits a single scalar -# log-hazard per bag, so attention stays a single shared channel: there is no -# per-class decomposition for a one-dimensional risk score. - -task: - name: finetune_survival - aggregation: mil_attention - bag_key: patch_id - bag_fallback_key: sample_id - max_instances_per_bag: 0 - instance_sampling: all - target_type: survival - # Column ORDER is load-bearing: the loss reads [:, 0] as time and [:, 1] as - # the event indicator. Rename to match your targets file, but keep the order. - target_columns: [time, event] - loss: survival_coxsgd - optimization: - backbone_lr: null - attention_lr: null - lr_warmup: - enabled: false - warmup_steps: 0 - start_factor: 0.1 - region_accumulation: - enabled: false - mode: manual_region_buffer - hyperbatch_size: 8 - flush_on_epoch_end: true - -model: - graph_head: - output_dim: 1 - attention: - n_classes: 1 - loss: - loss_type: survival_coxsgd - # Risk-set truncation; 0 uses the full risk set. - top_n: 5 - # Shrinks the spread of predicted risks. Leaving this at 0 lets attention - # over-sharpen on small validation cohorts. - regularizer_weight: 0.05 diff --git a/src/grass_mil/configs/task/pretrain_bgrl.yaml b/src/grass_mil/configs/task/pretrain_bgrl.yaml deleted file mode 100644 index fbbc2c3..0000000 --- a/src/grass_mil/configs/task/pretrain_bgrl.yaml +++ /dev/null @@ -1,30 +0,0 @@ -# @package _global_ - -defaults: - - override /scheduler: cosine_step - -task: - name: pretrain_bgrl - aggregation: mean - bag_key: patch_id - bag_fallback_key: sample_id - max_instances_per_bag: 0 - instance_sampling: all - target_type: binary - target_columns: null - loss: categorical_bce - lr_warmup: - enabled: true - warmup_steps: 1000 - start_factor: 0.1 - warmup_steps: 1000 - total_steps: 10000 - momentum: 0.99 - momentum_min: 0.99 - # Seeds the augmentation draws so pretraining views are reproducible - # independently of the global RNG. - augmentation_seed: 0 - drop_edge_p1: 0.2 - drop_edge_p2: 0.2 - drop_feat_p1: 0.2 - drop_feat_p2: 0.2 diff --git a/src/grass_mil/configs/task/pretrain_bgrl_a3.yaml b/src/grass_mil/configs/task/pretrain_bgrl_a3.yaml deleted file mode 100644 index 103c661..0000000 --- a/src/grass_mil/configs/task/pretrain_bgrl_a3.yaml +++ /dev/null @@ -1,53 +0,0 @@ -# @package _global_ -# BGRL pretraining, preprint "A3" recipe: AdamW 5e-5 / weight decay 1e-5, -# linear LR warmup over 1000 steps then cosine to 1e-7 at 200k steps, target -# EMA momentum 0.99 -> 1.0 (cosine, no warmup), symmetric augmentations -# DropImportance(mu 0.25, p_lambda 0.45, degree importance, unassigned token) -# + Gaussian noise (std 0.2) on the cell-size feature, checkpoint every 5000 -# steps, all regions used label-free. -defaults: - - override /scheduler: cosine_step - - override /optim: adamw - -optim: - lr: 5.0e-5 - weight_decay: 1.0e-5 - -scheduler: - T_max: auto # = total_steps - lr warmup - eta_min: 1.0e-7 - -task: - name: pretrain_bgrl - aggregation: mean - bag_key: patch_id - bag_fallback_key: sample_id - max_instances_per_bag: 0 - instance_sampling: all - target_type: binary - target_columns: null - loss: categorical_bce - lr_warmup: - enabled: true - warmup_steps: 1000 - start_factor: 0.0 # linear from ~0 - warmup_steps: 1000 - total_steps: 200000 - momentum: 0.99 - momentum_final: 1.0 - momentum_warmup_steps: 0 - augmentation_seed: 0 - augmentation: - mode: importance - mu: 0.25 - p_lambda: 0.45 - feature_noise_std: 0.2 - # Continuous node feature(s) to noise, by name (molecular_columns) or index. - # Set to [] for cohorts without a cell-size feature. - feature_noise_columns: [SIZE] - -trainer: - max_steps: ${task.total_steps} - max_epochs: -1 - precision: 16-mixed - val_check_interval: null diff --git a/src/grass_mil/configs/task/regression.yaml b/src/grass_mil/configs/task/regression.yaml new file mode 100644 index 0000000..c16ffc9 --- /dev/null +++ b/src/grass_mil/configs/task/regression.yaml @@ -0,0 +1,9 @@ +# Regression of one or more continuous unit-level targets. +target_type: regression +target_columns: ??? +loss: mse # mse | huber +huber_delta: 1.0 +selection: + metric: val/rmse + mode: min + patience: 8 diff --git a/src/grass_mil/configs/task/survival.yaml b/src/grass_mil/configs/task/survival.yaml new file mode 100644 index 0000000..028c423 --- /dev/null +++ b/src/grass_mil/configs/task/survival.yaml @@ -0,0 +1,8 @@ +# Survival: one risk score per unit, Cox partial likelihood, concordance index. +target_type: survival +target_columns: ??? # [time, event] +cox_regularization: 0.0 # L2 penalty on the risk scores +selection: + metric: val/c_index + mode: max + patience: 8 diff --git a/src/grass_mil/configs/train.yaml b/src/grass_mil/configs/train.yaml index f1e320d..ae6aa78 100644 --- a/src/grass_mil/configs/train.yaml +++ b/src/grass_mil/configs/train.yaml @@ -1,59 +1,30 @@ -# @package _global_ - -# specify here default configuration -# order of defaults determines the order in which configs override each other +# grass-mil-train: one training run (fine-tuning by default, `regime=pretrain` +# for label-free encoder pretraining). Every key below can be overridden on the +# command line; see docs/configuration.md. defaults: - _self_ - data: spatial_omics - - model: supervised_module + - model: grass_mil + - task: classification - optim: adamw - - scheduler: cosine_epoch - # `task` is merged after `model`, `optim` and `scheduler` so that a task - # preset can override them. finetune_mil sets the per-class head and - # attention widths, and pretrain_bgrl overrides the scheduler; both are - # silently discarded if `task` is listed first. - - task: finetune_mean - - callbacks: default - - logger: null # set logger here or use command line (e.g. `python train.py logger=tensorboard`) + - scheduler: cosine + - callbacks: selection + - logger: null - trainer: default - paths: default - - extras: default - hydra: default - - # experiment configs allow for version control of specific hyperparameters - # e.g. best hyperparameters for given model and datamodule + # A cohort description (configs/experiment/example_cohort.yaml as template). - experiment: null - # Training regime presets (configs/regime/*.yaml): mil, mil_frozen, - # linear_probehead, linear_ftall, linear_ftall_reg. Select with `regime=mil`. - - regime: null - - # config for hyperparameter optimization - - hparams_search: null - - # optional local config for machine/user specific settings - # it's optional since it doesn't need to exist and is excluded from version control - - optional local: default - - # debugging config (enable through command line, e.g. `python train.py debug=default) - - debug: null - -# task name, determines output directory path -task_name: "train" - -# tags to help you identify your experiments -# you can overwrite this in experiment configs -# overwrite from command line with `python train.py tags="[first_tag, second_tag]"` -tags: ["dev"] - -# set False to skip model training -train: True - -# evaluate on test set, using best model weights achieved during training -# lightning chooses best weights based on the metric specified in checkpoint callback -test: True - -# simply provide checkpoint path to resume training + # Training procedure: finetune (default), finetune_reg, attention_only, + # mlp_head, mlp_head_frozen, or pretrain. + - regime: finetune + +task_name: train +seed: 0 +# Fit the model, then score the test split with the selected checkpoint. +train: true +test: true +# Resume an interrupted run from this checkpoint. ckpt_path: null - -# seed for random number generators in pytorch, numpy and python.random -seed: null +# Print the composed config at the start of the run. +print_config: false diff --git a/src/grass_mil/configs/trainer/cpu.yaml b/src/grass_mil/configs/trainer/cpu.yaml index b7d6767..5618c41 100644 --- a/src/grass_mil/configs/trainer/cpu.yaml +++ b/src/grass_mil/configs/trainer/cpu.yaml @@ -1,5 +1,2 @@ defaults: - default - -accelerator: cpu -devices: 1 diff --git a/src/grass_mil/configs/trainer/ddp.yaml b/src/grass_mil/configs/trainer/ddp.yaml deleted file mode 100644 index ab8f890..0000000 --- a/src/grass_mil/configs/trainer/ddp.yaml +++ /dev/null @@ -1,9 +0,0 @@ -defaults: - - default - -strategy: ddp - -accelerator: gpu -devices: 4 -num_nodes: 1 -sync_batchnorm: True diff --git a/src/grass_mil/configs/trainer/ddp_sim.yaml b/src/grass_mil/configs/trainer/ddp_sim.yaml deleted file mode 100644 index 8404419..0000000 --- a/src/grass_mil/configs/trainer/ddp_sim.yaml +++ /dev/null @@ -1,7 +0,0 @@ -defaults: - - default - -# simulate DDP on CPU, useful for debugging -accelerator: cpu -devices: 2 -strategy: ddp_spawn diff --git a/src/grass_mil/configs/trainer/default.yaml b/src/grass_mil/configs/trainer/default.yaml index 50905e7..eb4a4b7 100644 --- a/src/grass_mil/configs/trainer/default.yaml +++ b/src/grass_mil/configs/trainer/default.yaml @@ -1,19 +1,9 @@ -_target_: lightning.pytorch.trainer.Trainer - +_target_: lightning.pytorch.Trainer default_root_dir: ${paths.output_dir} - -min_epochs: 1 # prevents early stopping -max_epochs: 10 - accelerator: cpu devices: 1 - -# mixed precision for extra speed-up -# precision: 16 - -# perform a validation loop every N training epochs +min_epochs: 1 +max_epochs: 10 +precision: 32 check_val_every_n_epoch: 1 - -# set True to to ensure deterministic results -# makes training slower but gives more reproducibility than just setting seeds -deterministic: False +deterministic: false diff --git a/src/grass_mil/configs/trainer/gpu.yaml b/src/grass_mil/configs/trainer/gpu.yaml index b238951..b685e5f 100644 --- a/src/grass_mil/configs/trainer/gpu.yaml +++ b/src/grass_mil/configs/trainer/gpu.yaml @@ -1,5 +1,7 @@ +# One GPU with fp16 mixed precision. defaults: - default accelerator: gpu devices: 1 +precision: 16-mixed diff --git a/src/grass_mil/configs/trainer/gpu_amp.yaml b/src/grass_mil/configs/trainer/gpu_amp.yaml deleted file mode 100644 index 6aef739..0000000 --- a/src/grass_mil/configs/trainer/gpu_amp.yaml +++ /dev/null @@ -1,6 +0,0 @@ -defaults: - - default - -accelerator: gpu -devices: 1 -precision: 16-mixed diff --git a/src/grass_mil/configs/trainer/mps.yaml b/src/grass_mil/configs/trainer/mps.yaml deleted file mode 100644 index 1ecf6d5..0000000 --- a/src/grass_mil/configs/trainer/mps.yaml +++ /dev/null @@ -1,5 +0,0 @@ -defaults: - - default - -accelerator: mps -devices: 1 diff --git a/src/grass_mil/contracts.py b/src/grass_mil/contracts.py deleted file mode 100644 index 030bf1c..0000000 --- a/src/grass_mil/contracts.py +++ /dev/null @@ -1,138 +0,0 @@ -"""Cross-cutting contracts. - -Invariants and payload shapes shared between the model, inference and -interpretability layers. Kept in one dependency-light module so any layer can -import it without creating a cycle -- ``grass_mil.models`` and -``grass_mil.inference`` already import each other through the supervised -module. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Dict, List, Optional - -import numpy as np -import pandas as pd - -__all__ = ["validate_attention_width"] - - -def validate_attention_width(*, attention_width: int, num_classes: int) -> None: - """Require one attention channel per class. - - Bag pooling computes ``L_c = sum_i A[i,c] * l[i,c]``, and the attribution - identity that the interpretability suite rests on is only exact when each - class has its own attention distribution. A binary head therefore needs - **two** attention channels, not one shared channel broadcast across both: - a shared channel cannot express a niche that pushes toward one class and - away from the other. - """ - if attention_width != num_classes: - raise ValueError( - f"Attention emits {attention_width} channel(s) but the head emits " - f"{num_classes} class logit(s). Attention is per-class: set " - f"model.attention.n_classes to {num_classes}." - ) - - -@dataclass(frozen=True) -class EmbeddingSet: - matrix: np.ndarray - feature_columns: List[str] - id_column: str = "instance_id" - - -@dataclass(frozen=True) -class InterpretabilityDataset: - instance_table: pd.DataFrame - bag_table: Optional[pd.DataFrame] = None - spatial_table: Optional[pd.DataFrame] = None - # Cell-level tables: one row per node of every sampled ego-graph, and one - # row per intra-subgraph edge. Required by the analyses that measure cells - # rather than neighbourhood summaries. - cell_table: Optional[pd.DataFrame] = None - cell_edge_table: Optional[pd.DataFrame] = None - cell_id_column: str = "cell_id" - id_column: str = "instance_id" - bag_id_column: str = "bag_id" - cell_type_column: Optional[str] = None - condition_column: Optional[str] = None - - -@dataclass(frozen=True) -class ReductionResult: - method: str - embedding: np.ndarray - params: Dict[str, Any] - fitted_object: Any = None - - -@dataclass(frozen=True) -class ClusteringResult: - method: str - labels: np.ndarray - params: Dict[str, Any] - fitted_object: Any = None - - -@dataclass(frozen=True) -class NicheSummary: - niche_labels: np.ndarray - composition: pd.DataFrame - enrichment: pd.DataFrame - niche_counts: pd.Series - - -@dataclass(frozen=True) -class NeighborhoodEnrichmentResult: - enrichment: pd.DataFrame - observed: pd.DataFrame - expected: pd.DataFrame - pvalues: Optional[pd.DataFrame] = None - - -@dataclass(frozen=True) -class FiltrationCurvesResult: - thresholds: np.ndarray - curves: Dict[str, Dict[str, np.ndarray]] - - -@dataclass(frozen=True) -class ReportSection: - title: str - description: str - figure_html: Optional[str] = None - tables: Dict[str, pd.DataFrame] = field(default_factory=dict) - metadata: Dict[str, Any] = field(default_factory=dict) - - -@dataclass(frozen=True) -class PluginResult: - name: str - payload: Dict[str, Any] - sections: List[ReportSection] - - -@dataclass(frozen=True) -class NicheTransferBundle: - embedding_columns: List[str] - id_column: str - cluster_on_pca: bool - pca_model: Any = None - knn_model: Any = None - label_values: np.ndarray = field(default_factory=lambda: np.asarray([])) - metadata: Dict[str, Any] = field(default_factory=dict) - - -@dataclass(frozen=True) -class ReportBundle: - dataset: InterpretabilityDataset - reduction: Optional[ReductionResult] - niche_feature_reduction: Optional[ReductionResult] - clustering: Optional[ClusteringResult] - niche_summary: Optional[NicheSummary] - plugin_results: Dict[str, PluginResult] - artifacts_dir: Path - metadata: Dict[str, Any] = field(default_factory=dict) diff --git a/src/grass_mil/data/__init__.py b/src/grass_mil/data/__init__.py index 332dfa3..2e3db99 100644 --- a/src/grass_mil/data/__init__.py +++ b/src/grass_mil/data/__init__.py @@ -1,3 +1,14 @@ -from .spatial_omics_datamodule import SpatialOmicsDataModule +"""Data layer: cell tables to graph units, labels, splits and ego-graph loaders.""" -__all__ = ["SpatialOmicsDataModule"] +from .datamodule import GraphUnitDataset, InputSpec, SpatialOmicsDataModule +from .labels import LabelTable +from .samplers import EgoGraphLoader, SamplerConfig + +__all__ = [ + "EgoGraphLoader", + "GraphUnitDataset", + "InputSpec", + "LabelTable", + "SamplerConfig", + "SpatialOmicsDataModule", +] diff --git a/src/grass_mil/data/components/__init__.py b/src/grass_mil/data/components/__init__.py deleted file mode 100644 index d1d0558..0000000 --- a/src/grass_mil/data/components/__init__.py +++ /dev/null @@ -1,43 +0,0 @@ -from .datasets import SpatialOmicsGraphDataset, TransformDataset -from .feature_reducers import FeatureReducerConfig, get_feature_reducer -from .graph_builders import GraphBuilderConfig, get_graph_builder -from .loaders import ( - CsvConfig, - H5adConfig, - SceConfig, - load_csv_table, - load_h5ad_table, - load_polygons_from_path, - load_sce_table, -) -from .patching import TileConfig, build_grid_tiles, build_patches_from_polygons -from .precompute import SpatialOmicsPreprocessor -from .samplers import SamplerConfig, get_sampler_strategy, register_sampler_strategy -from .spatial_types import SpatialOmicsTable -from .transforms import CompositionVector, instantiate_transforms - -__all__ = [ - "CompositionVector", - "CsvConfig", - "FeatureReducerConfig", - "GraphBuilderConfig", - "H5adConfig", - "SamplerConfig", - "SceConfig", - "SpatialOmicsGraphDataset", - "SpatialOmicsPreprocessor", - "SpatialOmicsTable", - "TileConfig", - "TransformDataset", - "build_grid_tiles", - "build_patches_from_polygons", - "get_feature_reducer", - "get_graph_builder", - "get_sampler_strategy", - "instantiate_transforms", - "load_csv_table", - "load_h5ad_table", - "load_polygons_from_path", - "load_sce_table", - "register_sampler_strategy", -] diff --git a/src/grass_mil/data/components/datasets.py b/src/grass_mil/data/components/datasets.py deleted file mode 100644 index a755eeb..0000000 --- a/src/grass_mil/data/components/datasets.py +++ /dev/null @@ -1,86 +0,0 @@ -from __future__ import annotations - -import json -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Dict, List, Optional - -import torch -from torch.utils.data import Dataset - - -@dataclass -class ProcessedIndexEntry: - path: str - sample_id: str - region_id: Optional[str] - patch_id: str - split: Optional[str] = None - patient_id: Optional[str] = None - graph_y: Optional[List[float]] = None - - -def _migrate_legacy_attributes(data): - """Upgrade graphs written by an older precompute. - - ``categorical_index`` was renamed to ``categorical_codes`` because - PyTorch Geometric treats *any* attribute whose name contains ``index`` as an - edge-index tensor and concatenates it along the last dimension, which - corrupts node-level codes as soon as more than one graph is batched. - Migrating on load means existing processed caches stay usable without a - forced re-precompute. - """ - legacy = getattr(data, "categorical_index", None) - if legacy is not None and getattr(data, "categorical_codes", None) is None: - data.categorical_codes = legacy - del data.categorical_index - return data - - -class SpatialOmicsGraphDataset(Dataset): - def __init__(self, index_path: str | Path) -> None: - index_path = Path(index_path) - if not index_path.exists(): - raise FileNotFoundError(f"Processed index not found: {index_path}") - self.index_path = index_path - data = json.loads(index_path.read_text()) - self.entries: List[ProcessedIndexEntry] = [ - ProcessedIndexEntry(**entry) for entry in data["entries"] - ] - self.label_maps: Dict[str, Dict[str, int]] = data.get("label_maps", {}) - self.label_maps_inv: Dict[str, Dict[int, str]] = { - label: {idx: name for name, idx in mapping.items()} - for label, mapping in self.label_maps.items() - } - self.graph_label_maps: Dict[str, Dict[str, int]] = data.get( - "graph_label_maps", {} - ) - self.reducer_state: Dict[str, Any] = data.get("reducer_state", {}) - - def __len__(self) -> int: - return len(self.entries) - - def __getitem__(self, idx: int) -> torch.nn.Module: - entry = self.entries[idx] - # weights_only=False is required: these are PyG `Data` objects written - # by our own precompute step, not bare tensors. torch>=2.6 defaults the - # flag to True, which cannot unpickle `DataEdgeAttr`/`GlobalStorage` - # without allowlisting a set of PyG internals that changes between - # releases. The file is first-party, produced by this pipeline. - data = torch.load(entry.path, weights_only=False) - return _migrate_legacy_attributes(data) - - -class TransformDataset(Dataset): - def __init__(self, dataset: Dataset, transforms) -> None: - self.dataset = dataset - self.transforms = transforms - - def __len__(self) -> int: - return len(self.dataset) - - def __getitem__(self, idx: int): - data = self.dataset[idx] - for transform in self.transforms: - data = transform(data) - return data diff --git a/src/grass_mil/data/components/ego_radius.py b/src/grass_mil/data/components/ego_radius.py deleted file mode 100644 index 6cffb75..0000000 --- a/src/grass_mil/data/components/ego_radius.py +++ /dev/null @@ -1,180 +0,0 @@ -"""Euclidean radius cutoff for k-hop ego-graphs. - -A k-hop ego-graph can reach far across sparse tissue, so hop depth alone does -not bound its physical extent. The published sampler additionally caps the -ego-graph at a radius that grows linearly with depth: - -.. math:: r_k = \\mathrm{radius\\_per\\_hop} \\cdot k + \\mathrm{radius\\_offset} - -Distances are in the datamodule's coordinate units (micrometres once -``coord_scale_um`` has been applied), *not* pixels: the published 75/55 are -pixel values at the NSCLC imaging scale and must be converted per cohort. - -Applied as a post-hoc transform on sampled subgraphs, so no sampler internals -are involved. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Mapping, Optional - -import numpy as np -import torch - -__all__ = ["EgoRadiusConfig", "EgoRadiusCutoff", "resolve_ego_radius", "apply_ego_radius_cutoff"] - - -@dataclass(frozen=True) -class EgoRadiusConfig: - enabled: bool = False - radius_per_hop: float = 75.0 - radius_offset: float = 55.0 - keep_root_component: bool = True - - @classmethod - def from_dict(cls, cfg: Optional[Mapping[str, Any]]) -> "EgoRadiusConfig": - cfg = dict(cfg or {}) - allowed = {"enabled", "radius_per_hop", "radius_offset", "keep_root_component"} - unknown = set(cfg) - allowed - if unknown: - raise ValueError( - f"Unknown ego radius keys: {sorted(unknown)}. " - f"Expected a subset of {sorted(allowed)}." - ) - return cls( - enabled=bool(cfg.get("enabled", False)), - radius_per_hop=float(cfg.get("radius_per_hop", 75.0)), - radius_offset=float(cfg.get("radius_offset", 55.0)), - keep_root_component=bool(cfg.get("keep_root_component", True)), - ) - - -def resolve_ego_radius(depth: int, *, radius_per_hop: float, radius_offset: float) -> float: - """``r_k = radius_per_hop * k + radius_offset``.""" - if depth < 0: - raise ValueError(f"depth must be >= 0, got {depth}.") - return float(radius_per_hop) * float(depth) + float(radius_offset) - - -def _root_index(data) -> Optional[int]: - root = getattr(data, "root_n_id", None) - if root is None: - return None - if isinstance(root, torch.Tensor): - if root.numel() == 0: - return None - return int(root.reshape(-1)[0]) - return int(root) - - -def apply_ego_radius_cutoff(data, *, radius: float, keep_root_component: bool = True): - """Drop nodes farther than ``radius`` from the ego-graph root. - - The root is always retained. When ``keep_root_component`` is set, nodes left - disconnected from the root by the cut are dropped too, so the ego-graph - stays a single connected neighbourhood rather than the root plus islands. - """ - from torch_geometric.utils import subgraph - - pos = getattr(data, "pos", None) - if pos is None or radius <= 0: - return data - root = _root_index(data) - if root is None: - return data - - coords = pos[:, :2].float() - num_nodes = int(coords.size(0)) - if num_nodes == 0 or root >= num_nodes: - return data - - distances = torch.linalg.norm(coords - coords[root].unsqueeze(0), dim=1) - keep = distances <= float(radius) - keep[root] = True - - if keep_root_component and bool(keep.any()): - keep = _root_connected_mask(data, keep, root=root, num_nodes=num_nodes) - - if bool(keep.all()): - return data - - keep_idx = torch.where(keep)[0] - edge_attr = getattr(data, "edge_attr", None) - new_edge_index, new_edge_attr = subgraph( - subset=keep_idx, - edge_index=data.edge_index, - edge_attr=edge_attr, - relabel_nodes=True, - num_nodes=num_nodes, - ) - - out = data.clone() - out.edge_index = new_edge_index - if edge_attr is not None: - out.edge_attr = new_edge_attr - for key, value in data: - if key in {"edge_index", "edge_attr"}: - continue - if isinstance(value, torch.Tensor) and value.dim() > 0 and value.size(0) == num_nodes: - setattr(out, key, value[keep_idx]) - out.num_nodes = int(keep_idx.numel()) - - # Root position shifts when earlier nodes are dropped. - remap = torch.full((num_nodes,), -1, dtype=torch.long) - remap[keep_idx] = torch.arange(keep_idx.numel(), dtype=torch.long) - out.root_n_id = remap[root].reshape(1) - return out - - -def _root_connected_mask(data, keep: torch.Tensor, *, root: int, num_nodes: int) -> torch.Tensor: - """Restrict ``keep`` to the connected component containing the root.""" - from scipy.sparse import coo_matrix - from scipy.sparse.csgraph import connected_components - - keep_np = keep.detach().cpu().numpy() - edge_index = data.edge_index.detach().cpu().numpy() - if edge_index.size == 0: - mask = np.zeros(num_nodes, dtype=bool) - mask[root] = True - return torch.from_numpy(mask) - - src, dst = edge_index[0], edge_index[1] - both_kept = keep_np[src] & keep_np[dst] - src, dst = src[both_kept], dst[both_kept] - rows = np.concatenate([src, dst]) - cols = np.concatenate([dst, src]) - adjacency = coo_matrix( - (np.ones(rows.shape[0], dtype=bool), (rows, cols)), - shape=(num_nodes, num_nodes), - ).tocsr() - - _, labels = connected_components(adjacency, directed=False) - mask = keep_np & (labels == labels[root]) - mask[root] = True - return torch.from_numpy(mask) - - -class EgoRadiusCutoff: - """Transform form of :func:`apply_ego_radius_cutoff`. - - Constructed with an explicit radius so the transform itself carries no - dependency on sampler internals. - """ - - def __init__(self, *, radius: float, keep_root_component: bool = True) -> None: - self.radius = float(radius) - self.keep_root_component = bool(keep_root_component) - - def __call__(self, data): - return apply_ego_radius_cutoff( - data, - radius=self.radius, - keep_root_component=self.keep_root_component, - ) - - def __repr__(self) -> str: - return ( - f"{type(self).__name__}(radius={self.radius}, " - f"keep_root_component={self.keep_root_component})" - ) diff --git a/src/grass_mil/data/components/feature_reducers.py b/src/grass_mil/data/components/feature_reducers.py deleted file mode 100644 index fce45ff..0000000 --- a/src/grass_mil/data/components/feature_reducers.py +++ /dev/null @@ -1,101 +0,0 @@ -from __future__ import annotations - -import importlib -from dataclasses import dataclass -from typing import Callable, Dict, Optional - -import numpy as np - - -class FeatureReducer: - def fit(self, features: np.ndarray) -> "FeatureReducer": - return self - - def transform(self, features: np.ndarray) -> np.ndarray: - return features - - def fit_transform(self, features: np.ndarray) -> np.ndarray: - self.fit(features) - return self.transform(features) - - def state_dict(self) -> Dict[str, object]: - return {} - - def load_state_dict(self, state: Dict[str, object]) -> None: - _ = state - - -@dataclass -class FeatureReducerConfig: - name: str - kwargs: Dict[str, object] - - @classmethod - def from_dict(cls, data: Dict[str, object]) -> "FeatureReducerConfig": - return cls( - name=data["name"], - kwargs=data.get("kwargs", {}), - ) - - -_FEATURE_REDUCERS: Dict[str, Callable[..., FeatureReducer]] = {} - - -def register_feature_reducer(name: str): - def decorator(cls: Callable[..., FeatureReducer]): - _FEATURE_REDUCERS[name] = cls - return cls - - return decorator - - -def get_feature_reducer(config: FeatureReducerConfig) -> FeatureReducer: - if config.name in _FEATURE_REDUCERS: - return _FEATURE_REDUCERS[config.name](**(config.kwargs or {})) - if "." in config.name: - module_name, class_name = config.name.rsplit(".", 1) - module = importlib.import_module(module_name) - cls = getattr(module, class_name) - return cls(**(config.kwargs or {})) - raise ValueError( - f"Unknown feature reducer '{config.name}'. Available: {list(_FEATURE_REDUCERS)}" - ) - - -@register_feature_reducer("identity") -class IdentityReducer(FeatureReducer): - pass - - -@register_feature_reducer("pca") -class PcaReducer(FeatureReducer): - def __init__(self, n_components: int = 64) -> None: - self.n_components = n_components - self.components_: Optional[np.ndarray] = None - self.mean_: Optional[np.ndarray] = None - - def fit(self, features: np.ndarray) -> "PcaReducer": - mean = features.mean(axis=0, keepdims=True) - centered = features - mean - _, _, vt = np.linalg.svd(centered, full_matrices=False) - self.components_ = vt[: self.n_components] - self.mean_ = mean - return self - - def transform(self, features: np.ndarray) -> np.ndarray: - if self.components_ is None or self.mean_ is None: - raise RuntimeError("PCA reducer has not been fitted.") - centered = features - self.mean_ - return centered @ self.components_.T - - def state_dict(self) -> Dict[str, object]: - return { - "n_components": self.n_components, - "components": self.components_, - "mean": self.mean_, - } - - def load_state_dict(self, state: Dict[str, object]) -> None: - self.n_components = int(state["n_components"]) - self.components_ = np.asarray(state["components"]) - self.mean_ = np.asarray(state["mean"]) diff --git a/src/grass_mil/data/components/graph_builders.py b/src/grass_mil/data/components/graph_builders.py deleted file mode 100644 index ea71cf6..0000000 --- a/src/grass_mil/data/components/graph_builders.py +++ /dev/null @@ -1,189 +0,0 @@ -from __future__ import annotations - -import importlib -from dataclasses import dataclass -from typing import Callable, Dict, List, Optional - -import numpy as np -import torch - - -@dataclass -class GraphBuildResult: - edge_index: torch.Tensor - edge_attr: Optional[torch.Tensor] = None - edge_attr_names: Optional[List[str]] = None - - -class GraphBuilder: - def build(self, coords: np.ndarray) -> GraphBuildResult: - raise NotImplementedError - - -@dataclass -class GraphBuilderConfig: - name: str - kwargs: Dict[str, object] - - @classmethod - def from_dict(cls, data: Dict[str, object]) -> "GraphBuilderConfig": - return cls( - name=data["name"], - kwargs=data.get("kwargs", {}), - ) - - -_GRAPH_BUILDERS: Dict[str, Callable[..., GraphBuilder]] = {} - - -def register_graph_builder(name: str): - def decorator(cls: Callable[..., GraphBuilder]): - _GRAPH_BUILDERS[name] = cls - return cls - - return decorator - - -def get_graph_builder(config: GraphBuilderConfig) -> GraphBuilder: - if config.name in _GRAPH_BUILDERS: - return _GRAPH_BUILDERS[config.name](**(config.kwargs or {})) - if "." in config.name: - module_name, class_name = config.name.rsplit(".", 1) - module = importlib.import_module(module_name) - cls = getattr(module, class_name) - return cls(**(config.kwargs or {})) - raise ValueError( - f"Unknown graph builder '{config.name}'. Available: {list(_GRAPH_BUILDERS)}" - ) - - -@register_graph_builder("delaunay") -class DelaunayGraphBuilder(GraphBuilder): - """Delaunay triangulation over cell centroids. - - Edge features (``edge_features``): - - - ``distance``: raw Euclidean distance in coordinate units. - - ``log_distance``: ``clip((ln(d + eps) - log_lower) / (log_upper - log_lower), 0, 1)``, - the bounded log-distance used as the GINE edge input in the GRASS-MIL preprint - (defaults ``log_lower=2.0``, ``log_upper=5.0``; with coordinates in pixels this maps - ~7 px -> 0 and ~148 px -> 1). - - ``neighbor``: 1 if ``distance <= neighbor_cutoff_um`` else 0. - """ - - def __init__( - self, - edge_features: Optional[List[str]] = None, - neighbor_cutoff_um: Optional[float] = None, - log_lower: float = 2.0, - log_upper: float = 5.0, - log_eps: float = 1e-5, - ) -> None: - self.edge_features = edge_features or [] - self.neighbor_cutoff_um = neighbor_cutoff_um - if log_upper <= log_lower: - raise ValueError("log_upper must be greater than log_lower") - self.log_lower = float(log_lower) - self.log_upper = float(log_upper) - self.log_eps = float(log_eps) - - def build(self, coords: np.ndarray) -> GraphBuildResult: - if coords.shape[0] < 2: - return GraphBuildResult(edge_index=torch.empty((2, 0), dtype=torch.long)) - if coords.shape[0] < 3: - edges = {(0, 1), (1, 0)} if coords.shape[0] == 2 else set() - edge_index = _edges_to_tensor(edges) - edge_attr, edge_attr_names = _build_edge_attr( - edge_index, - coords, - self.edge_features, - self.neighbor_cutoff_um, - log_lower=self.log_lower, - log_upper=self.log_upper, - log_eps=self.log_eps, - ) - return GraphBuildResult( - edge_index=edge_index, - edge_attr=edge_attr, - edge_attr_names=edge_attr_names, - ) - - from scipy.spatial import Delaunay - - tri = Delaunay(coords) - edges = set() - for simplex in tri.simplices: - for i in range(len(simplex)): - for j in range(i + 1, len(simplex)): - a = int(simplex[i]) - b = int(simplex[j]) - edges.add((a, b)) - edges.add((b, a)) - edge_index = _edges_to_tensor(edges) - edge_attr, edge_attr_names = _build_edge_attr( - edge_index, - coords, - self.edge_features, - self.neighbor_cutoff_um, - log_lower=self.log_lower, - log_upper=self.log_upper, - log_eps=self.log_eps, - ) - return GraphBuildResult( - edge_index=edge_index, edge_attr=edge_attr, edge_attr_names=edge_attr_names - ) - - -def _edges_to_tensor(edges: set[tuple[int, int]]) -> torch.Tensor: - if not edges: - return torch.empty((2, 0), dtype=torch.long) - edge_index = torch.tensor(list(edges), dtype=torch.long).t().contiguous() - return edge_index - - -def log_distance_feature( - distances: np.ndarray, log_lower: float = 2.0, log_upper: float = 5.0, eps: float = 1e-5 -) -> np.ndarray: - """Bounded log-distance edge feature: clip((ln(d + eps) - lower) / (upper - lower), 0, 1).""" - log_d = np.log(np.asarray(distances, dtype=np.float64) + eps) - return np.clip((log_d - log_lower) / (log_upper - log_lower), 0.0, 1.0).astype(np.float32) - - -def _build_edge_attr( - edge_index: torch.Tensor, - coords: np.ndarray, - edge_features: List[str], - neighbor_cutoff_um: Optional[float], - log_lower: float = 2.0, - log_upper: float = 5.0, - log_eps: float = 1e-5, -) -> tuple[Optional[torch.Tensor], Optional[List[str]]]: - if not edge_features: - return None, None - if edge_index.numel() == 0: - return torch.empty((0, len(edge_features)), dtype=torch.float), list( - edge_features - ) - - src = edge_index[0].cpu().numpy() - dst = edge_index[1].cpu().numpy() - deltas = coords[src] - coords[dst] - distances = np.linalg.norm(deltas, axis=1) - - columns: List[np.ndarray] = [] - for feat in edge_features: - if feat == "distance": - columns.append(distances.astype(np.float32)) - elif feat == "log_distance": - columns.append(log_distance_feature(distances, log_lower, log_upper, log_eps)) - elif feat == "neighbor": - if neighbor_cutoff_um is None: - raise ValueError( - "neighbor_cutoff_um must be set to compute neighbor edge feature" - ) - columns.append((distances <= neighbor_cutoff_um).astype(np.float32)) - else: - raise ValueError(f"Unsupported edge feature '{feat}'") - - edge_attr = np.stack(columns, axis=1) - return torch.from_numpy(edge_attr), list(edge_features) diff --git a/src/grass_mil/data/components/loaders.py b/src/grass_mil/data/components/loaders.py deleted file mode 100644 index 5a8015c..0000000 --- a/src/grass_mil/data/components/loaders.py +++ /dev/null @@ -1,382 +0,0 @@ -from __future__ import annotations - -import json -import pickle -from dataclasses import dataclass -from pathlib import Path -from typing import Dict, Iterable, List, Optional, Sequence, Tuple - -import numpy as np -import pandas as pd - -from .spatial_types import SpatialOmicsTable - - -@dataclass -class CsvConfig: - coord_columns: Tuple[str, str] - cell_id_column: Optional[str] - categorical_label_columns: Sequence[str] - molecular_columns: Optional[Sequence[str]] - use_molecular_features: bool = True - sep: str = "," - - @classmethod - def from_dict( - cls, data: Dict[str, object], use_molecular_features: bool - ) -> "CsvConfig": - return cls( - coord_columns=tuple(data["coord_columns"]), - cell_id_column=data.get("cell_id_column"), - categorical_label_columns=data.get("categorical_label_columns", []), - molecular_columns=data.get("molecular_columns"), - use_molecular_features=use_molecular_features, - sep=data.get("sep", ","), - ) - - -@dataclass -class H5adConfig: - coord_source: str # "obs" or "obsm" - coord_key: Optional[str] - coord_columns: Tuple[str, str] - cell_id_column: Optional[str] - categorical_label_columns: Sequence[str] - molecular_layer: Optional[str] - molecular_features: Optional[Sequence[str]] - use_molecular_features: bool = True - - @classmethod - def from_dict( - cls, data: Dict[str, object], use_molecular_features: bool - ) -> "H5adConfig": - return cls( - coord_source=data.get("coord_source", "obsm"), - coord_key=data.get("coord_key"), - coord_columns=tuple(data["coord_columns"]), - cell_id_column=data.get("cell_id_column"), - categorical_label_columns=data.get("categorical_label_columns", []), - molecular_layer=data.get("molecular_layer"), - molecular_features=data.get("molecular_features"), - use_molecular_features=use_molecular_features, - ) - - -@dataclass -class SceConfig: - assay_name: Optional[str] - coord_source: str # "colData" or "reducedDims" - coord_key: Optional[str] - coord_columns: Optional[Tuple[str, str]] - cell_id_column: Optional[str] - categorical_label_columns: Sequence[str] - transpose_assay: bool = True - molecular_features: Optional[Sequence[str]] = None - use_molecular_features: bool = True - - @classmethod - def from_dict( - cls, data: Dict[str, object], use_molecular_features: bool - ) -> "SceConfig": - return cls( - assay_name=data.get("assay_name"), - coord_source=data.get("coord_source", "colData"), - coord_key=data.get("coord_key"), - coord_columns=tuple(data["coord_columns"]) - if data.get("coord_columns") - else None, - cell_id_column=data.get("cell_id_column"), - categorical_label_columns=data.get("categorical_label_columns", []), - transpose_assay=data.get("transpose_assay", True), - molecular_features=data.get("molecular_features"), - use_molecular_features=use_molecular_features, - ) - - -def _normalize_columns(df: pd.DataFrame) -> pd.DataFrame: - df.columns = [c.strip() for c in df.columns] - return df - - -def load_csv_table( - path: str | Path, - config: CsvConfig, - sample_id: Optional[str] = None, - region_id: Optional[str] = None, -) -> SpatialOmicsTable: - df = pd.read_csv(path, sep=config.sep) - df = _normalize_columns(df) - - x_col, y_col = config.coord_columns - if x_col not in df.columns or y_col not in df.columns: - raise ValueError(f"Missing coord columns {config.coord_columns} in {path}") - - if config.cell_id_column and config.cell_id_column in df.columns: - cell_ids = df[config.cell_id_column].astype(str).to_numpy() - else: - cell_ids = df.index.astype(str).to_numpy() - - coord_df = df[[x_col, y_col]].astype(float) - coords = coord_df.to_numpy() - - categorical_labels: Dict[str, np.ndarray] = {} - for label_col in config.categorical_label_columns: - if label_col not in df.columns: - raise ValueError(f"Missing categorical label column {label_col} in {path}") - categorical_labels[label_col] = df[label_col].astype(str).to_numpy() - - molecular_features = None - molecular_feature_names: Optional[List[str]] = None - if config.use_molecular_features: - if config.molecular_columns is None: - exclude = {x_col, y_col, *(config.categorical_label_columns or [])} - if config.cell_id_column: - exclude.add(config.cell_id_column) - molecular_cols = [c for c in df.columns if c not in exclude] - else: - molecular_cols = list(config.molecular_columns) - - if not molecular_cols: - raise ValueError(f"No molecular feature columns found in {path}") - - molecular_features = df[molecular_cols].astype(float).to_numpy() - molecular_feature_names = list(molecular_cols) - - return SpatialOmicsTable( - coords=coords, - molecular_features=molecular_features, - molecular_feature_names=molecular_feature_names, - categorical_labels=categorical_labels, - cell_ids=cell_ids, - sample_id=sample_id, - region_id=region_id, - ) - - -def load_h5ad_table( - path: str | Path, - config: H5adConfig, - sample_id: Optional[str] = None, - region_id: Optional[str] = None, -) -> SpatialOmicsTable: - import anndata as ad - - adata = ad.read_h5ad(path) - if config.coord_source == "obsm": - key = config.coord_key or "spatial" - if key not in adata.obsm: - raise ValueError(f"Missing obsm['{key}'] for spatial coords in {path}") - coords = np.asarray(adata.obsm[key], dtype=float) - if coords.ndim != 2 or coords.shape[1] < 2: - raise ValueError( - f"obsm['{key}'] must be 2D with at least 2 columns in {path}" - ) - coords = coords[:, :2] - elif config.coord_source == "obs": - x_col, y_col = config.coord_columns - if x_col not in adata.obs.columns or y_col not in adata.obs.columns: - raise ValueError(f"Missing coord columns {config.coord_columns} in {path}") - coords = adata.obs[[x_col, y_col]].to_numpy(dtype=float) - else: - raise ValueError("h5ad coord_source must be 'obs' or 'obsm'") - - if config.cell_id_column and config.cell_id_column in adata.obs.columns: - cell_ids = adata.obs[config.cell_id_column].astype(str).to_numpy() - else: - cell_ids = adata.obs_names.astype(str).to_numpy() - - categorical_labels: Dict[str, np.ndarray] = {} - for label_col in config.categorical_label_columns: - if label_col not in adata.obs.columns: - raise ValueError(f"Missing categorical label column {label_col} in {path}") - categorical_labels[label_col] = adata.obs[label_col].astype(str).to_numpy() - - molecular_features = None - molecular_feature_names: Optional[List[str]] = None - if config.use_molecular_features: - if config.molecular_features is not None: - requested = list(config.molecular_features) - missing = [f for f in requested if f not in set(adata.var_names)] - if missing: - raise ValueError( - f"Requested molecular features not found in {path}: {missing[:10]}" - ) - adata = adata[:, requested] - - if config.molecular_layer: - data = adata.layers[config.molecular_layer] - else: - data = adata.X - - if hasattr(data, "toarray"): - data = data.toarray() - molecular_features = np.asarray(data, dtype=float) - molecular_feature_names = list(map(str, adata.var_names)) - - return SpatialOmicsTable( - coords=coords, - molecular_features=molecular_features, - molecular_feature_names=molecular_feature_names, - categorical_labels=categorical_labels, - cell_ids=cell_ids, - sample_id=sample_id, - region_id=region_id, - ) - - -def load_sce_table( - path: str | Path, - config: SceConfig, - sample_id: Optional[str] = None, - region_id: Optional[str] = None, -) -> SpatialOmicsTable: - try: - import rpy2.robjects as ro - from rpy2.robjects import numpy2ri, pandas2ri - except ImportError as exc: - raise ImportError( - "rpy2 is required to load SingleCellExperiment objects. " - "Install rpy2 and make sure R + Bioconductor are available." - ) from exc - - pandas2ri.activate() - numpy2ri.activate() - - r = ro.r - sce = r["readRDS"](str(path)) - - molecular_features = None - molecular_feature_names: Optional[List[str]] = None - if config.use_molecular_features: - if config.assay_name: - assay = r["assay"](sce, config.assay_name) - else: - assay = r["assay"](sce) - molecular_features = np.asarray(assay, dtype=float) - try: - rn = r["rownames"](assay) - full_names = [str(x) for x in list(rn)] - except Exception: - full_names = None - - if config.molecular_features is not None: - if full_names is None: - raise ValueError( - "Cannot subset SCE features by name: rownames(assay) not available." - ) - name_to_idx = {n: i for i, n in enumerate(full_names)} - requested = list(config.molecular_features) - missing = [f for f in requested if f not in name_to_idx] - if missing: - raise ValueError( - f"Requested molecular features not found in {path}: {missing[:10]}" - ) - idx = [name_to_idx[f] for f in requested] - molecular_features = molecular_features[idx, :] - molecular_feature_names = requested - else: - molecular_feature_names = full_names - if config.transpose_assay: - molecular_features = molecular_features.T - - coldata_df = None - try: - coldata_df = pandas2ri.rpy2py(r["as.data.frame"](r["colData"](sce))) - except Exception: - coldata_df = None - - if config.coord_source == "colData": - if coldata_df is None: - raise ValueError( - "coord_source=colData but colData could not be loaded from SCE." - ) - if not config.coord_columns: - raise ValueError("coord_columns must be set when coord_source=colData.") - x_col, y_col = config.coord_columns - if x_col not in coldata_df.columns or y_col not in coldata_df.columns: - raise ValueError( - f"Missing coord columns {config.coord_columns} in colData for {path}" - ) - coords = coldata_df[[x_col, y_col]].to_numpy(dtype=float) - elif config.coord_source == "reducedDims": - if not config.coord_key: - raise ValueError("coord_key must be set when coord_source=reducedDims.") - reduced = r["reducedDim"](sce, config.coord_key) - coords = np.asarray(reduced, dtype=float) - if coords.shape[1] < 2: - raise ValueError(f"reducedDim '{config.coord_key}' has <2 columns.") - coords = coords[:, :2] - else: - raise ValueError("coord_source must be 'colData' or 'reducedDims'.") - - if ( - config.cell_id_column - and coldata_df is not None - and config.cell_id_column in coldata_df.columns - ): - cell_ids = coldata_df[config.cell_id_column].astype(str).to_numpy() - else: - cell_ids = np.asarray(r["colnames"](sce), dtype=str) - - categorical_labels: Dict[str, np.ndarray] = {} - if config.categorical_label_columns and coldata_df is None: - raise ValueError( - "categorical_label_columns requested but colData could not be loaded." - ) - for label_col in config.categorical_label_columns: - if label_col not in coldata_df.columns: - raise ValueError( - f"Missing categorical label column {label_col} in colData for {path}" - ) - categorical_labels[label_col] = coldata_df[label_col].astype(str).to_numpy() - - return SpatialOmicsTable( - coords=coords, - molecular_features=molecular_features, - molecular_feature_names=molecular_feature_names, - categorical_labels=categorical_labels, - cell_ids=cell_ids, - sample_id=sample_id, - region_id=region_id, - ) - - -def load_polygons_from_path(path: str | Path) -> List[object]: - from shapely import wkt - from shapely.geometry import shape - - path = Path(path) - if not path.exists(): - raise FileNotFoundError(f"Polygon file not found: {path}") - - if path.suffix in {".pkl", ".pickle"}: - with path.open("rb") as f: - polygons = pickle.load(f) - return _ensure_polygon_list(polygons) - - if path.suffix in {".wkt", ".txt"}: - polygons = [] - for line in path.read_text().splitlines(): - if not line.strip(): - continue - polygons.append(wkt.loads(line.strip())) - return _ensure_polygon_list(polygons) - - if path.suffix in {".json", ".geojson"}: - raw = json.loads(path.read_text()) - if raw.get("type") == "FeatureCollection": - polygons = [shape(feature["geometry"]) for feature in raw["features"]] - elif raw.get("type") in {"Polygon", "MultiPolygon"}: - polygons = [shape(raw)] - else: - raise ValueError(f"Unsupported GeoJSON type in {path}") - return _ensure_polygon_list(polygons) - - raise ValueError(f"Unsupported polygon file type: {path.suffix}") - - -def _ensure_polygon_list(polygons: Iterable[object]) -> List[object]: - polygons_list = list(polygons) - if not polygons_list: - raise ValueError("No polygons found in polygon file") - return polygons_list diff --git a/src/grass_mil/data/components/merge.py b/src/grass_mil/data/components/merge.py deleted file mode 100644 index e146633..0000000 --- a/src/grass_mil/data/components/merge.py +++ /dev/null @@ -1,79 +0,0 @@ -"""Merge several graph units into one disconnected graph. - -Used by the pretraining sampler (``runtime.merge_units``: ego-graph roots are -drawn from the union of all regions, as the preprint's BGRL does) and by the -patient-level "merged bag" scoring, where the cores of one patient are pooled -into a single bag. -""" - -from __future__ import annotations - -from typing import Any, Dict, List, Optional, Sequence - -import torch -from torch_geometric.data import Data - -__all__ = ["merge_graph_units"] - -_NODE_KEYS = ("x", "pos", "x_raw", "categorical_codes") -_EDGE_KEYS = ("edge_attr",) -_COPY_FIRST = ("edge_attr_names", "categorical_labels", "categorical_slices", "molecular_feature_names") - - -def merge_graph_units( - units: Sequence[Data], - *, - keep_graph_attrs: Optional[Sequence[str]] = ("graph_y", "graph_w"), -) -> Data: - """Concatenate node and edge tensors with offset ``edge_index``. - - The merged graph carries ``unit_index`` (node -> position in ``units``), - ``unit_region_ids`` / ``unit_sample_ids`` / ``unit_patch_ids`` lists and - ``unit_ptr`` (node offsets). Graph-level attributes listed in - ``keep_graph_attrs`` are stacked per unit when present on every unit. - """ - if len(units) == 0: - raise ValueError("merge_graph_units needs at least one unit.") - node_parts: Dict[str, List[torch.Tensor]] = {} - edge_parts: Dict[str, List[torch.Tensor]] = {} - edge_index_parts: List[torch.Tensor] = [] - unit_index: List[torch.Tensor] = [] - offsets = [0] - for i, unit in enumerate(units): - n = int(unit.num_nodes) - for key in _NODE_KEYS: - val = getattr(unit, key, None) - if isinstance(val, torch.Tensor): - node_parts.setdefault(key, []).append(val) - for key in _EDGE_KEYS: - val = getattr(unit, key, None) - if isinstance(val, torch.Tensor): - edge_parts.setdefault(key, []).append(val) - ei = unit.edge_index - edge_index_parts.append(ei + offsets[-1] if ei is not None else torch.zeros(2, 0, dtype=torch.long)) - unit_index.append(torch.full((n,), i, dtype=torch.long)) - offsets.append(offsets[-1] + n) - merged = Data(edge_index=torch.cat(edge_index_parts, dim=1)) - for key, parts in node_parts.items(): - if len(parts) != len(units): - raise ValueError(f"Node attribute '{key}' is missing on some units; cannot merge.") - setattr(merged, key, torch.cat(parts, dim=0)) - for key, parts in edge_parts.items(): - if len(parts) != len(units): - raise ValueError(f"Edge attribute '{key}' is missing on some units; cannot merge.") - setattr(merged, key, torch.cat(parts, dim=0)) - merged.num_nodes = offsets[-1] - merged.unit_index = torch.cat(unit_index) - merged.unit_ptr = torch.tensor(offsets, dtype=torch.long) - merged.unit_region_ids = [None if getattr(u, "region_id", None) is None else str(u.region_id) for u in units] - merged.unit_sample_ids = [None if getattr(u, "sample_id", None) is None else str(u.sample_id) for u in units] - merged.unit_patch_ids = [None if getattr(u, "patch_id", None) is None else str(u.patch_id) for u in units] - for key in _COPY_FIRST: - val = getattr(units[0], key, None) - if val is not None: - setattr(merged, key, val) - for key in keep_graph_attrs or (): - vals = [getattr(u, key, None) for u in units] - if all(isinstance(v, torch.Tensor) for v in vals): - setattr(merged, f"unit_{key}", torch.cat([v.view(1, -1) for v in vals], dim=0)) - return merged diff --git a/src/grass_mil/data/components/patching.py b/src/grass_mil/data/components/patching.py deleted file mode 100644 index 9e93764..0000000 --- a/src/grass_mil/data/components/patching.py +++ /dev/null @@ -1,91 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import List, Sequence - -import numpy as np - - -@dataclass -class TileConfig: - tile_size_um: float - stride_um: float - min_cells: int - - @classmethod - def from_dict(cls, data: dict) -> "TileConfig": - return cls( - tile_size_um=data["tile_size_um"], - stride_um=data["stride_um"], - min_cells=data["min_cells"], - ) - - -def build_patches_from_polygons( - coords: np.ndarray, polygons: Sequence[object] -) -> List[np.ndarray]: - from shapely import vectorized - from shapely.geometry import Point - from shapely.prepared import prep - - patches: List[np.ndarray] = [] - xs = coords[:, 0] - ys = coords[:, 1] - - for poly in polygons: - try: - mask = vectorized.contains(poly, xs, ys) - except Exception: - prepared = prep(poly) - mask = np.array([prepared.contains(Point(x, y)) for x, y in coords]) - indices = np.where(mask)[0] - patches.append(indices) - return patches - - -def filter_coords_by_union( - coords: np.ndarray, polygons: Sequence[object] -) -> np.ndarray: - from shapely import vectorized - from shapely.geometry import Point - from shapely.ops import unary_union - from shapely.prepared import prep - - union = unary_union(list(polygons)) - xs = coords[:, 0] - ys = coords[:, 1] - try: - mask = vectorized.contains(union, xs, ys) - except Exception: - prepared = prep(union) - mask = np.array([prepared.contains(Point(x, y)) for x, y in coords]) - return np.where(mask)[0] - - -def build_grid_tiles(coords: np.ndarray, tile_config: TileConfig) -> List[np.ndarray]: - if coords.size == 0: - return [] - - x_min, y_min = coords.min(axis=0) - x_max, y_max = coords.max(axis=0) - tile_size = tile_config.tile_size_um - stride = tile_config.stride_um - - patches: List[np.ndarray] = [] - x_starts = np.arange(x_min, x_max + 1e-6, stride) - y_starts = np.arange(y_min, y_max + 1e-6, stride) - - for x0 in x_starts: - for y0 in y_starts: - x1 = x0 + tile_size - y1 = y0 + tile_size - mask = ( - (coords[:, 0] >= x0) - & (coords[:, 0] < x1) - & (coords[:, 1] >= y0) - & (coords[:, 1] < y1) - ) - idx = np.where(mask)[0] - if idx.size >= tile_config.min_cells: - patches.append(idx) - return patches diff --git a/src/grass_mil/data/components/precompute.py b/src/grass_mil/data/components/precompute.py deleted file mode 100644 index 3138c34..0000000 --- a/src/grass_mil/data/components/precompute.py +++ /dev/null @@ -1,1181 +0,0 @@ -from __future__ import annotations - -import json -import warnings -from dataclasses import dataclass -from pathlib import Path -from typing import Dict, List, Optional, Sequence, Set, Tuple - -import numpy as np -import torch -from torch_geometric.data import Data - -from .feature_reducers import FeatureReducerConfig, get_feature_reducer -from .graph_builders import GraphBuilderConfig, GraphBuildResult, get_graph_builder -from .loaders import ( - CsvConfig, - H5adConfig, - SceConfig, - load_csv_table, - load_h5ad_table, - load_polygons_from_path, - load_sce_table, -) -from .patching import ( - TileConfig, - build_grid_tiles, - build_patches_from_polygons, - filter_coords_by_union, -) -from .spatial_types import SpatialOmicsTable - - -@dataclass -class ManifestConfig: - sample_id: str = "sample_id" - input_path: str = "input_path" - input_type: str = "input_type" - region_id: str = "region_id" - polygons_path: str = "polygons_path" - - @classmethod - def from_dict(cls, data: Optional[Dict[str, object]]) -> "ManifestConfig": - return cls(**(data or {})) - - -@dataclass -class GraphLabelConfig: - label_file: Optional[str] = None - id_column: str = "id" - tasks: Optional[Sequence[str]] = None - scope: str = "region" # region | patch | sample - # Optional grouping column (patient / donor). It is never used as a target; - # the value is attached as ``data.patient_id`` and recorded in the processed - # index so patient-grouped splits can be built without loading graphs. - patient_column: Optional[str] = None - - @classmethod - def from_dict(cls, data: Optional[Dict[str, object]]) -> "GraphLabelConfig": - return cls(**(data or {})) - - -@dataclass -class PrecomputeConfig: - raw_manifest_path: str - processed_dir: str - coord_scale_um: float - sample_unit: str - reducer_scope: str - reducer_fit_mode: str - keep_raw_molecular: bool - force: bool - min_cells: int - use_molecular_features: bool - split_by: str - split_ratios: Tuple[float, float, float] - split_seed: int - loocv: "LoocvConfig" - - @classmethod - def from_args( - cls, - raw_manifest_path: str, - processed_dir: str, - coord_scale_um: float, - sample_unit: str, - reducer_scope: str, - reducer_fit_mode: str, - keep_raw_molecular: bool, - force: bool, - min_cells: int, - use_molecular_features: bool, - split_by: str, - split_ratios: Tuple[float, float, float], - split_seed: int, - loocv: Optional[Dict[str, object]] = None, - ) -> "PrecomputeConfig": - return cls( - raw_manifest_path=raw_manifest_path, - processed_dir=processed_dir, - coord_scale_um=coord_scale_um, - sample_unit=sample_unit, - reducer_scope=reducer_scope, - reducer_fit_mode=reducer_fit_mode, - keep_raw_molecular=keep_raw_molecular, - force=force, - min_cells=min_cells, - use_molecular_features=use_molecular_features, - split_by=split_by, - split_ratios=split_ratios, - split_seed=split_seed, - loocv=LoocvConfig.from_dict(loocv), - ) - - -@dataclass -class LoocvConfig: - enabled: bool = False - fold_unit: str = "region" # region | sample - holdout_id: Optional[str] = None - validation_strategy: str = ( - "heldout_fold_items" # heldout_fold_items | patches_from_train_items - ) - val_ratio: float = 0.1 - seed: int = 42 - - @classmethod - def from_dict(cls, data: Optional[Dict[str, object]]) -> "LoocvConfig": - return cls(**(data or {})) - - -@dataclass -class CategoricalFeatureConfig: - include_labels: Sequence[str] - - @classmethod - def from_dict(cls, data: Optional[Dict[str, object]]) -> "CategoricalFeatureConfig": - data = data or {} - return cls(include_labels=data.get("include_labels", [])) - - -@dataclass -class PreparedSample: - table_idx: int - table: SpatialOmicsTable - sample_id: str - region_id: Optional[str] - patch_indices: List[np.ndarray] - - -@dataclass -class GraphUnitSpec: - table_idx: int - sample_id: str - region_id: Optional[str] - patch_idx: int - patch_id: str - indices: np.ndarray - split: str - - -class SpatialOmicsPreprocessor: - def __init__( - self, - manifest_config: ManifestConfig, - precompute_config: PrecomputeConfig, - csv_config: CsvConfig, - h5ad_config: H5adConfig, - sce_config: SceConfig, - graph_builder_config: GraphBuilderConfig, - feature_reducer_config: FeatureReducerConfig, - categorical_feature_config: CategoricalFeatureConfig, - tile_config: TileConfig, - graph_label_config: Optional[GraphLabelConfig] = None, - ) -> None: - self.manifest_config = manifest_config - self.precompute_config = precompute_config - self.csv_config = csv_config - self.h5ad_config = h5ad_config - self.sce_config = sce_config - self.graph_builder_config = graph_builder_config - self.feature_reducer_config = feature_reducer_config - self.categorical_feature_config = categorical_feature_config - self.tile_config = tile_config - self.graph_label_config = graph_label_config or GraphLabelConfig() - self._manifest_dir: Optional[Path] = None - - def precompute(self) -> Path: - processed_dir = Path(self.precompute_config.processed_dir) - processed_dir.mkdir(parents=True, exist_ok=True) - index_path = processed_dir / "processed_index.json" - metadata_path = processed_dir / "metadata.json" - - if index_path.exists() and not self.precompute_config.force: - _validate_existing_coord_scale( - index_path=index_path, - metadata_path=metadata_path, - expected_coord_scale_um=float(self.precompute_config.coord_scale_um), - ) - return index_path - - manifest_rows = self._load_manifest( - Path(self.precompute_config.raw_manifest_path) - ) - raw_tables = self._load_tables(manifest_rows) - - prepared = self._prepare_samples(manifest_rows, raw_tables) - units = self._build_unit_specs(prepared) - label_maps = self._build_label_maps([item.table for item in prepared]) - - graph_label_maps: Dict[str, Dict[str, int]] = {} - graph_label_df = self._load_graph_labels() - - reducer_state, reducer = self._prepare_reducer(prepared, units) - - table_embeddings: Dict[int, Optional[np.ndarray]] = {} - for sample in prepared: - table_embeddings[sample.table_idx] = self._compute_molecular_embedding( - table=sample.table, - reducer=reducer, - reducer_state=reducer_state, - sample_id=sample.sample_id, - ) - - entries: List[Dict[str, object]] = [] - for unit in units: - sample = prepared[unit.table_idx] - embedded = table_embeddings[unit.table_idx] - data = self._build_pyg_data( - table=sample.table, - indices=unit.indices, - embedded=embedded[unit.indices] if embedded is not None else None, - sample_id=unit.sample_id, - region_id=unit.region_id, - patch_id=unit.patch_id, - label_maps=label_maps, - graph_label_df=graph_label_df, - graph_label_maps=graph_label_maps, - ) - - graph_path = processed_dir / "graphs" / unit.sample_id - graph_path.mkdir(parents=True, exist_ok=True) - data_path = graph_path / f"{unit.patch_id}.pt" - torch.save(data, data_path) - entries.append( - { - "path": str(data_path), - "sample_id": unit.sample_id, - "region_id": unit.region_id, - "patch_id": unit.patch_id, - "split": unit.split, - "patient_id": getattr(data, "patient_id", None), - "graph_y": ( - data.graph_y.view(-1).tolist() if hasattr(data, "graph_y") else None - ), - } - ) - - index_payload = { - "entries": entries, - "label_maps": label_maps, - "graph_label_maps": graph_label_maps, - "reducer_state": reducer_state, - "preprocessing": { - "coord_scale_um": float(self.precompute_config.coord_scale_um), - "coord_unit": "um", - }, - "categorical_features": { - "include_labels": list(self.categorical_feature_config.include_labels), - }, - "split": { - "split_by": self.precompute_config.split_by, - "ratios": list(self.precompute_config.split_ratios), - "seed": int(self.precompute_config.split_seed), - "loocv": { - "enabled": bool(self.precompute_config.loocv.enabled), - "fold_unit": str(self.precompute_config.loocv.fold_unit), - "holdout_id": self.precompute_config.loocv.holdout_id, - "validation_strategy": str( - self.precompute_config.loocv.validation_strategy - ), - "val_ratio": float(self.precompute_config.loocv.val_ratio), - "seed": int(self.precompute_config.loocv.seed), - }, - }, - } - index_path.write_text(json.dumps(index_payload, indent=2)) - - metadata = self._build_metadata( - entries=entries, - tables=[item.table for item in prepared], - label_maps=label_maps, - graph_label_maps=graph_label_maps, - ) - metadata_path.write_text(json.dumps(metadata, indent=2)) - return index_path - - def _prepare_samples( - self, - manifest_rows: List[Dict[str, str]], - tables: List[SpatialOmicsTable], - ) -> List[PreparedSample]: - prepared: List[PreparedSample] = [] - for table_idx, (row, table) in enumerate(zip(manifest_rows, tables)): - sample_id = table.sample_id or str(row[self.manifest_config.sample_id]) - region_id = table.region_id or row.get(self.manifest_config.region_id) - polygons_path = self._resolve_manifest_relative_path( - row.get(self.manifest_config.polygons_path) - ) - - table = _apply_coord_scale(table, self.precompute_config.coord_scale_um) - table, polygons = self._apply_polygons( - table, - polygons_path, - coord_scale_um=float(self.precompute_config.coord_scale_um), - ) - patch_indices = self._compute_patch_indices(table, polygons) - prepared.append( - PreparedSample( - table_idx=table_idx, - table=table, - sample_id=sample_id, - region_id=region_id, - patch_indices=patch_indices, - ) - ) - return prepared - - def _build_unit_specs( - self, prepared: Sequence[PreparedSample] - ) -> List[GraphUnitSpec]: - units: List[GraphUnitSpec] = [] - for sample in prepared: - for patch_idx, indices in enumerate(sample.patch_indices): - patch_id = _build_patch_id( - sample.sample_id, sample.region_id, patch_idx - ) - units.append( - GraphUnitSpec( - table_idx=sample.table_idx, - sample_id=sample.sample_id, - region_id=sample.region_id, - patch_idx=patch_idx, - patch_id=patch_id, - indices=indices, - split="train", # assigned below - ) - ) - - split_labels = _assign_split_labels( - units=units, - split_by=self.precompute_config.split_by, - split_ratios=self.precompute_config.split_ratios, - seed=self.precompute_config.split_seed, - loocv_config=self.precompute_config.loocv, - ) - for i, split in enumerate(split_labels): - units[i].split = split - return units - - def _load_tables( - self, manifest_rows: List[Dict[str, str]] - ) -> List[SpatialOmicsTable]: - return [self._load_table(row) for row in manifest_rows] - - def _prepare_reducer( - self, - prepared: Sequence[PreparedSample], - units: Sequence[GraphUnitSpec], - ) -> tuple[Dict[str, object], object]: - reducer_state: Dict[str, object] = { - "scope": self.precompute_config.reducer_scope, - "fit_mode": self.precompute_config.reducer_fit_mode, - } - reducer = get_feature_reducer(self.feature_reducer_config) - - reducer_state["by_sample"] = {} - if not self.precompute_config.use_molecular_features: - return reducer_state, reducer - - if self.precompute_config.reducer_scope != "dataset": - return reducer_state, reducer - - fit_mode = str(self.precompute_config.reducer_fit_mode).strip().lower() - if fit_mode not in {"global", "train_only"}: - raise ValueError( - "feature_reducer.fit_mode must be one of ['global', 'train_only']." - ) - - if fit_mode == "global": - features = [ - sample.table.molecular_features - for sample in prepared - if sample.table.molecular_features is not None - ] - fitted_on = "global" - else: - selected_by_table: Dict[int, List[np.ndarray]] = {} - for unit in units: - if unit.split != "train": - continue - selected_by_table.setdefault(unit.table_idx, []).append(unit.indices) - features = [] - for sample in prepared: - if sample.table.molecular_features is None: - continue - train_indices = selected_by_table.get(sample.table_idx, []) - if not train_indices: - continue - uniq = np.unique(np.concatenate(train_indices, axis=0)) - if uniq.size == 0: - continue - features.append(sample.table.molecular_features[uniq]) - fitted_on = "train_only" - - if not features: - raise ValueError( - "No molecular features available for reducer fitting with the current split and fit_mode." - ) - - reducer.fit(np.concatenate(features, axis=0)) - reducer_state["global"] = reducer.state_dict() - reducer_state["fitted_on"] = fitted_on - return reducer_state, reducer - - def _apply_polygons( - self, - table: SpatialOmicsTable, - polygons_path: Optional[str], - *, - coord_scale_um: float, - ) -> tuple[SpatialOmicsTable, Optional[Sequence[object]]]: - polygons = load_polygons_from_path(polygons_path) if polygons_path else None - polygons = _apply_coord_scale_to_polygons(polygons, coord_scale_um) - if polygons: - union_idx = filter_coords_by_union(table.coords, polygons) - table = _subset_table(table, union_idx) - return table, polygons - - def _compute_molecular_embedding( - self, - table: SpatialOmicsTable, - reducer, - reducer_state: Dict[str, object], - sample_id: str, - ) -> Optional[np.ndarray]: - if not self.precompute_config.use_molecular_features: - return None - reducer_instance = reducer - if self.precompute_config.reducer_scope == "sample": - reducer_instance = get_feature_reducer(self.feature_reducer_config) - reducer_instance.fit(table.molecular_features) - reducer_state["by_sample"][sample_id] = reducer_instance.state_dict() - return reducer_instance.transform(table.molecular_features) - - def _compute_patch_indices( - self, table: SpatialOmicsTable, polygons: Optional[Sequence[object]] - ) -> List[np.ndarray]: - if polygons and self.precompute_config.sample_unit != "full": - patch_indices = build_patches_from_polygons(table.coords, polygons) - elif self.precompute_config.sample_unit == "tile": - patch_indices = build_grid_tiles(table.coords, self.tile_config) - else: - patch_indices = [np.arange(table.coords.shape[0])] - return [ - idx for idx in patch_indices if idx.size >= self.precompute_config.min_cells - ] - - def _load_manifest(self, path: Path) -> List[Dict[str, str]]: - import pandas as pd - - if not path.exists(): - raise FileNotFoundError(f"Manifest file not found: {path}") - self._manifest_dir = path.parent - df = pd.read_csv(path) - df.columns = [c.strip() for c in df.columns] - - required = [ - self.manifest_config.sample_id, - self.manifest_config.input_path, - self.manifest_config.input_type, - ] - for col in required: - if col not in df.columns: - raise ValueError(f"Manifest missing required column: {col}") - self._validate_manifest_regions(df) - return df.to_dict(orient="records") - - def _validate_manifest_regions(self, df) -> None: - sample_col = self.manifest_config.sample_id - region_col = self.manifest_config.region_id - if region_col not in df.columns: - return - - region_vals = df[region_col].fillna("").astype(str).str.strip() - non_empty_regions = region_vals[region_vals.ne("")] - duplicated_regions = non_empty_regions[non_empty_regions.duplicated(keep=False)] - if not duplicated_regions.empty: - duplicated_ids = sorted(set(duplicated_regions.tolist())) - raise ValueError( - "Manifest region_id values must be globally unique. " - f"Repeated region_id values: {duplicated_ids}." - ) - - duplicates = df[sample_col].duplicated(keep=False) - if not duplicates.any(): - return - - for sample_id, group in df[duplicates].groupby(sample_col): - sample_region_vals = group[region_col].fillna("").astype(str).str.strip() - if sample_region_vals.eq("").any(): - raise ValueError( - f"Manifest has duplicate sample_id '{sample_id}' with empty region_id." - ) - if sample_region_vals.duplicated().any(): - raise ValueError( - f"Manifest has duplicate sample_id '{sample_id}' with repeated region_id values." - ) - - def _load_table(self, row: Dict[str, str]) -> SpatialOmicsTable: - sample_id = row[self.manifest_config.sample_id] - region_id = row.get(self.manifest_config.region_id) - input_path = row[self.manifest_config.input_path] - input_path = self._resolve_manifest_relative_path(input_path) - input_type = row[self.manifest_config.input_type].lower() - - if input_type in {"csv", "tsv"}: - cfg = self.csv_config - if input_type == "tsv": - cfg = CsvConfig( - coord_columns=cfg.coord_columns, - cell_id_column=cfg.cell_id_column, - categorical_label_columns=cfg.categorical_label_columns, - molecular_columns=cfg.molecular_columns, - use_molecular_features=cfg.use_molecular_features, - sep="\t", - ) - return load_csv_table( - input_path, cfg, sample_id=sample_id, region_id=region_id - ) - if input_type == "h5ad": - return load_h5ad_table( - input_path, self.h5ad_config, sample_id=sample_id, region_id=region_id - ) - if input_type in {"sce", "rds"}: - return load_sce_table( - input_path, self.sce_config, sample_id=sample_id, region_id=region_id - ) - raise ValueError(f"Unsupported input type: {input_type}") - - def _resolve_manifest_relative_path(self, value: Optional[str]) -> Optional[str]: - if value is None: - return None - # handle pandas NaN and common empty sentinels - try: - import pandas as pd - - if pd.isna(value): - return None - except Exception: - pass - - value = str(value).strip() - if value == "": - return None - if value.lower() in {"nan", "none", "null"}: - return None - p = Path(value) - if p.is_absolute(): - return str(p) - if self._manifest_dir is None: - return str(p) - return str((self._manifest_dir / p).resolve()) - - def _build_label_maps( - self, tables: Sequence[SpatialOmicsTable] - ) -> Dict[str, Dict[str, int]]: - label_maps: Dict[str, Dict[str, int]] = {} - for table in tables: - for label_name, labels in table.categorical_labels.items(): - label_maps.setdefault(label_name, {}) - for label in labels: - if label not in label_maps[label_name]: - label_maps[label_name][label] = len(label_maps[label_name]) - return label_maps - - def _load_graph_labels(self): - if not self.graph_label_config.label_file: - return None - import pandas as pd - - label_path = Path(self.graph_label_config.label_file) - if not label_path.exists(): - raise FileNotFoundError(f"Graph label file not found: {label_path}") - df = pd.read_csv(label_path) - if self.graph_label_config.id_column not in df.columns: - raise ValueError( - f"Graph label file missing id column '{self.graph_label_config.id_column}'" - ) - df = df.set_index(self.graph_label_config.id_column) - return df - - def _build_pyg_data( - self, - table: SpatialOmicsTable, - indices: np.ndarray, - embedded: Optional[np.ndarray], - sample_id: str, - region_id: Optional[str], - patch_id: str, - label_maps: Dict[str, Dict[str, int]], - graph_label_df, - graph_label_maps: Dict[str, Dict[str, int]], - ) -> Data: - coords = table.coords[indices] - graph_result: GraphBuildResult = get_graph_builder( - self.graph_builder_config - ).build(coords) - - data_kwargs = { - "pos": torch.from_numpy(coords).float(), - "edge_index": graph_result.edge_index, - } - if embedded is not None: - data_kwargs["x"] = torch.from_numpy(embedded).float() - else: - data_kwargs["x"] = torch.zeros((coords.shape[0], 0), dtype=torch.float) - data = Data(**data_kwargs) - if graph_result.edge_attr is not None: - data.edge_attr = graph_result.edge_attr.float() - data.edge_attr_names = graph_result.edge_attr_names - - if ( - self.precompute_config.keep_raw_molecular - and table.molecular_features is not None - ): - data.x_raw = torch.from_numpy(table.molecular_features[indices]).float() - if ( - embedded is not None - and table.molecular_feature_names is not None - and self.feature_reducer_config.name == "identity" - and len(table.molecular_feature_names) == int(embedded.shape[1]) - ): - # Column names of `x` (identity reducer only), so augmentations can - # address continuous features by name (e.g. feature_noise_columns: [SIZE]). - data.molecular_feature_names = list(map(str, table.molecular_feature_names)) - - self._attach_categorical_indices( - data, table.categorical_labels, label_maps, indices - ) - - if graph_label_df is not None: - label_id = _select_graph_label_id( - scope=self.graph_label_config.scope, - sample_id=sample_id, - region_id=region_id, - patch_id=patch_id, - ) - if label_id in graph_label_df.index: - row = graph_label_df.loc[label_id] - patient_column = self.graph_label_config.patient_column - tasks = self.graph_label_config.tasks or [ - c for c in row.index if c != patient_column - ] - if patient_column: - if patient_column not in row.index: - raise ValueError( - f"graph_labels.patient_column '{patient_column}' not found in " - f"{self.graph_label_config.label_file}" - ) - data.patient_id = str(row[patient_column]) - graph_values = [] - for task in tasks: - value = row[task] - if isinstance(value, str): - graph_label_maps.setdefault(task, {}) - mapping = graph_label_maps[task] - if value not in mapping: - mapping[value] = len(mapping) - graph_values.append(mapping[value]) - else: - graph_values.append(value) - data.graph_y = torch.tensor(graph_values, dtype=torch.float).unsqueeze( - 0 - ) - data.graph_label_names = tasks - - data.sample_id = sample_id - data.region_id = region_id - data.patch_id = patch_id - return data - - def _attach_categorical_indices( - self, - data: Data, - labels: Dict[str, np.ndarray], - label_maps: Dict[str, Dict[str, int]], - indices: np.ndarray, - ) -> None: - config = self.categorical_feature_config - if not config.include_labels: - return - - encoded_cols: List[np.ndarray] = [] - slices: Dict[str, int] = {} - for col_idx, label_name in enumerate(config.include_labels): - if label_name not in labels: - raise ValueError(f"Categorical label '{label_name}' not found in data.") - mapping = label_maps.get(label_name, {}) - encoded = np.array( - [mapping[val] for val in labels[label_name][indices]], dtype=np.int64 - ) - encoded_cols.append(encoded) - slices[label_name] = col_idx - - categorical_codes = np.stack(encoded_cols, axis=1) - data.categorical_codes = torch.from_numpy(categorical_codes).long() - data.categorical_labels = list(config.include_labels) - data.categorical_slices = slices - - def _build_metadata( - self, - entries: List[Dict[str, object]], - tables: List[SpatialOmicsTable], - label_maps: Dict[str, Dict[str, int]], - graph_label_maps: Dict[str, Dict[str, int]], - ) -> Dict[str, object]: - molecular_dim = 0 - molecular_feature_names = None - if self.precompute_config.use_molecular_features: - for table in tables: - if table.molecular_features is not None: - molecular_dim = int(table.molecular_features.shape[1]) - molecular_feature_names = ( - list(table.molecular_feature_names) - if table.molecular_feature_names is not None - else None - ) - break - categorical_cardinalities = { - label: len(mapping) for label, mapping in label_maps.items() - } - return { - "num_graphs": len(entries), - "num_samples": len({entry["sample_id"] for entry in entries}), - "molecular_feature_dim": molecular_dim, - "molecular_feature_names": molecular_feature_names, - "coord_scale_um": float(self.precompute_config.coord_scale_um), - "coord_unit": "um", - "categorical_labels": list(self.categorical_feature_config.include_labels), - "categorical_cardinalities": categorical_cardinalities, - "label_maps": label_maps, - "graph_label_maps": graph_label_maps, - "use_molecular_features": self.precompute_config.use_molecular_features, - "reducer_scope": self.precompute_config.reducer_scope, - "reducer_fit_mode": self.precompute_config.reducer_fit_mode, - "sample_unit": self.precompute_config.sample_unit, - "split_by": self.precompute_config.split_by, - "split_ratios": list(self.precompute_config.split_ratios), - "split_seed": int(self.precompute_config.split_seed), - } - - -def _apply_coord_scale( - table: SpatialOmicsTable, coord_scale_um: float -) -> SpatialOmicsTable: - if coord_scale_um == 1.0: - return table - scaled = table.coords * coord_scale_um - return SpatialOmicsTable( - coords=scaled, - molecular_features=table.molecular_features, - molecular_feature_names=table.molecular_feature_names, - categorical_labels=table.categorical_labels, - cell_ids=table.cell_ids, - sample_id=table.sample_id, - region_id=table.region_id, - ) - - -def _apply_coord_scale_to_polygons( - polygons: Optional[Sequence[object]], coord_scale_um: float -) -> Optional[List[object]]: - if not polygons: - return polygons - if coord_scale_um == 1.0: - return list(polygons) - from shapely import affinity - - return [ - affinity.scale(poly, xfact=coord_scale_um, yfact=coord_scale_um, origin=(0.0, 0.0)) - for poly in polygons - ] - - -def _read_coord_scale_from_payload(payload: object) -> Optional[float]: - if not isinstance(payload, dict): - return None - preprocessing = payload.get("preprocessing") - if isinstance(preprocessing, dict): - value = preprocessing.get("coord_scale_um") - if value is not None: - try: - return float(value) - except (TypeError, ValueError): - return None - value = payload.get("coord_scale_um") - if value is not None: - try: - return float(value) - except (TypeError, ValueError): - return None - return None - - -def _validate_existing_coord_scale( - *, - index_path: Path, - metadata_path: Path, - expected_coord_scale_um: float, -) -> None: - stored_scale = None - try: - stored_scale = _read_coord_scale_from_payload(json.loads(index_path.read_text())) - except Exception: - stored_scale = None - if stored_scale is None and metadata_path.exists(): - try: - stored_scale = _read_coord_scale_from_payload(json.loads(metadata_path.read_text())) - except Exception: - stored_scale = None - if stored_scale is None: - warnings.warn( - "Existing processed artifacts do not store coord_scale_um metadata. " - "Cannot verify unit consistency for cache reuse; consider force_precompute=true.", - UserWarning, - stacklevel=2, - ) - return - if not np.isclose(float(stored_scale), float(expected_coord_scale_um)): - raise ValueError( - "Existing processed artifacts were built with " - f"coord_scale_um={stored_scale}, but current config has " - f"coord_scale_um={expected_coord_scale_um}. " - "Set data.force_precompute=true or use a new data.processed_dir." - ) - - -def _subset_table(table: SpatialOmicsTable, indices: np.ndarray) -> SpatialOmicsTable: - return SpatialOmicsTable( - coords=table.coords[indices], - molecular_features=( - None - if table.molecular_features is None - else table.molecular_features[indices] - ), - molecular_feature_names=table.molecular_feature_names, - categorical_labels={k: v[indices] for k, v in table.categorical_labels.items()}, - cell_ids=table.cell_ids[indices], - sample_id=table.sample_id, - region_id=table.region_id, - ) - - -def _split_indices( - total: int, - ratios: Tuple[float, float, float], - seed: int, -) -> Tuple[List[int], List[int], List[int]]: - train_r, val_r, test_r = ratios - if not np.isclose(train_r + val_r + test_r, 1.0): - raise ValueError("train_val_test_split must sum to 1.0") - if total <= 0: - return [], [], [] - - ratio_arr = np.asarray([train_r, val_r, test_r], dtype=float) - raw_counts = ratio_arr * float(total) - counts = np.floor(raw_counts).astype(int) - - remainder = int(total - counts.sum()) - if remainder > 0: - fractions = raw_counts - counts - fractions[ratio_arr <= 0.0] = -1.0 - order = np.argsort(-fractions) - ptr = 0 - while remainder > 0 and ptr < len(order): - idx = int(order[ptr]) - if fractions[idx] < 0: - break - counts[idx] += 1 - remainder -= 1 - ptr += 1 - while remainder > 0: - counts[0] += 1 - remainder -= 1 - - if ratio_arr[0] > 0.0 and counts[0] == 0: - donor_idx = int(np.argmax(counts[1:]) + 1) if counts[1:].sum() > 0 else -1 - if donor_idx >= 0 and counts[donor_idx] > 0: - counts[donor_idx] -= 1 - counts[0] += 1 - else: - counts[0] = 1 - overflow = int(counts.sum() - total) - for idx in (2, 1): - if overflow <= 0: - break - take = min(overflow, counts[idx]) - counts[idx] -= take - overflow -= take - - # Keep val split non-empty whenever val ratio is non-zero. - if ratio_arr[1] > 0.0 and counts[1] == 0 and total > 0: - donor_idx = -1 - if counts[2] > 1: - donor_idx = 2 - elif counts[0] > 1 or ratio_arr[0] <= 0.0: - donor_idx = 0 - else: - # Tiny totals can make non-empty train and val impossible. - raise ValueError( - "Tiny total makes non-empty train and val impossible. Use a bigger dataset or remove the validation split." - ) - counts[donor_idx] -= 1 - counts[1] += 1 - - indices = np.arange(total) - rng = np.random.default_rng(seed) - rng.shuffle(indices) - - train_end = int(counts[0]) - val_end = train_end + int(counts[1]) - return ( - indices[:train_end].tolist(), - indices[train_end:val_end].tolist(), - indices[val_end:].tolist(), - ) - - -def _assign_split_labels( - units: Sequence[GraphUnitSpec], - split_by: str, - split_ratios: Tuple[float, float, float], - seed: int, - loocv_config: Optional[LoocvConfig] = None, -) -> List[str]: - labels = ["train"] * len(units) - if not units: - return labels - - loocv = loocv_config or LoocvConfig() - if loocv.enabled: - return _assign_loocv_split_labels(units, loocv) - - if split_by == "sample": - by_sample: Dict[str, List[int]] = {} - for idx, unit in enumerate(units): - by_sample.setdefault(unit.sample_id, []).append(idx) - sample_keys = list(by_sample) - train_idx, val_idx, test_idx = _split_indices( - total=len(sample_keys), ratios=split_ratios, seed=seed - ) - for i in train_idx: - for unit_idx in by_sample[sample_keys[i]]: - labels[unit_idx] = "train" - for i in val_idx: - for unit_idx in by_sample[sample_keys[i]]: - labels[unit_idx] = "val" - for i in test_idx: - for unit_idx in by_sample[sample_keys[i]]: - labels[unit_idx] = "test" - return labels - - if split_by == "region": - by_region: Dict[Tuple[str, Optional[str]], List[int]] = {} - for idx, unit in enumerate(units): - by_region.setdefault((unit.sample_id, unit.region_id), []).append(idx) - region_keys = list(by_region) - train_idx, val_idx, test_idx = _split_indices( - total=len(region_keys), ratios=split_ratios, seed=seed - ) - for i in train_idx: - for unit_idx in by_region[region_keys[i]]: - labels[unit_idx] = "train" - for i in val_idx: - for unit_idx in by_region[region_keys[i]]: - labels[unit_idx] = "val" - for i in test_idx: - for unit_idx in by_region[region_keys[i]]: - labels[unit_idx] = "test" - return labels - - if split_by != "patch": - raise ValueError("split_by must be one of ['sample', 'region', 'patch']") - - train_idx, val_idx, test_idx = _split_indices( - total=len(units), ratios=split_ratios, seed=seed - ) - for i in train_idx: - labels[i] = "train" - for i in val_idx: - labels[i] = "val" - for i in test_idx: - labels[i] = "test" - return labels - - -def _assign_loocv_split_labels( - units: Sequence[GraphUnitSpec], loocv: LoocvConfig -) -> List[str]: - if loocv.fold_unit not in {"region", "sample"}: - raise ValueError("split.loocv.fold_unit must be one of ['region', 'sample'].") - if loocv.validation_strategy not in { - "heldout_fold_items", - "patches_from_train_items", - }: - raise ValueError( - "split.loocv.validation_strategy must be one of " - "['heldout_fold_items', 'patches_from_train_items']." - ) - if loocv.val_ratio < 0.0 or loocv.val_ratio >= 1.0: - raise ValueError("split.loocv.val_ratio must be in [0.0, 1.0).") - if not loocv.holdout_id: - raise ValueError( - "split.loocv.holdout_id must be set when split.loocv.enabled=true." - ) - - by_fold_item: Dict[object, List[int]] = {} - fold_alias_to_key: Dict[str, object] = {} - ambiguous_region_aliases = set() - - for idx, unit in enumerate(units): - key = ( - (unit.sample_id, unit.region_id) - if loocv.fold_unit == "region" - else unit.sample_id - ) - by_fold_item.setdefault(key, []).append(idx) - - canonical = _loocv_fold_id_from_key(key=key, fold_unit=loocv.fold_unit) - fold_alias_to_key[canonical] = key - if loocv.fold_unit == "region" and unit.region_id is not None: - alias = str(unit.region_id) - existing = fold_alias_to_key.get(alias) - if existing is None: - fold_alias_to_key[alias] = key - elif existing != key: - ambiguous_region_aliases.add(alias) - - for alias in ambiguous_region_aliases: - if alias in fold_alias_to_key: - del fold_alias_to_key[alias] - - if len(by_fold_item) < 2: - raise ValueError( - "LOOCV requires at least two fold items after grouping by " - f"'{loocv.fold_unit}'. Found {len(by_fold_item)}." - ) - if loocv.holdout_id not in fold_alias_to_key: - available = sorted( - _loocv_fold_id_from_key(key=key, fold_unit=loocv.fold_unit) - for key in by_fold_item - ) - raise ValueError( - "split.loocv.holdout_id not found in grouped fold items. " - f"Provided '{loocv.holdout_id}'. Available canonical ids: {available}" - ) - - holdout_key = fold_alias_to_key[loocv.holdout_id] - labels = ["train"] * len(units) - for unit_idx in by_fold_item[holdout_key]: - labels[unit_idx] = "test" - - remaining_keys = [key for key in by_fold_item if key != holdout_key] - if not remaining_keys: - raise ValueError("LOOCV left no non-test fold items.") - - if loocv.validation_strategy == "heldout_fold_items": - val_group_indices = _pick_validation_fold_items( - total=len(remaining_keys), - val_ratio=float(loocv.val_ratio), - seed=int(loocv.seed), - ) - for group_idx in val_group_indices: - for unit_idx in by_fold_item[remaining_keys[group_idx]]: - labels[unit_idx] = "val" - return labels - - warnings.warn( - "Using patch-level validation from training fold items. " - "This is useful in low-data settings, but validation is less independent " - "than held-out fold-item validation.", - stacklevel=2, - ) - candidate_indices = [ - unit_idx for key in remaining_keys for unit_idx in by_fold_item[key] - ] - train_idx, val_idx, _ = _split_indices( - total=len(candidate_indices), - ratios=(1.0 - float(loocv.val_ratio), float(loocv.val_ratio), 0.0), - seed=int(loocv.seed), - ) - for local_idx in train_idx: - labels[candidate_indices[local_idx]] = "train" - for local_idx in val_idx: - labels[candidate_indices[local_idx]] = "val" - return labels - - -def _pick_validation_fold_items(total: int, val_ratio: float, seed: int) -> List[int]: - if total <= 1: - raise ValueError( - "LOOCV validation strategy 'heldout_fold_items' requires at least two " - "non-test fold items. Use 'patches_from_train_items' for tiny-data regimes." - ) - _, val_idx, _ = _split_indices( - total=total, - ratios=(1.0 - float(val_ratio), float(val_ratio), 0.0), - seed=seed, - ) - if val_idx: - return val_idx - if float(val_ratio) == 0.0: - return [] - return [int(np.random.default_rng(seed).integers(0, total))] - - -def _loocv_fold_id_from_key(key: object, fold_unit: str) -> str: - if fold_unit == "sample": - return str(key) - sample_id, region_id = key - region_value = "__NONE__" if region_id is None else str(region_id) - return f"{sample_id}::{region_value}" - - -def discover_viable_loocv_fold_ids( - preprocessor: SpatialOmicsPreprocessor, fold_unit: str -) -> List[str]: - if fold_unit not in {"region", "sample"}: - raise ValueError("split.loocv.fold_unit must be one of ['region', 'sample'].") - - manifest_rows = preprocessor._load_manifest( - Path(preprocessor.precompute_config.raw_manifest_path) - ) - tables = preprocessor._load_tables(manifest_rows) - prepared = preprocessor._prepare_samples(manifest_rows, tables) - - fold_keys: Set[object] = set() - for sample in prepared: - if not sample.patch_indices: - continue - key = ( - (sample.sample_id, sample.region_id) - if fold_unit == "region" - else sample.sample_id - ) - fold_keys.add(key) - - return sorted( - _loocv_fold_id_from_key(key=key, fold_unit=fold_unit) for key in fold_keys - ) - - -def _select_graph_label_id( - scope: str, - sample_id: str, - region_id: Optional[str], - patch_id: str, -) -> str: - if scope == "patch": - return patch_id - if scope == "sample": - return sample_id - if region_id is None: - raise ValueError("region_id required for region-scope graph labels") - return region_id - - -def _build_patch_id(sample_id: str, region_id: Optional[str], patch_idx: int) -> str: - if region_id: - return f"{sample_id}_{region_id}_patch_{patch_idx}" - return f"{sample_id}_patch_{patch_idx}" diff --git a/src/grass_mil/data/components/samplers.py b/src/grass_mil/data/components/samplers.py deleted file mode 100644 index 72fb2b5..0000000 --- a/src/grass_mil/data/components/samplers.py +++ /dev/null @@ -1,1060 +0,0 @@ -from __future__ import annotations - -import copy -from dataclasses import dataclass, field -from typing import Any, Callable, Dict, Iterable, Optional - -import torch -from torch.utils.data import DataLoader as TorchDataLoader -from torch.utils.data import Dataset - -from grass_mil.data.components.ego_radius import ( - EgoRadiusConfig, - EgoRadiusCutoff, - resolve_ego_radius, -) -from grass_mil.data.components.seed_sampling import ( - InteriorSeedConfig, - InteriorSeedIndexCache, - cache_key_for_unit, - compute_interior_seed_indices, -) - -try: - from torch_geometric.data import Batch, Data - - from grass_mil.batching import apply_per_subgraph - from torch_geometric.loader import DataLoader as PyGDataLoader - from torch_geometric.loader import ShaDowKHopSampler -except Exception as exc: # pragma: no cover - raise ImportError("torch_geometric is required for sampler strategies") from exc - -try: - from torch_geometric.typing import WITH_TORCH_SPARSE, SparseTensor -except Exception: # pragma: no cover - WITH_TORCH_SPARSE = False - SparseTensor = None # type: ignore - - -@dataclass -class SamplerConfig: - name: str - kwargs: Dict[str, Any] - runtime: Dict[str, Any] - - @classmethod - def from_dict(cls, cfg: Optional[Dict[str, Any]]) -> "SamplerConfig": - if cfg is None: - return cls(name="identity", kwargs={}, runtime={}) - return cls( - name=cfg.get("name", "identity"), - kwargs=cfg.get("kwargs", {}), - runtime=cfg.get("runtime", {}), - ) - - -@dataclass -class RuntimeShadowConfig: - enabled: bool = False - depth: int = 2 - num_neighbors: int = 8 - subgraph_batch_size: int = 32 - replace: bool = False - shuffle_subgraphs: bool = True - node_idx: Optional[torch.Tensor] = None - proportional_root_sampling: bool = True - property_name: str = "cell_type" - weight_mode: str = "inverse" - min_weight: float = 1e-6 - subsample_fraction: float = 1.0 - subsample_seed: Optional[int] = None - # Absolute number of root cells drawn per unit per epoch (uniformly, without - # replacement). Overrides `subsample_fraction` when set; `null` keeps every - # candidate root. The preprint trains on 256 (fine-tuning) / 512 (frozen - # encoder) roots per region per step and validates on a fixed 1024. - subsample_count: Optional[int] = None - # Seed that makes the per-unit root subset identical on every pass (the - # generator is re-seeded from `seed` and the unit index). Use for validation - # so epoch-to-epoch metric differences reflect the model, not the sample. - fixed_per_unit_seed: Optional[int] = None - # One batch per unit holding all of its selected roots, so the MIL bag, the - # attention softmax and the loss all cover the whole region at once. - # `subgraph_batch_size` is then only an upper bound (chunks are concatenated). - unit_as_bag: bool = False - # Pretraining: sample roots from the union of all units rather than one unit - # at a time (the preprint's BGRL draws its 64-root batches this way). - merge_units: bool = False - interior_seeds: InteriorSeedConfig = field(default_factory=InteriorSeedConfig) - ego_radius: EgoRadiusConfig = field(default_factory=EgoRadiusConfig) - - @classmethod - def from_dict(cls, cfg: Optional[Dict[str, Any]]) -> "RuntimeShadowConfig": - cfg = cfg or {} - num_neighbors = int(cfg.get("num_neighbors", 8)) - if num_neighbors < 0: - # -1 = unlimited: every neighbour is kept at every hop (full k-hop ego-graph). - num_neighbors = UNLIMITED_NEIGHBORS - subsample_count = cfg.get("subsample_count") - if subsample_count is not None: - subsample_count = int(subsample_count) - if subsample_count <= 0: - raise ValueError("runtime.subsample_count must be a positive integer or null.") - fixed_per_unit_seed = cfg.get("fixed_per_unit_seed") - if fixed_per_unit_seed is not None: - fixed_per_unit_seed = int(fixed_per_unit_seed) - weight_mode = str(cfg.get("weight_mode", "inverse")).strip().lower() - allowed_weight_modes = {"inverse", "sqrt_inverse", "proportional"} - if weight_mode not in allowed_weight_modes: - raise ValueError( - f"Unsupported weight_mode '{weight_mode}'. " - f"Expected one of {sorted(allowed_weight_modes)}." - ) - min_weight = float(cfg.get("min_weight", 1e-6)) - if min_weight <= 0: - min_weight = 1e-6 - subsample_fraction = float(cfg.get("subsample_fraction", 1.0)) - if subsample_fraction == 0.0: - raise ValueError( - "runtime.subsample_fraction=0.0 selects no roots and is invalid." - ) - if not (0.0 < subsample_fraction <= 1.0): - raise ValueError( - "runtime.subsample_fraction must be within (0, 1]. " - f"Got {subsample_fraction}." - ) - subsample_seed = cfg.get("subsample_seed") - if subsample_seed is not None: - subsample_seed = int(subsample_seed) - return cls( - enabled=bool(cfg.get("enabled", False)), - depth=int(cfg.get("depth", 2)), - num_neighbors=num_neighbors, - subgraph_batch_size=int(cfg.get("subgraph_batch_size", 32)), - replace=bool(cfg.get("replace", False)), - shuffle_subgraphs=bool(cfg.get("shuffle_subgraphs", True)), - node_idx=cfg.get("node_idx"), - proportional_root_sampling=bool( - cfg.get("proportional_root_sampling", True) - ), - property_name=str(cfg.get("property_name", "cell_type")), - weight_mode=weight_mode, - min_weight=min_weight, - subsample_fraction=subsample_fraction, - subsample_seed=subsample_seed, - subsample_count=subsample_count, - fixed_per_unit_seed=fixed_per_unit_seed, - unit_as_bag=bool(cfg.get("unit_as_bag", False)), - merge_units=bool(cfg.get("merge_units", False)), - interior_seeds=InteriorSeedConfig.from_dict(cfg.get("interior_seeds")), - ego_radius=EgoRadiusConfig.from_dict(cfg.get("ego_radius")), - ) - - -UNLIMITED_NEIGHBORS = 10**8 - - -def _build_weighted_node_idx( - data: Data, - *, - property_name: str, - weight_mode: str, - min_weight: float, - allowed_indices: Optional[torch.Tensor] = None, -) -> Optional[torch.Tensor]: - if not hasattr(data, "categorical_codes") or not hasattr( - data, "categorical_slices" - ): - return None - categorical_codes = getattr(data, "categorical_codes") - categorical_slices = getattr(data, "categorical_slices") - if not isinstance(categorical_slices, dict): - return None - if property_name not in categorical_slices: - return None - if ( - not isinstance(categorical_codes, torch.Tensor) - or categorical_codes.numel() == 0 - ): - return None - if categorical_codes.dim() != 2: - return None - - prop_col = int(categorical_slices[property_name]) - if prop_col < 0 or prop_col >= int(categorical_codes.size(1)): - return None - - node_labels = categorical_codes[:, prop_col].long().view(-1) - if node_labels.numel() == 0: - return None - - # Restrict the candidate pool *before* weighting, so class frequencies are - # computed over the eligible nodes rather than the whole graph. - pool = None - if allowed_indices is not None: - pool = allowed_indices.view(-1).long() - if pool.numel() == 0: - return None - node_labels = node_labels[pool] - - unique_labels, inverse, counts = torch.unique( - node_labels, sorted=False, return_inverse=True, return_counts=True - ) - if unique_labels.numel() == 0: - return None - - counts = counts.float() - if weight_mode == "inverse": - class_weights = 1.0 / counts - elif weight_mode == "sqrt_inverse": - class_weights = 1.0 / torch.sqrt(counts) - elif weight_mode == "proportional": - class_weights = counts / counts.sum() - else: - return None - - class_weights = torch.clamp(class_weights, min=float(min_weight)) - node_weights = class_weights[inverse] - weight_sum = node_weights.sum() - if not torch.isfinite(weight_sum) or float(weight_sum) <= 0.0: - return None - probs = node_weights / weight_sum - if not torch.isfinite(probs).all(): - return None - - num_samples = int(getattr(data, "num_nodes", node_labels.numel())) - if pool is not None: - num_samples = int(pool.numel()) - if num_samples <= 0: - num_samples = int(node_labels.numel()) - if num_samples <= 0: - return None - sampled = torch.multinomial(probs, num_samples=num_samples, replacement=True).long() - if pool is not None: - # `sampled` indexes into the restricted pool; map back to node ids. - return pool[sampled] - return sampled - - -def _resolve_interior_indices( - data: Data, - runtime: "RuntimeShadowConfig", - cache: Optional[InteriorSeedIndexCache] = None, -) -> Optional[torch.Tensor]: - """Interior seed pool for one graph unit, or None when unrestricted.""" - if not runtime.interior_seeds.enabled: - return None - hops = ( - int(runtime.interior_seeds.n_hops) - if runtime.interior_seeds.n_hops is not None - else int(runtime.depth) - ) - key = cache_key_for_unit(data, n_hops=hops) - if cache is not None and key: - cached = cache.load(key) - if cached is not None: - return torch.from_numpy(cached).long() - indices = compute_interior_seed_indices(data, runtime.interior_seeds, n_hops=hops) - if cache is not None and key: - cache.store(key, indices) - return torch.from_numpy(indices).long() - - -def _compose_ego_radius_transform( - transform: Optional[Callable[[Data], Data]], - runtime: "RuntimeShadowConfig", -) -> Optional[Callable[[Data], Data]]: - """Append the ego-radius cutoff after any user transform.""" - if not runtime.ego_radius.enabled: - return transform - radius = resolve_ego_radius( - int(runtime.depth), - radius_per_hop=runtime.ego_radius.radius_per_hop, - radius_offset=runtime.ego_radius.radius_offset, - ) - cutoff = EgoRadiusCutoff( - radius=radius, - keep_root_component=runtime.ego_radius.keep_root_component, - ) - if transform is None: - return cutoff - - def _composed(data: Data) -> Data: - return cutoff(transform(data)) - - return _composed - - -def _resolve_root_candidates(data: Data, node_idx: Optional[torch.Tensor]) -> torch.Tensor: - if node_idx is None: - num_nodes = int(getattr(data, "num_nodes", 0)) - if num_nodes <= 0: - return torch.empty(0, dtype=torch.long) - return torch.arange(num_nodes, dtype=torch.long) - if isinstance(node_idx, torch.Tensor): - if node_idx.dtype == torch.bool: - return node_idx.nonzero(as_tuple=False).view(-1).long() - return node_idx.view(-1).long() - if hasattr(node_idx, "__len__"): - return torch.as_tensor(list(node_idx), dtype=torch.long) - return torch.empty(0, dtype=torch.long) - - -def _subsample_count(num_candidates: int, subsample_fraction: float) -> int: - if num_candidates <= 0: - return 0 - if subsample_fraction >= 1.0: - return num_candidates - n_keep = int(num_candidates * subsample_fraction) - if n_keep <= 0: - n_keep = 1 - return min(n_keep, num_candidates) - - -def _resolve_num_selected_roots( - num_candidates: int, *, subsample_fraction: float, subsample_count: Optional[int] -) -> int: - if subsample_count is not None: - return min(int(subsample_count), int(num_candidates)) - return _subsample_count(num_candidates, subsample_fraction) - - -def _subsample_root_candidates( - candidates: torch.Tensor, - *, - subsample_fraction: float, - generator: Optional[torch.Generator], - subsample_count: Optional[int] = None, -) -> torch.Tensor: - n_candidates = int(candidates.numel()) - n_keep = _resolve_num_selected_roots( - n_candidates, subsample_fraction=subsample_fraction, subsample_count=subsample_count - ) - if n_keep >= n_candidates: - return candidates - perm = torch.randperm(n_candidates, generator=generator) - selected = perm[:n_keep].sort().values - return candidates.index_select(0, selected) - - -def _per_unit_generator(seed: Optional[int], unit_index: int) -> Optional[torch.Generator]: - """Generator that yields the same root subset for a unit on every pass.""" - if seed is None: - return None - gen = torch.Generator() - gen.manual_seed(int(seed) * 1000003 + int(unit_index)) - return gen - - -class BaseSamplerStrategy: - """Data-layer agnostic sampler strategy. - - A strategy consumes graph units emitted by the data layer and returns iterables/loaders. - """ - - def build_dataset_loader( - self, - dataset: Dataset, - batch_size: int, - num_workers: int = 0, - pin_memory: bool = False, - shuffle: bool = True, - **kwargs: Any, - ): - raise NotImplementedError - - def build_unit_loader(self, data: Data, **kwargs: Any): - raise NotImplementedError - - -class _RuntimeUnitShadowDatasetLoader: - def __init__( - self, - *, - dataset: Dataset, - strategy: BaseSamplerStrategy, - batch_size: int, - runtime: RuntimeShadowConfig, - num_workers: int = 0, - pin_memory: bool = False, - shuffle: bool = True, - **kwargs: Any, - ): - self.dataset = dataset - self.strategy = strategy - self.dataset_batch_size = batch_size - self.runtime = runtime - self.num_workers = num_workers - self.pin_memory = pin_memory - self.shuffle = shuffle - self.kwargs = kwargs - self._interior_cache = ( - InteriorSeedIndexCache(runtime.interior_seeds.cache_dir) - if runtime.interior_seeds.enabled - else None - ) - - def __len__(self) -> int: - if len(self.dataset) == 0: - return 0 - subgraph_batch_size = int(self.runtime.subgraph_batch_size) - if subgraph_batch_size <= 0: - subgraph_batch_size = max(1, int(self.dataset_batch_size)) - - if self.runtime.merge_units: - total_roots = sum( - self._num_selected_roots_for_unit(self.dataset[idx]) - for idx in range(len(self.dataset)) - ) - return int((total_roots + subgraph_batch_size - 1) // subgraph_batch_size) - - total_batches = 0 - for idx in range(len(self.dataset)): - unit_data = self.dataset[idx] - effective_node_idx = self.runtime.node_idx - interior_idx = _resolve_interior_indices( - unit_data, self.runtime, self._interior_cache - ) - if interior_idx is not None: - effective_node_idx = interior_idx - if self.runtime.proportional_root_sampling: - computed_node_idx = _build_weighted_node_idx( - unit_data, - property_name=self.runtime.property_name, - weight_mode=self.runtime.weight_mode, - min_weight=self.runtime.min_weight, - allowed_indices=interior_idx, - ) - if computed_node_idx is not None: - effective_node_idx = computed_node_idx - num_candidate_roots = _resolve_num_roots(unit_data, effective_node_idx) - num_roots = _resolve_num_selected_roots( - num_candidate_roots, - subsample_fraction=float(self.runtime.subsample_fraction), - subsample_count=self.runtime.subsample_count, - ) - if self.runtime.unit_as_bag: - total_batches += 1 if num_roots > 0 else 0 - else: - total_batches += int( - (num_roots + subgraph_batch_size - 1) // subgraph_batch_size - ) - return total_batches - - def _num_selected_roots_for_unit(self, unit_data: Data) -> int: - effective_node_idx = self.runtime.node_idx - interior_idx = _resolve_interior_indices(unit_data, self.runtime, self._interior_cache) - if interior_idx is not None: - effective_node_idx = interior_idx - num_candidate_roots = _resolve_num_roots(unit_data, effective_node_idx) - return _resolve_num_selected_roots( - num_candidate_roots, - subsample_fraction=float(self.runtime.subsample_fraction), - subsample_count=self.runtime.subsample_count, - ) - - def _iter_merged(self, subgraph_batch_size: int): - """One pass over the roots of ALL units pooled into a single graph. - - The preprint's BGRL draws every 64-root batch from the merged - all-region graph, so a batch mixes ego-graphs of different regions. - Root selection per unit follows the same rules as the per-unit path - (interior seeds, subsample fraction/count, fixed seeds); weighted - (cell-type proportional) root sampling is not applied in merged mode. - """ - from grass_mil.data.components.merge import merge_graph_units - - units = [self.dataset[idx] for idx in range(len(self.dataset))] - merged = merge_graph_units(units) - roots: list[torch.Tensor] = [] - subsample_generator = None - if self.runtime.subsample_seed is not None: - subsample_generator = torch.Generator() - subsample_generator.manual_seed(int(self.runtime.subsample_seed)) - offsets = merged.unit_ptr - for idx, unit_data in enumerate(units): - effective_node_idx = self.runtime.node_idx - interior_idx = _resolve_interior_indices( - unit_data, self.runtime, self._interior_cache - ) - if interior_idx is not None: - effective_node_idx = interior_idx - candidates = _resolve_root_candidates(unit_data, effective_node_idx) - unit_generator = _per_unit_generator(self.runtime.fixed_per_unit_seed, idx) - candidates = _subsample_root_candidates( - candidates, - subsample_fraction=float(self.runtime.subsample_fraction), - generator=unit_generator if unit_generator is not None else subsample_generator, - subsample_count=self.runtime.subsample_count, - ) - if candidates.numel() == 0: - continue - roots.append(candidates.long() + int(offsets[idx])) - if not roots: - return - node_idx = torch.cat(roots) - if self.shuffle: - node_idx = node_idx[torch.randperm(node_idx.numel())] - unit_loader_kwargs: Dict[str, Any] = { - "num_workers": self.num_workers, - "pin_memory": self.pin_memory, - "persistent_workers": self.kwargs.get("persistent_workers", False) - if self.num_workers > 0 - else False, - } - if self.num_workers > 0 and "prefetch_factor" in self.kwargs: - unit_loader_kwargs["prefetch_factor"] = self.kwargs["prefetch_factor"] - loader = self.strategy.build_unit_loader( - data=merged, - depth=int(self.runtime.depth), - num_neighbors=int(self.runtime.num_neighbors), - batch_size=subgraph_batch_size, - node_idx=node_idx, - replace=bool(self.runtime.replace), - shuffle=False, - transform=_compose_ego_radius_transform(self.kwargs.get("transform"), self.runtime), - **unit_loader_kwargs, - ) - for sub_batch in loader: - yield _normalize_shadow_batch(sub_batch, merged) - - def __iter__(self): - if len(self.dataset) == 0: - return - - if self.runtime.merge_units: - subgraph_batch_size = int(self.runtime.subgraph_batch_size) - if subgraph_batch_size <= 0: - subgraph_batch_size = max(1, int(self.dataset_batch_size)) - yield from self._iter_merged(subgraph_batch_size) - return - - indices = torch.arange(len(self.dataset), dtype=torch.long) - if self.shuffle: - indices = indices[torch.randperm(len(indices))] - - subgraph_batch_size = int(self.runtime.subgraph_batch_size) - if subgraph_batch_size <= 0: - subgraph_batch_size = max(1, int(self.dataset_batch_size)) - subsample_generator = None - if self.runtime.subsample_seed is not None: - subsample_generator = torch.Generator() - subsample_generator.manual_seed(int(self.runtime.subsample_seed)) - - for idx in indices.tolist(): - unit_data = self.dataset[idx] - effective_node_idx = self.runtime.node_idx - interior_idx = _resolve_interior_indices( - unit_data, self.runtime, self._interior_cache - ) - if interior_idx is not None: - effective_node_idx = interior_idx - if self.runtime.proportional_root_sampling: - computed_node_idx = _build_weighted_node_idx( - unit_data, - property_name=self.runtime.property_name, - weight_mode=self.runtime.weight_mode, - min_weight=self.runtime.min_weight, - allowed_indices=interior_idx, - ) - if computed_node_idx is not None: - effective_node_idx = computed_node_idx - effective_root_candidates = _resolve_root_candidates( - unit_data, effective_node_idx - ) - unit_generator = _per_unit_generator(self.runtime.fixed_per_unit_seed, idx) - effective_root_candidates = _subsample_root_candidates( - effective_root_candidates, - subsample_fraction=float(self.runtime.subsample_fraction), - generator=unit_generator if unit_generator is not None else subsample_generator, - subsample_count=self.runtime.subsample_count, - ) - if effective_root_candidates.numel() == 0: - continue - unit_batch_size = subgraph_batch_size - if self.runtime.unit_as_bag: - unit_batch_size = int(effective_root_candidates.numel()) - unit_loader_kwargs: Dict[str, Any] = { - "num_workers": self.num_workers, - "pin_memory": self.pin_memory, - "persistent_workers": self.kwargs.get("persistent_workers", False) - if self.num_workers > 0 - else False, - } - if self.num_workers > 0 and "prefetch_factor" in self.kwargs: - unit_loader_kwargs["prefetch_factor"] = self.kwargs["prefetch_factor"] - loader = self.strategy.build_unit_loader( - data=unit_data, - depth=int(self.runtime.depth), - num_neighbors=int(self.runtime.num_neighbors), - batch_size=unit_batch_size, - node_idx=effective_root_candidates, - replace=bool(self.runtime.replace), - shuffle=bool(self.runtime.shuffle_subgraphs and self.shuffle), - transform=_compose_ego_radius_transform( - self.kwargs.get("transform"), self.runtime - ), - **unit_loader_kwargs, - ) - for sub_batch in loader: - yield _normalize_shadow_batch(sub_batch, unit_data) - - -class IdentityBatchStrategy(BaseSamplerStrategy): - def build_dataset_loader( - self, - dataset: Dataset, - batch_size: int, - num_workers: int = 0, - pin_memory: bool = False, - shuffle: bool = True, - **kwargs: Any, - ): - loader_cls = PyGDataLoader or TorchDataLoader - loader_kwargs = { - "dataset": dataset, - "batch_size": batch_size, - "num_workers": num_workers, - "pin_memory": pin_memory, - "shuffle": shuffle, - "persistent_workers": kwargs.get("persistent_workers", False) - if num_workers > 0 - else False, - } - if num_workers > 0 and "prefetch_factor" in kwargs: - loader_kwargs["prefetch_factor"] = kwargs["prefetch_factor"] - return loader_cls( - **loader_kwargs, - ) - - def build_unit_loader(self, data: Data, **kwargs: Any): - # Identity strategy on a single unit is just one-element iteration. - return [data] - - -class ShadowNativeStrategy(BaseSamplerStrategy): - """Native PyG ShaDow sampler, to be used per graph unit.""" - - def __init__(self, runtime: Optional[Dict[str, Any]] = None, **kwargs: Any): - self.runtime = RuntimeShadowConfig.from_dict(runtime) - self.kwargs = kwargs - - def build_dataset_loader( - self, - dataset: Dataset, - batch_size: int, - num_workers: int = 0, - pin_memory: bool = False, - shuffle: bool = True, - **kwargs: Any, - ): - if self.runtime.enabled: - return _RuntimeUnitShadowDatasetLoader( - dataset=dataset, - strategy=self, - batch_size=batch_size, - runtime=self.runtime, - num_workers=num_workers, - pin_memory=pin_memory, - shuffle=shuffle, - **kwargs, - ) - identity = IdentityBatchStrategy() - return identity.build_dataset_loader( - dataset=dataset, - batch_size=batch_size, - num_workers=num_workers, - pin_memory=pin_memory, - shuffle=shuffle, - **kwargs, - ) - - def build_unit_loader( - self, - data: Data, - depth: int, - num_neighbors: int, - batch_size: int, - node_idx: Optional[torch.Tensor] = None, - replace: bool = False, - shuffle: bool = True, - transform: Optional[Callable[[Data], Data]] = None, - **kwargs: Any, - ): - if transform is not None: - # Native loader does not guarantee per-subgraph post-collate transform behavior. - # Caller should choose ShadowCustomStrategy when strict semantics are required. - raise ValueError( - "ShadowNativeStrategy does not support strict post-hoc per-subgraph transforms. " - "Use ShadowCustomStrategy instead." - ) - num_workers = kwargs.get("num_workers", 0) - loader_kwargs: Dict[str, Any] = { - "data": data, - "depth": depth, - "num_neighbors": num_neighbors, - "node_idx": node_idx, - "replace": replace, - "batch_size": batch_size, - "shuffle": shuffle, - "num_workers": num_workers, - "pin_memory": kwargs.get("pin_memory", False), - "persistent_workers": kwargs.get("persistent_workers", False) - if num_workers > 0 - else False, - } - if num_workers > 0 and "prefetch_factor" in kwargs: - loader_kwargs["prefetch_factor"] = kwargs["prefetch_factor"] - return ShaDowKHopSampler(**loader_kwargs) - - -class ShadowCustomStrategy(BaseSamplerStrategy): - """Custom ShaDow sampler with optional per-subgraph transform support.""" - - def __init__(self, runtime: Optional[Dict[str, Any]] = None, **kwargs: Any): - self.runtime = RuntimeShadowConfig.from_dict(runtime) - self.kwargs = kwargs - - def build_dataset_loader( - self, - dataset: Dataset, - batch_size: int, - num_workers: int = 0, - pin_memory: bool = False, - shuffle: bool = True, - **kwargs: Any, - ): - if self.runtime.enabled: - return _RuntimeUnitShadowDatasetLoader( - dataset=dataset, - strategy=self, - batch_size=batch_size, - runtime=self.runtime, - num_workers=num_workers, - pin_memory=pin_memory, - shuffle=shuffle, - **kwargs, - ) - identity = IdentityBatchStrategy() - return identity.build_dataset_loader( - dataset=dataset, - batch_size=batch_size, - num_workers=num_workers, - pin_memory=pin_memory, - shuffle=shuffle, - **kwargs, - ) - - def build_unit_loader( - self, - data: Data, - depth: int, - num_neighbors: int, - batch_size: int, - node_idx: Optional[torch.Tensor] = None, - replace: bool = False, - shuffle: bool = True, - transform: Optional[Callable[[Data], Data]] = None, - **kwargs: Any, - ): - if not WITH_TORCH_SPARSE: - if transform is not None: - raise ValueError( - "ShadowCustomStrategy transform support requires torch-sparse. " - "Install torch-sparse or disable transforms." - ) - native = ShadowNativeStrategy(runtime={"enabled": False}) - return native.build_unit_loader( - data=data, - depth=depth, - num_neighbors=num_neighbors, - batch_size=batch_size, - node_idx=node_idx, - replace=replace, - shuffle=shuffle, - transform=None, - **kwargs, - ) - num_workers = kwargs.get("num_workers", 0) - loader_kwargs: Dict[str, Any] = { - "data": data, - "depth": depth, - "num_neighbors": num_neighbors, - "node_idx": node_idx, - "replace": replace, - "batch_size": batch_size, - "shuffle": shuffle, - "transform": transform, - "num_workers": num_workers, - "pin_memory": kwargs.get("pin_memory", False), - "persistent_workers": kwargs.get("persistent_workers", False) - if num_workers > 0 - else False, - } - if num_workers > 0 and "prefetch_factor" in kwargs: - loader_kwargs["prefetch_factor"] = kwargs["prefetch_factor"] - return _ShaDowKHopSamplerWithTransform(**loader_kwargs) - - -def _expand_graph_level_value(value: Any, count: int) -> Any: - if count <= 0: - return value - if isinstance(value, torch.Tensor): - if value.dim() == 0: - return value.repeat(count) - if value.size(0) == count: - return value - if value.size(0) == 1: - repeats = [count] + [1] * (value.dim() - 1) - return value.repeat(*repeats) - return value - if isinstance(value, (str, bytes)): - return [value] * count - if isinstance(value, (list, tuple)): - if len(value) == count: - return list(value) - if len(value) == 1: - return [value[0]] * count - return list(value) - return [value] * count - - -def _ensure_edge_index(batch: Batch) -> Batch: - if hasattr(batch, "edge_index") and batch.edge_index is not None: - return batch - if not hasattr(batch, "adj_t"): - return batch - row, col, edge_val = batch.adj_t.t().coo() - batch.edge_index = torch.stack([row, col], dim=0) - if edge_val is not None and not hasattr(batch, "edge_attr"): - batch.edge_attr = edge_val - return batch - - -def _resolve_batch_root_node_ids(batch: Batch) -> Optional[torch.Tensor]: - root_n_id = getattr(batch, "root_n_id", None) - if not isinstance(root_n_id, torch.Tensor): - return None - roots = root_n_id.detach().cpu().long().view(-1) - if roots.numel() == 0: - return roots - - roots_abs = roots - if hasattr(batch, "ptr") and isinstance(batch.ptr, torch.Tensor): - ptr = batch.ptr.detach().cpu().long().view(-1) - if ptr.numel() == roots.numel() + 1: - counts = ptr[1:] - ptr[:-1] - if bool(torch.all((roots >= 0) & (roots < counts))): - roots_abs = ptr[:-1] + roots - elif bool(torch.all((roots >= ptr[:-1]) & (roots < ptr[1:]))): - roots_abs = roots - - n_id = getattr(batch, "n_id", None) - if isinstance(n_id, torch.Tensor): - n_id = n_id.detach().cpu().long().view(-1) - if n_id.numel() > 0 and bool(torch.all((roots_abs >= 0) & (roots_abs < n_id.numel()))): - return n_id.index_select(0, roots_abs) - return roots_abs - - -def _normalize_shadow_batch(batch: Batch, source_data: Optional[Data] = None) -> Batch: - batch = _ensure_edge_index(batch) - num_subgraphs = int(len(batch.ptr) - 1) if hasattr(batch, "ptr") else 1 - - for key in ["graph_y", "graph_w", "sample_id", "region_id", "patch_id"]: - if hasattr(batch, key): - setattr( - batch, - key, - _expand_graph_level_value(getattr(batch, key), num_subgraphs), - ) - elif source_data is not None and hasattr(source_data, key): - setattr( - batch, - key, - _expand_graph_level_value(getattr(source_data, key), num_subgraphs), - ) - # Idempotent: this normalisation runs more than once over the same batch, - # and _resolve_batch_root_node_ids maps root indices through `n_id`. Applied - # twice, it maps already-global ids through `n_id` a second time and returns - # unrelated nodes -- after which a root is no longer found in its own - # subgraph. Convert only when the ids are not already global. - already_global = getattr(batch, "root_n_id_is_global", None) - if already_global is not None: - flags = torch.as_tensor(already_global).view(-1) - if flags.numel() > 0 and bool(flags[0]): - return batch - - root_node_ids = _resolve_batch_root_node_ids(batch) - if isinstance(root_node_ids, torch.Tensor): - batch.root_n_id = root_node_ids - batch.root_n_id_is_global = torch.ones( - int(root_node_ids.numel()), dtype=torch.bool - ) - return batch - - -def _is_schema_name_key(key: str) -> bool: - return str(key).endswith("_names") - - -def _select_sequence_value_for_subgraph( - *, - key: str, - value: list[Any] | tuple[Any, ...], - subgraph_index: int, - num_subgraphs: int, -) -> list[Any] | tuple[Any, ...]: - if _is_schema_name_key(key): - if ( - len(value) == num_subgraphs - and len(value) > subgraph_index - and isinstance(value[subgraph_index], (list, tuple)) - ): - item = value[subgraph_index] - return list(item) if isinstance(value, list) else tuple(item) - return list(value) if isinstance(value, list) else tuple(value) - - if len(value) == num_subgraphs and len(value) > subgraph_index: - item = value[subgraph_index] - return [item] if isinstance(value, list) else (item,) - return value - - -_SAMPLER_REGISTRY: Dict[str, Callable[..., BaseSamplerStrategy]] = { - "identity": IdentityBatchStrategy, - "shadow_native": ShadowNativeStrategy, - "shadow_custom": ShadowCustomStrategy, -} - - -def register_sampler_strategy(name: str): - def _decorator(cls): - _SAMPLER_REGISTRY[name] = cls - return cls - - return _decorator - - -def get_sampler_strategy( - config: SamplerConfig | Dict[str, Any] | None, -) -> BaseSamplerStrategy: - cfg = ( - config if isinstance(config, SamplerConfig) else SamplerConfig.from_dict(config) - ) - name = cfg.name.lower() - if name not in _SAMPLER_REGISTRY: - raise ValueError( - f"Unknown sampler strategy '{cfg.name}'. Available: {sorted(_SAMPLER_REGISTRY)}" - ) - if name == "identity": - return _SAMPLER_REGISTRY[name](**cfg.kwargs) - return _SAMPLER_REGISTRY[name](runtime=cfg.runtime, **cfg.kwargs) - - -class _ShaDowKHopSamplerWithTransform(torch.utils.data.DataLoader): - def __init__( - self, - data: Data, - depth: int, - num_neighbors: int, - node_idx: Optional[torch.Tensor] = None, - replace: bool = False, - transform: Optional[Callable[[Data], Data]] = None, - **kwargs: Any, - ): - if not WITH_TORCH_SPARSE: - raise ImportError( - "Custom ShaDow sampler requires torch-sparse; install torch-sparse or use shadow_native." - ) - self.data = copy.copy(data) - self.depth = depth - self.num_neighbors = num_neighbors - self.replace = replace - self.transform = transform - - if data.edge_index is not None: - self.is_sparse_tensor = False - row, col = data.edge_index.cpu() - self.adj_t = SparseTensor( - row=row, - col=col, - value=torch.arange(col.size(0)), - sparse_sizes=(data.num_nodes, data.num_nodes), - ).t() - else: - self.is_sparse_tensor = True - self.adj_t = data.adj_t.cpu() - - if node_idx is None: - node_idx = torch.arange(self.adj_t.sparse_size(0)) - elif node_idx.dtype == torch.bool: - node_idx = node_idx.nonzero(as_tuple=False).view(-1) - self.node_idx = node_idx - - super().__init__(node_idx.tolist(), collate_fn=self._collate, **kwargs) - - def _collate(self, n_id: Iterable[int]): - n_id = torch.tensor(list(n_id)) - rowptr, col, value = self.adj_t.csr() - out = torch.ops.torch_sparse.ego_k_hop_sample_adj( - rowptr, col, n_id, self.depth, self.num_neighbors, self.replace - ) - rowptr, col, n_id, e_id, ptr, root_n_id = out - - adj_t = SparseTensor( - rowptr=rowptr, - col=col, - value=value[e_id] if value is not None else None, - sparse_sizes=(n_id.numel(), n_id.numel()), - is_sorted=True, - trust_data=True, - ) - - batch = Batch(batch=torch.ops.torch_sparse.ptr2ind(ptr, n_id.numel()), ptr=ptr) - batch.root_n_id = root_n_id - batch.n_id = n_id - - if self.is_sparse_tensor: - batch.adj_t = adj_t - else: - row, col, e_id = adj_t.t().coo() - batch.edge_index = torch.stack([row, col], dim=0) - - for key, val in self.data: - if key in ["edge_index", "adj_t", "num_nodes", "batch", "ptr"]: - continue - if ( - key == "y" - and isinstance(val, torch.Tensor) - and val.size(0) == self.data.num_nodes - ): - batch[key] = val[n_id][root_n_id] - elif isinstance(val, torch.Tensor) and val.size(0) == self.data.num_nodes: - batch[key] = val[n_id] - elif isinstance(val, torch.Tensor) and val.size(0) == self.data.num_edges: - batch[key] = val[e_id] - else: - batch[key] = val - - batch = _normalize_shadow_batch(batch, self.data) - if self.transform is not None: - batch = self._apply_transform_to_each_subgraph(batch) - return batch - - def _apply_transform_to_each_subgraph(self, batch: Batch) -> Batch: - # Splitting lives in grass_mil.batching so the inference path can reuse - # it: it needs the same per-subgraph Data objects to export spatial - # edges, and previously had no way to get them from a ShaDow batch. - return apply_per_subgraph(batch, self.transform) - - -def _resolve_num_roots(data: Data, node_idx: Optional[torch.Tensor]) -> int: - if node_idx is None: - return int(getattr(data, "num_nodes", 0)) - if isinstance(node_idx, torch.Tensor): - if node_idx.dtype == torch.bool: - return int(node_idx.sum().item()) - return int(node_idx.numel()) - if hasattr(node_idx, "__len__"): - return int(len(node_idx)) - return int(getattr(data, "num_nodes", 0)) diff --git a/src/grass_mil/data/components/seed_sampling.py b/src/grass_mil/data/components/seed_sampling.py deleted file mode 100644 index ec0cca7..0000000 --- a/src/grass_mil/data/components/seed_sampling.py +++ /dev/null @@ -1,277 +0,0 @@ -"""Interior-only seed sampling for ego-graph extraction. - -Ego-graphs rooted on cells at the edge of a tissue section are truncated: their -neighbourhood is cut off by the section boundary rather than by biology, so -they are not comparable to interior ego-graphs. Restricting seeds to the -interior removes that artefact. - -"Interior" here excludes only the ``n_hops`` graph-hop rim around the **convex -hull** of the cell coordinates -- the outer tissue perimeter. Internal holes and -low-cellularity gaps are deliberately *not* treated as boundary. - -Ported from ``boundary_sampling.convex_hull_interior_mask`` in the NSCLC tuning -study, keeping the mask definition identical so that models trained against the -legacy definition remain comparable. The graph traversal is expressed with -sparse matrix operations instead of a Python BFS; the result is the same. -""" - -from __future__ import annotations - -import os -from dataclasses import dataclass -from pathlib import Path -from typing import Optional - -import numpy as np - -__all__ = [ - "InteriorSeedConfig", - "InteriorSeedIndexCache", - "convex_hull_interior_mask", - "distance_to_hull_polygon", - "compute_interior_seed_indices", -] - - -@dataclass(frozen=True) -class InteriorSeedConfig: - """Configuration for interior seed restriction. - - Attributes: - enabled: Restrict seeds to the interior mask. - n_hops: Width of the excluded rim, in graph hops. ``None`` uses the - sampler depth, so the rim matches the ego-graph radius. - seed_tolerance_multiplier: Cells within this multiple of the median - nearest-neighbour distance of the hull polygon start the rim. - keep_largest_component: Discard interior islands disconnected from the - main interior body. - cache_dir: Directory for the on-disk index cache. ``None`` disables it. - """ - - enabled: bool = False - n_hops: Optional[int] = None - seed_tolerance_multiplier: float = 1.0 - keep_largest_component: bool = True - cache_dir: Optional[str] = None - - @classmethod - def from_dict(cls, cfg: Optional[dict]) -> "InteriorSeedConfig": - cfg = dict(cfg or {}) - allowed = { - "enabled", - "n_hops", - "seed_tolerance_multiplier", - "keep_largest_component", - "cache_dir", - } - unknown = set(cfg) - allowed - if unknown: - raise ValueError( - f"Unknown interior seed keys: {sorted(unknown)}. " - f"Expected a subset of {sorted(allowed)}." - ) - n_hops = cfg.get("n_hops") - return cls( - enabled=bool(cfg.get("enabled", False)), - n_hops=None if n_hops is None else int(n_hops), - seed_tolerance_multiplier=float(cfg.get("seed_tolerance_multiplier", 1.0)), - keep_largest_component=bool(cfg.get("keep_largest_component", True)), - cache_dir=cfg.get("cache_dir"), - ) - - -def distance_to_hull_polygon(coords: np.ndarray, hull_vertices: np.ndarray) -> np.ndarray: - """Distance from every point to the convex-hull polygon's edge segments.""" - coords = np.asarray(coords, dtype=float) - hull_vertices = np.asarray(hull_vertices) - starts = coords[hull_vertices] - ends = coords[np.roll(hull_vertices, -1)] - - distances = np.full(coords.shape[0], np.inf) - for start, end in zip(starts, ends): - segment = end - start - denom = float(segment @ segment) + 1e-12 - t = np.clip(((coords - start) @ segment) / denom, 0.0, 1.0) - projection = start + np.outer(t, segment) - distances = np.minimum(distances, np.linalg.norm(coords - projection, axis=1)) - return distances - - -def _adjacency(edge_index: np.ndarray, num_nodes: int): - from scipy.sparse import coo_matrix - - edge_index = np.asarray(edge_index) - if edge_index.size == 0: - return coo_matrix((num_nodes, num_nodes), dtype=bool).tocsr() - rows = np.concatenate([edge_index[0], edge_index[1]]).astype(np.int64) - cols = np.concatenate([edge_index[1], edge_index[0]]).astype(np.int64) - data = np.ones(rows.shape[0], dtype=bool) - return coo_matrix((data, (rows, cols)), shape=(num_nodes, num_nodes)).tocsr() - - -def convex_hull_interior_mask( - coords: np.ndarray, - edge_index: np.ndarray, - n_hops: int, - *, - seed_tolerance_multiplier: float = 1.0, - keep_largest_component: bool = True, -) -> np.ndarray: - """Boolean per-cell mask; ``True`` marks interior cells. - - Degenerate inputs (fewer than 4 points, collinear point sets, ``n_hops <= - 0``) fall back to marking every cell interior, so callers always get a - usable seed pool. - """ - coords = np.asarray(coords, dtype=float) - num_nodes = int(coords.shape[0]) - if n_hops <= 0 or num_nodes < 4: - return np.ones(num_nodes, dtype=bool) - - try: - from scipy.spatial import ConvexHull, cKDTree - except Exception: # pragma: no cover - scipy is a hard dependency - return np.ones(num_nodes, dtype=bool) - - try: - hull = np.asarray(ConvexHull(coords).vertices) - except Exception: - # Degenerate or collinear point set: no meaningful hull. - return np.ones(num_nodes, dtype=bool) - - nearest = float(np.median(cKDTree(coords).query(coords, k=2)[0][:, 1])) - edge_distance = distance_to_hull_polygon(coords, hull) - tolerance = seed_tolerance_multiplier * max(nearest, 1e-6) - - rim = edge_distance <= tolerance - if not rim.any(): - rim = np.zeros(num_nodes, dtype=bool) - rim[np.unique(hull)] = True - - adjacency = _adjacency(edge_index, num_nodes) - # Grey out everything within (n_hops - 1) further hops of the hull rim, - # i.e. graph distance 0..n_hops-1 from a rim cell. - reached = rim.copy() - frontier = rim.copy() - for _ in range(max(int(n_hops) - 1, 0)): - if not frontier.any(): - break - neighbours = adjacency.dot(frontier.astype(np.int8)).astype(bool) - frontier = neighbours & ~reached - reached |= frontier - - interior = ~reached - if keep_largest_component and interior.any(): - interior = _largest_component(adjacency, interior) - return interior - - -def _largest_component(adjacency, interior: np.ndarray) -> np.ndarray: - """Keep only the largest connected component of the interior subgraph.""" - from scipy.sparse.csgraph import connected_components - - indices = np.where(interior)[0] - if indices.size == 0: - return interior - sub = adjacency[indices][:, indices] - n_components, labels = connected_components(sub, directed=False) - if n_components <= 1: - return interior - counts = np.bincount(labels) - keep = labels == int(np.argmax(counts)) - mask = np.zeros_like(interior) - mask[indices[keep]] = True - return mask - - -def compute_interior_seed_indices( - data, - cfg: InteriorSeedConfig, - *, - n_hops: int, -) -> np.ndarray: - """Interior node indices for one graph unit. - - Falls back to every node when coordinates are unavailable or the mask is - empty, so seeding never silently yields an empty pool. - """ - pos = getattr(data, "pos", None) - num_nodes = int(getattr(data, "num_nodes", 0) or 0) - if pos is None: - return np.arange(num_nodes, dtype=np.int64) - - coords = pos.detach().cpu().numpy() if hasattr(pos, "detach") else np.asarray(pos) - coords = np.asarray(coords, dtype=float)[:, :2] - edge_index = getattr(data, "edge_index", None) - edge_index = ( - edge_index.detach().cpu().numpy() - if hasattr(edge_index, "detach") - else np.asarray(edge_index if edge_index is not None else np.empty((2, 0))) - ) - - hops = int(cfg.n_hops) if cfg.n_hops is not None else int(n_hops) - mask = convex_hull_interior_mask( - coords, - edge_index, - hops, - seed_tolerance_multiplier=cfg.seed_tolerance_multiplier, - keep_largest_component=cfg.keep_largest_component, - ) - indices = np.where(mask)[0].astype(np.int64) - if indices.size == 0: - return np.arange(coords.shape[0], dtype=np.int64) - return indices - - -class InteriorSeedIndexCache: - """On-disk cache of interior indices, keyed by graph identity and rim width. - - The hull mask is deterministic given (coordinates, edges, n_hops), and - recomputing it per epoch is wasteful for large regions. - """ - - def __init__(self, cache_dir: Optional[str]) -> None: - self.cache_dir = Path(cache_dir) if cache_dir else None - if self.cache_dir is not None: - self.cache_dir.mkdir(parents=True, exist_ok=True) - - def _path(self, key: str) -> Optional[Path]: - if self.cache_dir is None: - return None - safe = "".join(char if char.isalnum() or char in "._-" else "_" for char in key) - return self.cache_dir / f"{safe}.npy" - - def load(self, key: str) -> Optional[np.ndarray]: - path = self._path(key) - if path is None or not path.exists(): - return None - try: - return np.load(path) - except Exception: - # A truncated or corrupt cache entry must not be fatal; recompute. - return None - - def store(self, key: str, indices: np.ndarray) -> None: - path = self._path(key) - if path is None: - return - # Write to a process-unique temporary file and rename, so concurrent - # dataloader workers cannot observe a half-written array. - tmp = path.with_suffix(f".tmp{os.getpid()}.npy") - try: - np.save(tmp, indices) - os.replace(tmp, path) - except Exception: - if tmp.exists(): - tmp.unlink(missing_ok=True) - - -def cache_key_for_unit(data, *, n_hops: int) -> str: - """Stable identity for a graph unit's interior mask.""" - parts = [ - str(getattr(data, "sample_id", "") or ""), - str(getattr(data, "region_id", "") or ""), - str(getattr(data, "patch_id", "") or ""), - f"hops{int(n_hops)}", - ] - return "__".join(part for part in parts if part) diff --git a/src/grass_mil/data/components/spatial_types.py b/src/grass_mil/data/components/spatial_types.py deleted file mode 100644 index 724f48b..0000000 --- a/src/grass_mil/data/components/spatial_types.py +++ /dev/null @@ -1,19 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Dict, Optional, Sequence - -import numpy as np - - -@dataclass -class SpatialOmicsTable: - """Normalized in-memory representation of spatial omics inputs.""" - - coords: np.ndarray # (n_cells, 2) float, in micrometers - molecular_features: Optional[np.ndarray] # (n_cells, n_features) float - molecular_feature_names: Optional[Sequence[str]] - categorical_labels: Dict[str, np.ndarray] # {label_name: (n_cells,)} - cell_ids: np.ndarray # (n_cells,) - sample_id: Optional[str] = None - region_id: Optional[str] = None diff --git a/src/grass_mil/data/components/splits.py b/src/grass_mil/data/components/splits.py deleted file mode 100644 index 5459e53..0000000 --- a/src/grass_mil/data/components/splits.py +++ /dev/null @@ -1,197 +0,0 @@ -"""Explicit splits, fold files and patient-grouped inner folds. - -The preprint's cross-validation uses fixed, patient-grouped outer folds stored -as a JSON list of region-id lists, and patient-grouped stratified inner folds -built from the development regions of each outer fold. These helpers read and -build those splits; the datamodule applies them through ``data.split.explicit``. -""" - -from __future__ import annotations - -import json -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple - -import numpy as np -import pandas as pd - -__all__ = [ - "ExplicitSplit", - "LabelTable", - "load_fold_json", - "load_labels", - "make_inner_folds", - "write_fold_json", -] - - -@dataclass -class ExplicitSplit: - """Region (or sample) ids per split; ids not listed are excluded.""" - - train: List[str] = field(default_factory=list) - val: List[str] = field(default_factory=list) - test: List[str] = field(default_factory=list) - id_kind: str = "region" # region | sample | patch - - @classmethod - def from_dict(cls, cfg: Optional[Mapping[str, Any]]) -> Optional["ExplicitSplit"]: - if not cfg: - return None - train = cfg.get("train") - val = cfg.get("val") - test = cfg.get("test") - if train is None and val is None and test is None: - return None - id_kind = str(cfg.get("id_kind", "region")) - if id_kind not in {"region", "sample", "patch"}: - raise ValueError("split.explicit.id_kind must be 'region', 'sample' or 'patch'.") - - def _ids(value: Any) -> List[str]: - if value is None: - return [] - if isinstance(value, (str, Path)): - path = Path(str(value)) - if path.suffix.lower() == ".json": - return [str(v) for v in json.loads(path.read_text())] - return [line.strip() for line in path.read_text().splitlines() if line.strip()] - return [str(v) for v in value] - - split = cls(train=_ids(train), val=_ids(val), test=_ids(test), id_kind=id_kind) - overlap = (set(split.train) & set(split.val)) | (set(split.train) & set(split.test)) - overlap |= set(split.val) & set(split.test) - if overlap and not (set(split.val) == set(split.test) and not (set(split.train) & overlap)): - # val == test is tolerated (refit runs score the held-out fold as "val" - # without using it for selection); any other overlap is a leak. - raise ValueError(f"split.explicit lists overlap: {sorted(overlap)[:10]}") - return split - - def lookup(self) -> Dict[str, str]: - out: Dict[str, str] = {} - for name in ("train", "val", "test"): - for rid in getattr(self, name): - out.setdefault(rid, name) - return out - - -def load_fold_json(path: str | Path) -> List[List[str]]: - """Outer folds: a JSON list of lists of region ids.""" - folds = json.loads(Path(path).read_text()) - if not isinstance(folds, list) or not all(isinstance(f, list) for f in folds): - raise ValueError(f"{path} must contain a JSON list of lists of region ids.") - folds = [[str(r) for r in fold] for fold in folds] - seen: set[str] = set() - for fold in folds: - dup = seen & set(fold) - if dup: - raise ValueError(f"Region ids appear in more than one fold: {sorted(dup)[:10]}") - seen |= set(fold) - return folds - - -def write_fold_json(folds: Sequence[Sequence[str]], path: str | Path) -> None: - Path(path).write_text(json.dumps([[str(r) for r in f] for f in folds])) - - -@dataclass -class LabelTable: - """Per-region label and grouping information.""" - - labels: pd.Series # index: region id (str) -> int label - groups: pd.Series # index: region id (str) -> group id (patient), str - - def label_of(self, rid: str) -> int: - return int(self.labels.loc[str(rid)]) - - def group_of(self, rid: str) -> str: - return str(self.groups.loc[str(rid)]) - - def subset(self, ids: Sequence[str]) -> "LabelTable": - ids = [str(r) for r in ids] - return LabelTable(labels=self.labels.loc[ids], groups=self.groups.loc[ids]) - - -def load_labels( - label_file: str | Path, - *, - id_column: str, - label_column: str, - patient_column: Optional[str] = None, - drop_unlabelled: bool = True, -) -> LabelTable: - """Read a label CSV into region -> label and region -> group series. - - Without ``patient_column`` every region is its own group (no grouping). - Non-numeric labels are mapped to codes in sorted order. - """ - df = pd.read_csv(label_file) - df.columns = [str(c).strip() for c in df.columns] - for col in (id_column, label_column) + ((patient_column,) if patient_column else ()): - if col not in df.columns: - raise ValueError(f"Column '{col}' not found in {label_file}; columns: {list(df.columns)}") - df[id_column] = df[id_column].astype(str) - df = df.drop_duplicates(id_column).set_index(id_column) - values = df[label_column] - if drop_unlabelled: - values = values.dropna() - if not np.issubdtype(values.dtype, np.number): - codes = {v: i for i, v in enumerate(sorted(values.astype(str).unique()))} - values = values.astype(str).map(codes) - labels = values.astype(float).round().astype(int) - if patient_column: - groups = df.loc[labels.index, patient_column].astype(str) - else: - groups = pd.Series(labels.index, index=labels.index, dtype=str) - return LabelTable(labels=labels, groups=groups) - - -def make_inner_folds( - dev_ids: Sequence[str], - table: LabelTable, - *, - k: int = 3, - seed: int = 0, - max_reseeds: int = 20, -) -> List[Tuple[List[str], List[str]]]: - """Stratified, group-aware inner folds over the SORTED development ids. - - Uses ``StratifiedGroupKFold(k, shuffle=True, random_state=seed)``; sorting the - ids first makes the folds a deterministic function of the id set, which the - preprint relies on for reproducible epoch selection. With very few groups - sklearn's greedy assignment can leave a fold without one class, which makes - AUROC undefined; such a split is rejected and the next seed is tried - (``max_reseeds`` times) before giving up. Large cohorts keep ``seed``. - """ - from sklearn.model_selection import StratifiedGroupKFold - - ids = np.array(sorted(str(r) for r in dev_ids)) - y = np.array([table.label_of(r) for r in ids]) - g = np.array([table.group_of(r) for r in ids]) - n_classes = len(np.unique(y)) - if n_classes < 2: - raise ValueError("Inner folds need at least two classes among the development regions.") - if len(np.unique(g)) < int(k): - raise ValueError( - f"Inner folds need at least k={k} patient groups; got {len(np.unique(g))}." - ) - last_error: Optional[str] = None - for attempt in range(int(max_reseeds) + 1): - splitter = StratifiedGroupKFold( - n_splits=int(k), shuffle=True, random_state=int(seed) + attempt - ) - folds: List[Tuple[List[str], List[str]]] = [] - ok = True - for tr, va in splitter.split(np.zeros((len(ids), 1)), y, g): - if len(np.unique(y[tr])) < n_classes or len(np.unique(y[va])) < 2: - ok = False - break - folds.append((ids[tr].tolist(), ids[va].tolist())) - if ok and len(folds) == int(k): - return folds - last_error = f"seed {int(seed) + attempt} left a fold without one class" - raise ValueError( - f"Could not build {k} stratified patient-grouped inner folds with every class in " - f"every split after {max_reseeds + 1} seeds ({last_error}). Use fewer inner folds " - "or more development patients." - ) diff --git a/src/grass_mil/data/components/transforms.py b/src/grass_mil/data/components/transforms.py deleted file mode 100644 index b4faf55..0000000 --- a/src/grass_mil/data/components/transforms.py +++ /dev/null @@ -1,78 +0,0 @@ -from __future__ import annotations - -import importlib -from typing import Any, Callable, Dict, Iterable, List, Optional - -import numpy as np -import torch - - -def instantiate_transforms( - configs: Optional[Iterable[Dict[str, Any]]], -) -> List[Callable]: - if not configs: - return [] - transforms: List[Callable] = [] - for cfg in configs: - name = cfg.get("name") - kwargs = cfg.get("kwargs", {}) - if not name: - continue - module_name, class_name = name.rsplit(".", 1) - module = importlib.import_module(module_name) - cls = getattr(module, class_name) - transforms.append(cls(**kwargs)) - return transforms - - -class CompositionVector: - """Convert node labels into a composition vector (graph-level).""" - - def __init__( - self, - label_attr: str = "label_cell_type", - num_classes: Optional[int] = None, - output_attr: str = "x", - normalize: bool = True, - keep_original: bool = False, - ) -> None: - self.label_attr = label_attr - self.num_classes = num_classes - self.output_attr = output_attr - self.normalize = normalize - self.keep_original = keep_original - - def __call__(self, data): - if hasattr(data, self.label_attr): - labels = getattr(data, self.label_attr) - labels_np = ( - labels.detach().cpu().numpy() - if isinstance(labels, torch.Tensor) - else np.asarray(labels) - ) - elif hasattr(data, "categorical_codes") and hasattr(data, "categorical_slices"): - label_name = self.label_attr.replace("label_", "") - if label_name not in data.categorical_slices: - raise AttributeError( - f"Missing label '{label_name}' in categorical_slices for composition vector" - ) - col_idx = data.categorical_slices[label_name] - labels_np = data.categorical_codes[:, col_idx].detach().cpu().numpy() - else: - raise AttributeError( - f"Missing label attribute '{self.label_attr}' or categorical_codes for composition vector" - ) - - labels_np = labels_np.reshape(-1) - num_classes = self.num_classes or ( - int(labels_np.max()) + 1 if labels_np.size > 0 else 0 - ) - counts = np.bincount(labels_np, minlength=num_classes).astype(np.float32) - if self.normalize and counts.sum() > 0: - counts = counts / counts.sum() - - vec = torch.from_numpy(counts).view(1, -1) - if self.keep_original and hasattr(data, "x"): - data.x_original = data.x - setattr(data, self.output_attr, vec) - return data diff --git a/src/grass_mil/data/datamodule.py b/src/grass_mil/data/datamodule.py new file mode 100644 index 0000000..20d4e8d --- /dev/null +++ b/src/grass_mil/data/datamodule.py @@ -0,0 +1,309 @@ +"""Lightning datamodule: processed graph units, labels, splits and ego-graph loaders.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Any, Mapping, Sequence + +import lightning as L +import torch +from torch.utils.data import Dataset + +from .graphs import GraphConfig +from .labels import LabelTable +from .loaders import InputConfig +from .precompute import BuildConfig, ProcessedIndex, UnitEntry, build_graphs +from .reducers import ReducerConfig +from .samplers import EgoGraphLoader, SamplerConfig +from .splits import random_split, read_ids +from .units import TilingConfig + +__all__ = ["SpatialOmicsDataModule", "GraphUnitDataset", "InputSpec"] + +log = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class InputSpec: + """What a model needs to know about the node and edge inputs.""" + + feature_dim: int + cell_types: tuple[str, ...] + edge_dim: int + + def as_dict(self) -> dict[str, Any]: + return {"feature_dim": self.feature_dim, "cell_types": list(self.cell_types), "edge_dim": self.edge_dim} + + @classmethod + def from_dict(cls, d: Mapping[str, Any]) -> InputSpec: + return cls(int(d["feature_dim"]), tuple(d["cell_types"]), int(d["edge_dim"])) + + +class GraphUnitDataset(Dataset): + """Graph units with their ids and, when labelled, a ``graph_y`` target row.""" + + def __init__(self, entries: Sequence[UnitEntry], labels: LabelTable | None = None) -> None: + self.entries = list(entries) + self.labels = labels + + def __len__(self) -> int: + return len(self.entries) + + def label_id(self, entry: UnitEntry) -> str | None: + if self.labels is None: + return None + return self.labels.resolve(entry.unit_id, entry.region_id, entry.sample_id) + + def load(self, idx: int) -> dict[str, Any]: + # The files are written by this package's own precompute step and hold + # PyG objects, which torch.load(weights_only=True) cannot unpickle. + return torch.load(self.entries[idx].path, weights_only=False) + + def __getitem__(self, idx: int): + entry = self.entries[idx] + data = self.load(idx)["data"] + data.unit_id = entry.unit_id + data.region_id = entry.region_id + data.sample_id = entry.sample_id + label_id = self.label_id(entry) + if label_id is not None: + data.graph_y = torch.from_numpy(self.labels.target(label_id)).view(1, -1) + return data + + +class SpatialOmicsDataModule(L.LightningDataModule): + """Cell graphs of a cohort, cut into units (regions, tiles or polygons). + + Splits are drawn when the data are set up, not when graphs are built, so + one set of processed graphs serves pretraining, every cross-validation fold + and prediction. With explicit ``split.train/val/test`` lists (unit, region + or sample ids, or files listing them) only the listed units are used; + otherwise units are split at random by ``split.ratios``, grouped by patient + (or region when the label file has no patient column). + """ + + def __init__( + self, + manifest_path: str, + processed_dir: str, + unit: str = "region", + coord_scale_um: float = 1.0, + min_cells: int = 10, + tiling: Mapping[str, Any] | None = None, + input: Mapping[str, Any] | None = None, + graph: Mapping[str, Any] | None = None, + feature_reducer: Mapping[str, Any] | None = None, + labels: Mapping[str, Any] | None = None, + labelled_only: bool = True, + split: Mapping[str, Any] | None = None, + sampler: Mapping[str, Any] | None = None, + val_roots_per_unit: int | None = 1024, + num_workers: int = 0, + pin_memory: bool = False, + force_precompute: bool = False, + ) -> None: + super().__init__() + self.save_hyperparameters(_plain(locals()), logger=False) + self.build_cfg = BuildConfig( + unit=str(unit), + coord_scale_um=float(coord_scale_um), + min_cells=int(min_cells), + tiling=TilingConfig.from_dict(tiling), + input=InputConfig.from_dict(input), + graph=GraphConfig.from_dict(graph), + feature_reducer=ReducerConfig.from_dict(feature_reducer), + ) + self.manifest_path = str(manifest_path) + self.processed_dir = str(processed_dir) + self.labels_cfg = dict(labels or {}) + self.labelled_only = bool(labelled_only) + self.split_cfg = dict(split or {}) + self.sampler_cfg = SamplerConfig.from_dict(sampler) + self.val_roots_per_unit = val_roots_per_unit + self.num_workers = int(num_workers) + self.pin_memory = bool(pin_memory) + self.force_precompute = bool(force_precompute) + unknown = set(self.split_cfg) - {"ratios", "seed", "train", "val", "test"} + if unknown: + raise ValueError(f"Unknown data.split keys: {sorted(unknown)}") + if self.build_cfg.feature_reducer.fit_on == "train" and self._explicit(): + raise ValueError( + "feature_reducer.fit_on=train cannot be combined with explicit splits (as used " + "by cross-validation): every fold would need its own graphs. Use fit_on=all." + ) + + self.index: ProcessedIndex | None = None + self.labels: LabelTable | None = None + self.dataset_train: GraphUnitDataset | None = None + self.dataset_val: GraphUnitDataset | None = None + self.dataset_test: GraphUnitDataset | None = None + + # ----------------------------------------------------------------- setup + def _explicit(self) -> bool: + return any(self.split_cfg.get(k) is not None for k in ("train", "val", "test")) + + def read_labels(self) -> LabelTable | None: + if self.labels is None and self.labels_cfg.get("file"): + columns = self.labels_cfg.get("columns") + self.labels = LabelTable.read( + self.labels_cfg["file"], + id_column=str(self.labels_cfg.get("id_column", "region_id")), + target_columns=list(columns) if columns else None, + patient_column=self.labels_cfg.get("patient_column"), + ) + return self.labels + + def prepare_data(self) -> None: + reducer = self.build_cfg.feature_reducer + train_fn = self._random_assignment if reducer.fit_on == "train" else None + split_key = repr(sorted(self.split_cfg.items())) if train_fn else "" + build_graphs( + self.manifest_path, + self.processed_dir, + self.build_cfg, + force=self.force_precompute, + training_units=(lambda units: {u for u, s in train_fn(units).items() if s == "train"}) + if train_fn + else None, + split_key=split_key, + ) + # A forced rebuild happens once per process, not once per setup(). + self.force_precompute = False + + def _group_of(self, entry: UnitEntry) -> str: + labels = self.read_labels() + label_id = labels.resolve(entry.unit_id, entry.region_id, entry.sample_id) if labels else None + if labels is not None and label_id is not None and self.labels_cfg.get("patient_column"): + return labels.patient(label_id) + return entry.region_id + + def _random_assignment(self, units: Sequence[UnitEntry]) -> dict[str, str]: + groups = {u.unit_id: self._group_of(u) for u in units} + # Stratify by the first target when the units carry discrete labels. + strata = None + labels = self.read_labels() + if labels is not None: + values = {} + for u in units: + label_id = labels.resolve(u.unit_id, u.region_id, u.sample_id) + if label_id is not None: + values[groups[u.unit_id]] = float(labels.targets.loc[label_id].iloc[0]) + if len(values) == len(set(groups.values())) and len(set(values.values())) <= 10: + strata = {g: str(v) for g, v in values.items()} + ratios = self.split_cfg.get("ratios", (0.8, 0.1, 0.1)) + return random_split(groups, ratios, int(self.split_cfg.get("seed", 0)), strata) + + def _members(self, units: Sequence[UnitEntry]) -> dict[str, list[UnitEntry]]: + if not self._explicit(): + assignment = self._random_assignment(units) + return {s: [u for u in units if assignment[u.unit_id] == s] for s in ("train", "val", "test")} + out: dict[str, list[UnitEntry]] = {} + for split in ("train", "val", "test"): + ids = read_ids(self.split_cfg.get(split)) or [] + wanted = set(ids) + out[split] = [u for u in units if {u.unit_id, u.region_id, u.sample_id} & wanted] + found = {i for u in out[split] for i in (u.unit_id, u.region_id, u.sample_id)} + missing = sorted(wanted - found) + if missing: + raise ValueError( + f"data.split.{split} names {len(missing)} id(s) with no processed unit " + f"(first: {missing[:5]})." + ) + return out + + def setup(self, stage: str | None = None) -> None: + self.index = ProcessedIndex.read(self.processed_dir) + labels = self.read_labels() + units = self.index.units + if self.labelled_only and labels is not None: + labelled = [u for u in units if labels.resolve(u.unit_id, u.region_id, u.sample_id)] + if len(labelled) < len(units): + log.warning(f"{len(units) - len(labelled)} unit(s) have no label and are skipped.") + units = labelled + members = self._members(units) + self.dataset_train = GraphUnitDataset(members["train"], labels) + self.dataset_val = GraphUnitDataset(members["val"], labels) + self.dataset_test = GraphUnitDataset(members["test"], labels) + + def all_units(self) -> GraphUnitDataset: + """Every processed unit (labelled or not), regardless of the split.""" + if self.index is None: + self.index = ProcessedIndex.read(self.processed_dir) + return GraphUnitDataset(self.index.units, self.read_labels()) + + def unit_table(self, dataset: GraphUnitDataset | None = None): + """Region, sample and (when labelled) patient id of every unit, indexed by unit id.""" + import pandas as pd + + dataset = dataset if dataset is not None else self.all_units() + rows = [] + for entry in dataset.entries: + row = {"unit_id": entry.unit_id, "region_id": entry.region_id, "sample_id": entry.sample_id} + label_id = dataset.label_id(entry) + if label_id is not None and self.labels_cfg.get("patient_column"): + row["patient_id"] = dataset.labels.patient(label_id) + rows.append(row) + return pd.DataFrame(rows).set_index("unit_id") + + @property + def input_spec(self) -> InputSpec: + if self.index is None: + self.index = ProcessedIndex.read(self.processed_dir) + return InputSpec( + feature_dim=len(self.index.feature_names), + cell_types=tuple(self.index.cell_types), + edge_dim=len(self.index.edge_features), + ) + + # ---------------------------------------------------------------- loaders + def loader( + self, + dataset: GraphUnitDataset, + *, + mode: str = "chunks", + roots_per_unit: int | None = None, + seed: int | None = None, + shuffle: bool = False, + ) -> EgoGraphLoader: + """An ego-graph loader over ``dataset`` with this datamodule's sampling settings.""" + return EgoGraphLoader( + dataset, + self.sampler_cfg, + mode=mode, + roots_per_unit=roots_per_unit, + seed=seed, + shuffle=shuffle, + num_workers=self.num_workers, + pin_memory=self.pin_memory, + ) + + def train_dataloader(self) -> EgoGraphLoader: + mode = "merged" if self.sampler_cfg.merge_units else "bags" + return self.loader( + self.dataset_train, mode=mode, roots_per_unit=self.sampler_cfg.roots_per_unit, shuffle=True + ) + + def val_dataloader(self) -> EgoGraphLoader: + # A fixed root subset per unit, so epoch-to-epoch differences reflect the model. + return self.loader(self.dataset_val, roots_per_unit=self.val_roots_per_unit, seed=0) + + def test_dataloader(self) -> EgoGraphLoader: + return self.loader(self.dataset_test) + + def predict_dataloader(self) -> EgoGraphLoader: + return self.loader(self.dataset_test) + + +def _plain(values: dict[str, Any]) -> dict[str, Any]: + """Init arguments as plain Python containers (checkpoint-safe hyperparameters).""" + from omegaconf import DictConfig, ListConfig, OmegaConf + + out = {} + for key, value in values.items(): + if key in ("self", "__class__"): + continue + if isinstance(value, (DictConfig, ListConfig)): + value = OmegaConf.to_container(value, resolve=True) + out[key] = value + return out diff --git a/src/grass_mil/data/ego_radius.py b/src/grass_mil/data/ego_radius.py new file mode 100644 index 0000000..e17c2ef --- /dev/null +++ b/src/grass_mil/data/ego_radius.py @@ -0,0 +1,95 @@ +"""Euclidean radius cap for k-hop ego-graphs. + +Hop depth alone does not bound the physical extent of an ego-graph in sparse +tissue. With the cap enabled, nodes farther than + +.. math:: r_k = \\mathrm{radius\\_per\\_hop} \\cdot k + \\mathrm{radius\\_offset} + +from the root are dropped (``k`` is the sampler depth), in coordinate units +after ``coord_scale_um``. Applied to every sampled ego-graph after sampling. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import Any, Mapping + +import numpy as np +import torch + +__all__ = ["EgoRadiusConfig", "EgoRadiusCutoff", "radius_for_depth"] + + +@dataclass(frozen=True) +class EgoRadiusConfig: + enabled: bool = False + radius_per_hop: float = 75.0 + radius_offset: float = 55.0 + keep_root_component: bool = True + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any] | None) -> EgoRadiusConfig: + cfg = dict(cfg or {}) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown ego_radius keys: {sorted(unknown)}") + return cls(**cfg) + + +def radius_for_depth(depth: int, cfg: EgoRadiusConfig) -> float: + return float(cfg.radius_per_hop) * int(depth) + float(cfg.radius_offset) + + +class EgoRadiusCutoff: + """Drop the nodes of an ego-graph farther than ``radius`` from its root. + + The root is always kept. With ``keep_root_component`` the nodes left + disconnected from the root by the cut are dropped too. + """ + + def __init__(self, *, radius: float, keep_root_component: bool = True) -> None: + self.radius = float(radius) + self.keep_root_component = bool(keep_root_component) + + def __repr__(self) -> str: + return f"EgoRadiusCutoff(radius={self.radius}, keep_root_component={self.keep_root_component})" + + def __call__(self, data, root: int): + """Return the cut ego-graph and the new index of its root.""" + from torch_geometric.utils import subgraph + + n = int(data.num_nodes) + dist = torch.linalg.norm(data.pos[:, :2] - data.pos[root, :2], dim=1) + keep = dist <= self.radius + keep[root] = True + if self.keep_root_component: + keep = self._root_component(data.edge_index, keep, root, n) + if bool(keep.all()): + return data, root + idx = torch.nonzero(keep).view(-1) + edge_index, edge_attr = subgraph( + idx, data.edge_index, getattr(data, "edge_attr", None), relabel_nodes=True, num_nodes=n + ) + out = data.clone() + for key, value in data: + if isinstance(value, torch.Tensor) and value.dim() > 0 and value.size(0) == n: + out[key] = value[idx] + out.edge_index = edge_index + if edge_attr is not None: + out.edge_attr = edge_attr + out.num_nodes = int(idx.numel()) + return out, int((idx < root).sum()) + + @staticmethod + def _root_component(edge_index: torch.Tensor, keep: torch.Tensor, root: int, n: int): + from scipy.sparse import coo_matrix + from scipy.sparse.csgraph import connected_components + + keep_np = keep.cpu().numpy() + src, dst = edge_index.cpu().numpy() + both = keep_np[src] & keep_np[dst] + adj = coo_matrix( + (np.ones(int(both.sum()), dtype=bool), (src[both], dst[both])), shape=(n, n) + ) + _, labels = connected_components(adj, directed=False) + return torch.from_numpy(keep_np & (labels == labels[root])) diff --git a/src/grass_mil/data/graphs.py b/src/grass_mil/data/graphs.py new file mode 100644 index 0000000..7e54e0f --- /dev/null +++ b/src/grass_mil/data/graphs.py @@ -0,0 +1,89 @@ +"""Cell graphs: Delaunay triangulation over cell centroids, with edge features.""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import Any, Mapping + +import numpy as np +import torch + +__all__ = ["GraphConfig", "delaunay_edges", "edge_features", "log_distance_feature"] + +EDGE_FEATURES = ("log_distance", "distance", "neighbor") + + +@dataclass(frozen=True) +class GraphConfig: + """Edge features of the cell graph. + + Attributes: + edge_features: Any of ``log_distance`` (bounded log edge length), + ``distance`` (raw length) and ``neighbor`` (1 when the length is at + most ``neighbor_cutoff``), in that column order. + neighbor_cutoff: Contact distance for ``neighbor``, in coordinate units + (after ``coord_scale_um``). + log_distance_bounds: ``(lower, upper)`` of the log length mapped to + ``[0, 1]``: ``clip((ln d - lower) / (upper - lower), 0, 1)``. The + default ``(2, 5)`` maps lengths of about 7 to 0 and 150 to 1. + """ + + edge_features: tuple[str, ...] = ("log_distance", "neighbor") + neighbor_cutoff: float = 20.0 + log_distance_bounds: tuple[float, float] = (2.0, 5.0) + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any] | None) -> GraphConfig: + cfg = dict(cfg or {}) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown data.graph keys: {sorted(unknown)}") + out = cls( + edge_features=tuple(cfg.get("edge_features", cls.edge_features)), + neighbor_cutoff=float(cfg.get("neighbor_cutoff", cls.neighbor_cutoff)), + log_distance_bounds=tuple(cfg.get("log_distance_bounds", cls.log_distance_bounds)), + ) + bad = [f for f in out.edge_features if f not in EDGE_FEATURES] + if bad: + raise ValueError(f"Unknown edge feature(s) {bad}; expected {list(EDGE_FEATURES)}.") + lower, upper = out.log_distance_bounds + if upper <= lower: + raise ValueError("data.graph.log_distance_bounds must be increasing.") + return out + + +def delaunay_edges(coords: np.ndarray) -> torch.Tensor: + """Undirected Delaunay edges as a ``[2, E]`` index with both directions.""" + n = int(coords.shape[0]) + if n < 2: + return torch.empty((2, 0), dtype=torch.long) + if n < 3: + pairs = np.array([[0, 1]]) + else: + from scipy.spatial import Delaunay + + simplices = Delaunay(coords).simplices + pairs = np.concatenate([simplices[:, [0, 1]], simplices[:, [1, 2]], simplices[:, [0, 2]]]) + pairs = np.sort(pairs, axis=1) + pairs = np.unique(pairs, axis=0) + both = np.concatenate([pairs, pairs[:, ::-1]]) + return torch.from_numpy(np.ascontiguousarray(both.T)).long() + + +def log_distance_feature(distances: np.ndarray, lower: float, upper: float) -> np.ndarray: + """Bounded log edge length ``clip((ln d - lower) / (upper - lower), 0, 1)``.""" + log_d = np.log(np.maximum(np.asarray(distances, dtype=np.float64), 1e-12)) + return np.clip((log_d - lower) / (upper - lower), 0.0, 1.0).astype(np.float32) + + +def edge_features(coords: np.ndarray, edge_index: torch.Tensor, cfg: GraphConfig) -> torch.Tensor: + """``[E, len(cfg.edge_features)]`` edge attributes in the configured order.""" + src, dst = edge_index.numpy() + lengths = np.linalg.norm(coords[src] - coords[dst], axis=1) + columns = { + "log_distance": lambda: log_distance_feature(lengths, *cfg.log_distance_bounds), + "distance": lambda: lengths.astype(np.float32), + "neighbor": lambda: (lengths <= cfg.neighbor_cutoff).astype(np.float32), + } + stacked = np.stack([columns[name]() for name in cfg.edge_features], axis=1) + return torch.from_numpy(stacked.reshape(len(lengths), len(cfg.edge_features))) diff --git a/src/grass_mil/data/interior.py b/src/grass_mil/data/interior.py new file mode 100644 index 0000000..e04f8a9 --- /dev/null +++ b/src/grass_mil/data/interior.py @@ -0,0 +1,152 @@ +"""Restricting ego-graph roots to the tissue interior. + +Ego-graphs rooted at the edge of a section are truncated by the boundary rather +than by biology, so they are not comparable to interior ones. "Interior" +excludes the ``n_hops`` graph-hop rim around the convex hull of the cells; holes +and low-cellularity gaps inside the tissue are not treated as boundary. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, fields +from pathlib import Path +from typing import Any, Mapping + +import numpy as np + +__all__ = ["InteriorConfig", "InteriorCache", "interior_mask", "interior_indices"] + + +@dataclass(frozen=True) +class InteriorConfig: + """Interior root restriction. + + Attributes: + enabled: Restrict roots to the interior. + n_hops: Width of the excluded rim in graph hops (``None``: the sampler + depth, so the rim matches the ego-graph radius). + hull_tolerance: Cells within this multiple of the median + nearest-neighbour distance of the hull edge start the rim. + keep_largest_component: Drop interior islands cut off from the main body. + cache_dir: Directory caching the masks (``None``: no cache). + """ + + enabled: bool = False + n_hops: int | None = None + hull_tolerance: float = 1.0 + keep_largest_component: bool = True + cache_dir: str | None = None + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any] | None) -> InteriorConfig: + cfg = dict(cfg or {}) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown interior_roots keys: {sorted(unknown)}") + return cls(**cfg) + + +def _hull_distance(coords: np.ndarray, hull: np.ndarray) -> np.ndarray: + """Distance from every point to the convex-hull polygon's edges.""" + starts = coords[hull] + ends = coords[np.roll(hull, -1)] + out = np.full(coords.shape[0], np.inf) + for start, end in zip(starts, ends): + segment = end - start + t = np.clip(((coords - start) @ segment) / (segment @ segment + 1e-12), 0.0, 1.0) + out = np.minimum(out, np.linalg.norm(coords - (start + np.outer(t, segment)), axis=1)) + return out + + +def _adjacency(edge_index: np.ndarray, n: int): + from scipy.sparse import coo_matrix + + rows = np.concatenate([edge_index[0], edge_index[1]]) + cols = np.concatenate([edge_index[1], edge_index[0]]) + return coo_matrix((np.ones(rows.size, dtype=bool), (rows, cols)), shape=(n, n)).tocsr() + + +def interior_mask( + coords: np.ndarray, + edge_index: np.ndarray, + n_hops: int, + *, + hull_tolerance: float = 1.0, + keep_largest_component: bool = True, +) -> np.ndarray: + """Boolean per-cell mask, ``True`` for interior cells. + + Degenerate inputs (fewer than 4 cells, collinear cells, ``n_hops <= 0``) + mark every cell interior. + """ + from scipy.sparse.csgraph import connected_components + from scipy.spatial import ConvexHull, QhullError, cKDTree + + n = int(coords.shape[0]) + if n_hops <= 0 or n < 4: + return np.ones(n, dtype=bool) + try: + hull = ConvexHull(coords).vertices + except QhullError: + return np.ones(n, dtype=bool) + + spacing = float(np.median(cKDTree(coords).query(coords, k=2)[0][:, 1])) + rim = _hull_distance(coords, hull) <= hull_tolerance * max(spacing, 1e-12) + adjacency = _adjacency(edge_index, n) + reached = rim.copy() + frontier = rim.copy() + for _ in range(int(n_hops) - 1): + frontier = adjacency.dot(frontier.astype(np.int8)).astype(bool) & ~reached + if not frontier.any(): + break + reached |= frontier + interior = ~reached + + if keep_largest_component and interior.any(): + idx = np.flatnonzero(interior) + n_comp, labels = connected_components(adjacency[idx][:, idx], directed=False) + if n_comp > 1: + interior = np.zeros(n, dtype=bool) + interior[idx[labels == np.argmax(np.bincount(labels))]] = True + return interior + + +def interior_indices(data, cfg: InteriorConfig, *, n_hops: int) -> np.ndarray: + """Interior node indices of a graph unit (every node when the mask is empty).""" + coords = data.pos.detach().cpu().numpy().astype(float)[:, :2] + mask = interior_mask( + coords, + data.edge_index.detach().cpu().numpy(), + int(cfg.n_hops) if cfg.n_hops is not None else int(n_hops), + hull_tolerance=cfg.hull_tolerance, + keep_largest_component=cfg.keep_largest_component, + ) + idx = np.flatnonzero(mask) + return idx if idx.size else np.arange(coords.shape[0]) + + +class InteriorCache: + """On-disk cache of interior indices, keyed by unit and every mask setting.""" + + def __init__(self, cache_dir: str | None) -> None: + self.cache_dir = Path(cache_dir) if cache_dir else None + + def _path(self, unit_id: str, cfg: InteriorConfig, n_hops: int) -> Path | None: + if self.cache_dir is None: + return None + key = f"{unit_id}__h{n_hops}_t{cfg.hull_tolerance}_c{int(cfg.keep_largest_component)}" + return self.cache_dir / ("".join(c if c.isalnum() or c in "._-" else "_" for c in key) + ".npy") + + def get(self, data, cfg: InteriorConfig, *, n_hops: int) -> np.ndarray: + path = self._path(str(data.unit_id), cfg, n_hops) + if path is not None and path.exists(): + return np.load(path) + idx = interior_indices(data, cfg, n_hops=n_hops) + if path is not None: + path.parent.mkdir(parents=True, exist_ok=True) + # Write-then-rename, so concurrent workers never read a partial file. + tmp = path.with_suffix(f".{os.getpid()}.tmp.npy") + np.save(tmp, idx) + os.replace(tmp, path) + return idx diff --git a/src/grass_mil/data/labels.py b/src/grass_mil/data/labels.py new file mode 100644 index 0000000..cc1e19f --- /dev/null +++ b/src/grass_mil/data/labels.py @@ -0,0 +1,109 @@ +"""Unit labels and patient grouping, read once from a label CSV. + +The label file has one row per labelled id (a region, a tile or a sample), the +target column(s) and optionally a patient column used for grouping. Every part +of the package (training, cross-validation, prediction, niche analysis) reads +labels through :class:`LabelTable`, so they are parsed and encoded in exactly +one way: numeric columns are kept as floats, text columns are encoded as +integer codes in sorted order. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Sequence + +import numpy as np +import pandas as pd + +__all__ = ["LabelTable", "class_index"] + + +@dataclass +class LabelTable: + """Targets and patient of every labelled id. + + ``targets`` is indexed by id (str) with one float column per target; + ``patients`` maps id to patient (every id is its own patient when the file + has no patient column). ``codes`` records the encoding of text columns. + """ + + targets: pd.DataFrame + patients: pd.Series + codes: dict[str, list[str]] = field(default_factory=dict) + + @classmethod + def read( + cls, + path: str | Path, + *, + id_column: str, + target_columns: Sequence[str] | None = None, + patient_column: str | None = None, + ) -> LabelTable: + df = pd.read_csv(path, dtype={id_column: str}) + df.columns = [str(c).strip() for c in df.columns] + if id_column not in df.columns: + raise ValueError(f"Label file {path} has no id column '{id_column}'.") + if patient_column and patient_column not in df.columns: + raise ValueError(f"Label file {path} has no patient column '{patient_column}'.") + df[id_column] = df[id_column].str.strip() + if df[id_column].duplicated().any(): + dup = sorted(df.loc[df[id_column].duplicated(), id_column])[:5] + raise ValueError(f"Label file {path} repeats ids: {dup}") + df = df.set_index(id_column) + if target_columns is None: + # Default: the first column that is not the patient column. + candidates = [c for c in df.columns if c != patient_column] + if not candidates: + raise ValueError(f"Label file {path} has no target column.") + target_columns = candidates[:1] + missing = [c for c in target_columns if c not in df.columns] + if missing: + raise ValueError(f"Label file {path} has no column(s) {missing}.") + + targets = pd.DataFrame(index=df.index) + codes: dict[str, list[str]] = {} + for column in target_columns: + values = df[column] + if not pd.api.types.is_numeric_dtype(values): + categories = sorted(values.dropna().astype(str).unique()) + codes[column] = categories + values = values.astype("string").map({c: i for i, c in enumerate(categories)}) + targets[column] = pd.to_numeric(values, errors="raise").astype(float) + targets = targets.dropna(how="any") + if patient_column: + patients = df.loc[targets.index, patient_column].astype(str) + else: + patients = pd.Series(targets.index, index=targets.index, dtype=str) + return cls(targets=targets, patients=patients, codes=codes) + + @property + def target_columns(self) -> list[str]: + return list(self.targets.columns) + + def __contains__(self, unit_id: str) -> bool: + return str(unit_id) in self.targets.index + + def resolve(self, *ids: str | None) -> str | None: + """The first of ``ids`` that has a label (unit id, then region, then sample).""" + for value in ids: + if value is not None and str(value) in self.targets.index: + return str(value) + return None + + def target(self, label_id: str) -> np.ndarray: + return self.targets.loc[str(label_id)].to_numpy(dtype=np.float32) + + def patient(self, label_id: str) -> str: + return str(self.patients.loc[str(label_id)]) + + def classes(self, label_id: str) -> int: + """Integer class of a classification label (first target column).""" + return class_index(float(self.targets.loc[str(label_id)].iloc[0])) + + +def class_index(value: float) -> int: + """Integer class of a label value: the nearest integer, halves rounded up (0.5 -> 1).""" + return int(np.floor(float(value) + 0.5)) diff --git a/src/grass_mil/data/loaders.py b/src/grass_mil/data/loaders.py new file mode 100644 index 0000000..1bc7753 --- /dev/null +++ b/src/grass_mil/data/loaders.py @@ -0,0 +1,243 @@ +"""Reading per-sample cell tables (CSV/TSV, h5ad, SingleCellExperiment). + +Every format is reduced to a :class:`CellTable`: coordinates, an optional +cell-type label per cell, optional continuous features and cell ids. Which +columns are read is described by one :class:`InputConfig` for all formats. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field, fields +from pathlib import Path +from typing import Any, Mapping, Sequence + +import numpy as np +import pandas as pd + +__all__ = ["CellTable", "InputConfig", "load_cell_table", "load_polygons", "input_type_of"] + + +@dataclass(frozen=True) +class InputConfig: + """Columns of the per-sample cell tables. + + Attributes: + coord_columns: x/y columns (CSV/TSV; h5ad ``obs`` and SCE ``colData`` + when no embedding key is set). + obsm_key: h5ad coordinates from ``obsm[obsm_key]``; ``None`` reads + ``coord_columns`` from ``obs``. + reduced_dim: SCE coordinates from ``reducedDim(sce, reduced_dim)``; + ``None`` reads ``coord_columns`` from ``colData``. + cell_id_column: Column of cell ids (``None``: row names / row number). + cell_type_column: Column of cell-type labels (``None``: no cell types). + features: Continuous node features: ``"all"`` (every remaining CSV + column, every h5ad variable / SCE feature), a list of names, or + ``None`` / ``[]`` for none (cell type only). + h5ad_layer: Read h5ad features from this layer instead of ``X``. + sce_assay: SCE assay name (``None``: the first assay). + sep: CSV separator (``.tsv`` files always use a tab). + """ + + coord_columns: tuple[str, str] = ("x", "y") + obsm_key: str | None = "spatial" + reduced_dim: str | None = None + cell_id_column: str | None = None + cell_type_column: str | None = "cell_type" + features: str | tuple[str, ...] | None = "all" + h5ad_layer: str | None = None + sce_assay: str | None = None + sep: str = "," + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any] | None) -> InputConfig: + cfg = dict(cfg or {}) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown data.input keys: {sorted(unknown)}") + if "coord_columns" in cfg: + cfg["coord_columns"] = tuple(cfg["coord_columns"]) + features = cfg.get("features", "all") + if features is not None and not isinstance(features, str): + cfg["features"] = tuple(str(f) for f in features) or None + elif isinstance(features, str) and features != "all": + raise ValueError("data.input.features must be 'all', a list of names or null.") + return cls(**cfg) + + +@dataclass +class CellTable: + coords: np.ndarray # (n, 2) float64 + cell_ids: np.ndarray # (n,) str + cell_types: np.ndarray | None = None # (n,) str + features: np.ndarray | None = None # (n, f) float32 + feature_names: list[str] = field(default_factory=list) + + def __len__(self) -> int: + return int(self.coords.shape[0]) + + def subset(self, indices: np.ndarray) -> CellTable: + return CellTable( + coords=self.coords[indices], + cell_ids=self.cell_ids[indices], + cell_types=None if self.cell_types is None else self.cell_types[indices], + features=None if self.features is None else self.features[indices], + feature_names=list(self.feature_names), + ) + + +_SUFFIX_TYPES = {".csv": "csv", ".tsv": "tsv", ".txt": "tsv", ".h5ad": "h5ad", ".rds": "sce"} + + +def input_type_of(path: str | Path, declared: str | None = None) -> str: + """``csv`` | ``tsv`` | ``h5ad`` | ``sce``, from the manifest or the file suffix.""" + if declared: + kind = str(declared).strip().lower() + kind = {"rds": "sce"}.get(kind, kind) + else: + kind = _SUFFIX_TYPES.get(Path(path).suffix.lower(), "") + if kind not in {"csv", "tsv", "h5ad", "sce"}: + raise ValueError(f"Cannot tell the input type of {path}; set input_type in the manifest.") + return kind + + +def load_cell_table(path: str | Path, kind: str, cfg: InputConfig) -> CellTable: + if kind in {"csv", "tsv"}: + return _load_csv(path, cfg, sep="\t" if kind == "tsv" else cfg.sep) + if kind == "h5ad": + return _load_h5ad(path, cfg) + if kind == "sce": + return _load_sce(path, cfg) + raise ValueError(f"Unsupported input type '{kind}'.") + + +def _select_features(names: Sequence[str], cfg: InputConfig, *, exclude: set[str]) -> list[str]: + if cfg.features is None: + return [] + if cfg.features == "all": + return [n for n in names if n not in exclude] + missing = [n for n in cfg.features if n not in set(names)] + if missing: + raise ValueError(f"Feature(s) not found: {missing[:10]}") + return list(cfg.features) + + +def _cell_types(values: pd.Series | None) -> np.ndarray | None: + return None if values is None else values.astype(str).to_numpy() + + +def _load_csv(path: str | Path, cfg: InputConfig, *, sep: str) -> CellTable: + df = pd.read_csv(path, sep=sep) + df.columns = [str(c).strip() for c in df.columns] + needed = [*cfg.coord_columns, *([cfg.cell_type_column] if cfg.cell_type_column else [])] + missing = [c for c in needed if c not in df.columns] + if missing: + raise ValueError(f"{path} has no column(s) {missing}") + exclude = {*cfg.coord_columns, cfg.cell_type_column, cfg.cell_id_column} + names = _select_features(list(df.columns), cfg, exclude=exclude) + return CellTable( + coords=df[list(cfg.coord_columns)].to_numpy(dtype=np.float64), + cell_ids=( + df[cfg.cell_id_column].astype(str).to_numpy() + if cfg.cell_id_column in df.columns + else np.arange(len(df)).astype(str) + ), + cell_types=_cell_types(df[cfg.cell_type_column] if cfg.cell_type_column else None), + features=df[names].to_numpy(dtype=np.float32) if names else None, + feature_names=names, + ) + + +def _load_h5ad(path: str | Path, cfg: InputConfig) -> CellTable: + import anndata as ad + + adata = ad.read_h5ad(path) + if cfg.obsm_key: + if cfg.obsm_key not in adata.obsm: + raise ValueError(f"{path} has no obsm['{cfg.obsm_key}']") + coords = np.asarray(adata.obsm[cfg.obsm_key], dtype=np.float64)[:, :2] + else: + coords = adata.obs[list(cfg.coord_columns)].to_numpy(dtype=np.float64) + if cfg.cell_type_column and cfg.cell_type_column not in adata.obs: + raise ValueError(f"{path} has no obs column '{cfg.cell_type_column}'") + names = _select_features(list(map(str, adata.var_names)), cfg, exclude=set()) + features = None + if names: + matrix = adata[:, names].layers[cfg.h5ad_layer] if cfg.h5ad_layer else adata[:, names].X + features = np.asarray(matrix.toarray() if hasattr(matrix, "toarray") else matrix) + features = features.astype(np.float32) + return CellTable( + coords=coords, + cell_ids=( + adata.obs[cfg.cell_id_column].astype(str).to_numpy() + if cfg.cell_id_column + else adata.obs_names.astype(str).to_numpy() + ), + cell_types=_cell_types(adata.obs[cfg.cell_type_column] if cfg.cell_type_column else None), + features=features, + feature_names=names, + ) + + +def _load_sce(path: str | Path, cfg: InputConfig) -> CellTable: + try: + import rpy2.robjects as ro + from rpy2.robjects import default_converter, numpy2ri, pandas2ri + from rpy2.robjects.conversion import localconverter + except ImportError as exc: + raise ImportError( + "Reading SingleCellExperiment (.rds) inputs needs rpy2 and an R installation " + "with Bioconductor (pip install 'grass-mil[sce]')." + ) from exc + + r = ro.r + with localconverter(default_converter + pandas2ri.converter + numpy2ri.converter): + sce = r["readRDS"](str(path)) + coldata = r["as.data.frame"](r["colData"](sce)) + if cfg.reduced_dim: + coords = np.asarray(r["reducedDim"](sce, cfg.reduced_dim), dtype=np.float64)[:, :2] + else: + coords = coldata[list(cfg.coord_columns)].to_numpy(dtype=np.float64) + assay = r["assay"](sce, cfg.sce_assay) if cfg.sce_assay else r["assay"](sce) + all_names = [str(n) for n in r["rownames"](sce)] + names = _select_features(all_names, cfg, exclude=set()) + features = None + if names: + rows = [all_names.index(n) for n in names] + features = np.asarray(assay, dtype=np.float32)[rows, :].T + cell_ids = ( + coldata[cfg.cell_id_column].astype(str).to_numpy() + if cfg.cell_id_column + else np.asarray(r["colnames"](sce), dtype=str) + ) + if cfg.cell_type_column and cfg.cell_type_column not in coldata.columns: + raise ValueError(f"{path} has no colData column '{cfg.cell_type_column}'") + return CellTable( + coords=coords, + cell_ids=cell_ids, + cell_types=_cell_types(coldata[cfg.cell_type_column] if cfg.cell_type_column else None), + features=features, + feature_names=names, + ) + + +def load_polygons(path: str | Path) -> list[Any]: + """Polygons from a GeoJSON (``.json``/``.geojson``) or WKT (one per line) file.""" + import json + + from shapely import wkt + from shapely.geometry import shape + + path = Path(path) + if path.suffix.lower() in {".json", ".geojson"}: + raw = json.loads(path.read_text()) + if raw.get("type") == "FeatureCollection": + polygons = [shape(feature["geometry"]) for feature in raw["features"]] + else: + polygons = [shape(raw)] + elif path.suffix.lower() in {".wkt", ".txt"}: + polygons = [wkt.loads(line) for line in path.read_text().splitlines() if line.strip()] + else: + raise ValueError(f"Unsupported polygon file {path}: use GeoJSON or WKT.") + if not polygons: + raise ValueError(f"No polygons in {path}") + return polygons diff --git a/src/grass_mil/data/precompute.py b/src/grass_mil/data/precompute.py new file mode 100644 index 0000000..1cd06e6 --- /dev/null +++ b/src/grass_mil/data/precompute.py @@ -0,0 +1,301 @@ +"""Building the cell graphs of a cohort once and storing them on disk. + +The manifest lists one row per region (``sample_id``, ``input_path`` and +optionally ``region_id``, ``input_type``, ``polygons_path``). Every region is +cut into graph units (the whole region, square tiles, or polygons), each unit +becomes a PyG ``Data`` with + +* ``x``: continuous node features (possibly zero columns), +* ``cell_type``: integer cell-type codes (when a cell-type column is set), +* ``pos``, ``edge_index`` and ``edge_attr`` of the Delaunay graph, + +and ``processed_index.json`` records the units, the cell-type vocabulary, the +feature names and a fingerprint of every setting that shaped the graphs. A +later run with different settings refuses to reuse the directory instead of +silently training on stale graphs. Labels are not part of the graphs; they are +attached when units are loaded (:mod:`grass_mil.data.datamodule`). +""" + +from __future__ import annotations + +import hashlib +import json +import logging +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Callable + +import numpy as np +import pandas as pd +import torch +from torch_geometric.data import Data + +from .graphs import GraphConfig, delaunay_edges, edge_features +from .loaders import CellTable, InputConfig, input_type_of, load_cell_table, load_polygons +from .reducers import Pca, ReducerConfig, load_reducer_state +from .units import TilingConfig, grid_tiles, inside_any, polygon_members + +__all__ = ["BuildConfig", "UnitEntry", "ProcessedIndex", "read_manifest", "build_graphs"] + +log = logging.getLogger(__name__) + +INDEX_NAME = "processed_index.json" +INDEX_VERSION = 2 +UNIT_KINDS = ("region", "tile", "polygon") + + +@dataclass(frozen=True) +class BuildConfig: + """Every setting that shapes the processed graphs.""" + + unit: str = "region" + coord_scale_um: float = 1.0 + min_cells: int = 10 + tiling: TilingConfig = field(default_factory=TilingConfig) + input: InputConfig = field(default_factory=InputConfig) + graph: GraphConfig = field(default_factory=GraphConfig) + feature_reducer: ReducerConfig = field(default_factory=ReducerConfig) + + def __post_init__(self) -> None: + if self.unit not in UNIT_KINDS: + raise ValueError(f"data.unit must be one of {list(UNIT_KINDS)}, got '{self.unit}'.") + + +@dataclass(frozen=True) +class UnitEntry: + unit_id: str + region_id: str + sample_id: str + path: str + n_cells: int + + +@dataclass +class ProcessedIndex: + units: list[UnitEntry] + cell_types: list[str] + feature_names: list[str] + edge_features: list[str] + fingerprint: str + feature_reducer: dict + + @classmethod + def read(cls, processed_dir: str | Path) -> ProcessedIndex: + path = Path(processed_dir) / INDEX_NAME + if not path.exists(): + raise FileNotFoundError(f"No processed graphs at {path}; run prepare_data() first.") + raw = json.loads(path.read_text()) + if raw.get("version") != INDEX_VERSION: + raise ValueError( + f"{path} was written by another version of grass-mil; " + "rebuild it with data.force_precompute=true." + ) + return cls( + units=[UnitEntry(**u) for u in raw["units"]], + cell_types=list(raw["cell_types"]), + feature_names=list(raw["feature_names"]), + edge_features=list(raw["edge_features"]), + fingerprint=str(raw["fingerprint"]), + feature_reducer=dict(raw.get("feature_reducer") or {}), + ) + + def write(self, processed_dir: str | Path) -> Path: + payload = {"version": INDEX_VERSION, **asdict(self)} + path = Path(processed_dir) / INDEX_NAME + path.write_text(json.dumps(payload, indent=1)) + return path + + +def read_manifest(path: str | Path) -> pd.DataFrame: + """Manifest rows with resolved input paths and a ``region_id`` per row.""" + path = Path(path) + if not path.exists(): + raise FileNotFoundError(f"Manifest not found: {path}") + df = pd.read_csv(path, dtype=str, keep_default_na=False) + df.columns = [c.strip() for c in df.columns] + for column in ("sample_id", "input_path"): + if column not in df.columns: + raise ValueError(f"Manifest {path} has no '{column}' column.") + df = df.apply(lambda col: col.str.strip()) + if "region_id" not in df.columns: + df["region_id"] = "" + df["region_id"] = df["region_id"].where(df["region_id"] != "", df["sample_id"]) + if df["region_id"].duplicated().any(): + dup = sorted(df.loc[df["region_id"].duplicated(), "region_id"])[:5] + raise ValueError(f"Manifest {path}: region ids must be unique; repeated: {dup}") + + def _resolve(value: str) -> str: + if not value: + return "" + p = Path(value) + return str(p if p.is_absolute() else (path.parent / p).resolve()) + + df["input_path"] = df["input_path"].map(_resolve) + if "polygons_path" in df.columns: + df["polygons_path"] = df["polygons_path"].map(_resolve) + return df + + +def fingerprint(cfg: BuildConfig, manifest: pd.DataFrame, split_key: str = "") -> str: + h = hashlib.sha1(json.dumps(asdict(cfg), sort_keys=True, default=list).encode()) + h.update(manifest.to_csv(index=False).encode()) + h.update(split_key.encode()) + return h.hexdigest()[:16] + + +def _units_of(table: CellTable, row: pd.Series, cfg: BuildConfig) -> list[tuple[str, np.ndarray]]: + region = str(row["region_id"]) + polygons_path = row.get("polygons_path") or None + polygons = load_polygons(polygons_path) if polygons_path else None + if polygons and cfg.coord_scale_um != 1.0: + from shapely import affinity + + scale = cfg.coord_scale_um + polygons = [affinity.scale(p, xfact=scale, yfact=scale, origin=(0, 0)) for p in polygons] + if cfg.unit == "polygon": + if not polygons: + raise ValueError(f"data.unit=polygon but region {region} has no polygons_path.") + parts = [(f"{region}_polygon{k:03d}", idx) for k, idx in enumerate(polygon_members(table.coords, polygons))] + else: + keep = inside_any(table.coords, polygons) if polygons else np.arange(len(table)) + if cfg.unit == "tile": + tiles = grid_tiles(table.coords[keep], cfg.tiling) + parts = [(f"{region}_tile{k:03d}", keep[idx]) for k, idx in enumerate(tiles)] + else: + parts = [(region, keep)] + return [(uid, idx) for uid, idx in parts if idx.size >= cfg.min_cells] + + +def build_graphs( + manifest_path: str | Path, + processed_dir: str | Path, + cfg: BuildConfig, + *, + force: bool = False, + training_units: Callable[[list[UnitEntry]], set[str]] | None = None, + split_key: str = "", +) -> ProcessedIndex: + """Build (or reuse) the processed graphs of a cohort. + + ``training_units`` returns the ids of the training units given all units; + it is needed only for ``feature_reducer.fit_on=train``, in which case + ``split_key`` describes that split so a different split triggers a rebuild. + """ + processed_dir = Path(processed_dir) + manifest = read_manifest(manifest_path) + fp = fingerprint(cfg, manifest, split_key) + if (processed_dir / INDEX_NAME).exists() and not force: + index = ProcessedIndex.read(processed_dir) + if index.fingerprint != fp: + raise ValueError( + f"{processed_dir} holds graphs built with different settings or manifest. " + "Set data.force_precompute=true to rebuild, or use another data.processed_dir." + ) + return index + + log.info(f"Building cell graphs for {len(manifest)} region(s) into {processed_dir}") + tables: list[tuple[pd.Series, CellTable, list[tuple[str, np.ndarray]]]] = [] + for _, row in manifest.iterrows(): + kind = input_type_of(row["input_path"], row.get("input_type") or None) + table = load_cell_table(row["input_path"], kind, cfg.input) + table.coords = table.coords * float(cfg.coord_scale_um) + tables.append((row, table, _units_of(table, row, cfg))) + + cell_types: list[str] = [] + if cfg.input.cell_type_column: + cell_types = sorted({str(t) for _, table, _ in tables for t in table.cell_types}) + code_of = {name: i for i, name in enumerate(cell_types)} + feature_names = tables[0][1].feature_names if tables else [] + if any(table.feature_names != feature_names for _, table, _ in tables): + raise ValueError("All regions must provide the same feature columns.") + + entries = [ + UnitEntry(uid, str(row["region_id"]), str(row["sample_id"]), "", int(idx.size)) + for row, _, units in tables + for uid, idx in units + ] + reducer_state, project = _fit_reducer(tables, entries, cfg.feature_reducer, training_units) + if cfg.feature_reducer.name == "pca" and feature_names: + feature_names = [f"pc{i + 1}" for i in range(min(cfg.feature_reducer.n_components, len(feature_names)))] + + graph_dir = processed_dir / "graphs" + graph_dir.mkdir(parents=True, exist_ok=True) + written: list[UnitEntry] = [] + for row, table, units in tables: + features = project(table, str(row["sample_id"])) + for uid, idx in units: + path = graph_dir / f"{_safe(uid)}.pt" + data = _unit_graph(table, features, idx, code_of, cfg.graph) + torch.save({"data": data, "cell_ids": table.cell_ids[idx].tolist()}, path) + written.append(UnitEntry(uid, str(row["region_id"]), str(row["sample_id"]), str(path), int(idx.size))) + index = ProcessedIndex( + units=written, + cell_types=cell_types, + feature_names=list(feature_names), + edge_features=list(cfg.graph.edge_features), + fingerprint=fp, + feature_reducer={**asdict(cfg.feature_reducer), "state": reducer_state}, + ) + index.write(processed_dir) + log.info(f"Wrote {len(written)} graph unit(s) to {processed_dir}") + return index + + +def _safe(name: str) -> str: + return "".join(c if c.isalnum() or c in "._-" else "_" for c in name) + + +def _fit_reducer(tables, entries, cfg: ReducerConfig, training_units): + """The fitted reducer state and a ``(table, sample_id) -> features`` projection.""" + if cfg.name == "identity" or not any(t.features is not None for _, t, _ in tables): + return None, lambda table, sample_id: table.features + if cfg.state_from: + pca = load_reducer_state(cfg.state_from) + return pca.state(), lambda table, sample_id: pca.transform(table.features) + if cfg.fit_on == "sample": + fitted = { + str(row["sample_id"]): Pca.fit(table.features, cfg.n_components) + for row, table, _ in tables + } + state = {sample: pca.state() for sample, pca in fitted.items()} + return state, lambda table, sample_id: fitted[sample_id].transform(table.features) + if cfg.fit_on == "train": + if training_units is None: + raise ValueError("feature_reducer.fit_on=train needs a training split.") + train = training_units(entries) + parts = [ + table.features[idx] + for _, table, units in tables + for uid, idx in units + if uid in train + ] + else: + parts = [table.features for _, table, _ in tables] + pca = Pca.fit(np.concatenate(parts), cfg.n_components) + return pca.state(), lambda table, sample_id: pca.transform(table.features) + + +def _unit_graph( + table: CellTable, + features: np.ndarray | None, + idx: np.ndarray, + code_of: dict[str, int], + graph: GraphConfig, +) -> Data: + coords = table.coords[idx] + edge_index = delaunay_edges(coords) + x = ( + torch.from_numpy(np.ascontiguousarray(features[idx], dtype=np.float32)) + if features is not None + else torch.zeros((idx.size, 0)) + ) + data = Data( + x=x, + pos=torch.from_numpy(coords).float(), + edge_index=edge_index, + edge_attr=edge_features(coords, edge_index, graph), + num_nodes=int(idx.size), + ) + if code_of: + data.cell_type = torch.tensor([code_of[str(t)] for t in table.cell_types[idx]]) + return data diff --git a/src/grass_mil/data/reducers.py b/src/grass_mil/data/reducers.py new file mode 100644 index 0000000..76ff38a --- /dev/null +++ b/src/grass_mil/data/reducers.py @@ -0,0 +1,83 @@ +"""Optional dimensionality reduction of the continuous node features. + +``identity`` keeps the configured features as they are. ``pca`` projects them +onto principal components fitted on: + +* ``all``: every cell of the cohort (label-free, so it never leaks a split); +* ``train``: the cells of the training units of the configured split. Only + valid for a single random split: cross-validation re-partitions the same + graphs, so a train-only fit would leak into other folds; +* ``sample``: each sample separately (one basis per sample). + +The fitted state is stored in ``processed_index.json`` and can be applied to a +new cohort with ``state_from`` (the index of the cohort it was fitted on), so +both are projected onto the same components. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, fields +from pathlib import Path +from typing import Any, Mapping + +import numpy as np + +__all__ = ["ReducerConfig", "Pca", "load_reducer_state"] + + +@dataclass(frozen=True) +class ReducerConfig: + name: str = "identity" # identity | pca + n_components: int = 32 + fit_on: str = "all" # all | train | sample + state_from: str | None = None + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any] | None) -> ReducerConfig: + cfg = dict(cfg or {}) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown data.feature_reducer keys: {sorted(unknown)}") + out = cls(**cfg) + if out.name not in {"identity", "pca"}: + raise ValueError("data.feature_reducer.name must be 'identity' or 'pca'.") + if out.fit_on not in {"all", "train", "sample"}: + raise ValueError("data.feature_reducer.fit_on must be 'all', 'train' or 'sample'.") + if out.state_from is not None and out.fit_on == "sample": + raise ValueError("feature_reducer.state_from needs a shared basis (fit_on all/train).") + return out + + +@dataclass +class Pca: + mean: np.ndarray # (f,) + components: np.ndarray # (k, f) + + @classmethod + def fit(cls, features: np.ndarray, n_components: int) -> Pca: + mean = features.mean(axis=0) + _, _, vt = np.linalg.svd(features - mean, full_matrices=False) + return cls(mean=mean, components=vt[:n_components]) + + def transform(self, features: np.ndarray) -> np.ndarray: + return ((features - self.mean) @ self.components.T).astype(np.float32) + + def state(self) -> dict[str, list]: + return {"mean": self.mean.tolist(), "components": self.components.tolist()} + + @classmethod + def from_state(cls, state: Mapping[str, Any]) -> Pca: + return cls( + mean=np.asarray(state["mean"], dtype=np.float64), + components=np.asarray(state["components"], dtype=np.float64), + ) + + +def load_reducer_state(index_path: str | Path) -> Pca: + """The shared PCA basis stored in another cohort's ``processed_index.json``.""" + index = json.loads(Path(index_path).read_text()) + state = (index.get("feature_reducer") or {}).get("state") + if not state or "mean" not in state: + raise ValueError(f"{index_path} holds no shared PCA basis.") + return Pca.from_state(state) diff --git a/src/grass_mil/data/samplers.py b/src/grass_mil/data/samplers.py new file mode 100644 index 0000000..87703eb --- /dev/null +++ b/src/grass_mil/data/samplers.py @@ -0,0 +1,291 @@ +"""Ego-graph sampling: every instance of a bag is the k-hop ego-graph of one cell. + +:class:`EgoGraphLoader` turns a dataset of graph units into batches of +ego-graphs (one subgraph per root cell, built by PyG's ``ShaDowKHopSampler``): + +* ``bags``: one batch per unit holding the ego-graphs of all of its selected + roots, so the attention softmax and the loss of a bag see the whole unit + (MIL training); +* ``chunks``: the selected roots of each unit in batches of ``batch_size`` + (evaluation, embedding extraction); +* ``merged``: the roots of all units pooled and shuffled into batches of + ``batch_size`` (label-free pretraining). + +Every batch carries, per ego-graph, ``root_cell`` (the root's node index in +its unit) and the unit's ids and ``graph_y`` target row. +""" + +from __future__ import annotations + +import math +import zlib +from dataclasses import dataclass, field, fields +from typing import Any, Iterator, Mapping + +import torch +from torch_geometric.data import Batch, Data +from torch_geometric.loader import ShaDowKHopSampler +from torch_geometric.typing import WITH_TORCH_SPARSE + +from .ego_radius import EgoRadiusConfig, EgoRadiusCutoff, radius_for_depth +from .interior import InteriorCache, InteriorConfig + +__all__ = ["SamplerConfig", "EgoGraphLoader"] + +MODES = ("bags", "chunks", "merged") +ID_KEYS = ("unit_id", "region_id", "sample_id") + + +@dataclass(frozen=True) +class SamplerConfig: + """Ego-graph sampling settings. + + Attributes: + depth: Hops of every ego-graph. + num_neighbors: Neighbours kept per node and hop (``-1``: all, i.e. the + full k-hop neighbourhood). + roots_per_unit: Root cells drawn per unit and epoch (``None``: all). + batch_size: Ego-graphs per batch in ``chunks`` and ``merged`` mode. + merge_units: Pool the roots of all units (pretraining) instead of one + bag per unit. + root_weighting: ``None`` draws roots uniformly; ``inverse`` / + ``sqrt_inverse`` weight them by the inverse (square-root) frequency + of their cell type within the unit. + interior_roots: Restrict roots to the tissue interior. + ego_radius: Cap the physical extent of every ego-graph. + """ + + depth: int = 4 + num_neighbors: int = -1 + roots_per_unit: int | None = 256 + batch_size: int = 64 + merge_units: bool = False + root_weighting: str | None = None + interior_roots: InteriorConfig = field(default_factory=InteriorConfig) + ego_radius: EgoRadiusConfig = field(default_factory=EgoRadiusConfig) + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any] | None) -> SamplerConfig: + cfg = dict(cfg or {}) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown data.sampler keys: {sorted(unknown)}") + cfg["interior_roots"] = InteriorConfig.from_dict(cfg.get("interior_roots")) + cfg["ego_radius"] = EgoRadiusConfig.from_dict(cfg.get("ego_radius")) + out = cls(**cfg) + if out.depth < 1: + raise ValueError("data.sampler.depth must be >= 1.") + if out.roots_per_unit is not None and out.roots_per_unit < 1: + raise ValueError("data.sampler.roots_per_unit must be positive or null.") + if out.root_weighting not in (None, "inverse", "sqrt_inverse"): + raise ValueError("data.sampler.root_weighting must be null, inverse or sqrt_inverse.") + return out + + def as_dict(self) -> dict[str, Any]: + from dataclasses import asdict + + return asdict(self) + + +class _EgoGraphBatches(ShaDowKHopSampler): + """``ShaDowKHopSampler`` that also records each ego-graph's root cell and + applies the optional radius cap.""" + + def __init__(self, data: Data, roots: torch.Tensor, *, depth: int, num_neighbors: int, + cutoff: EgoRadiusCutoff | None, **loader_kwargs: Any) -> None: + data = data.clone() + data.node_position = torch.arange(int(data.num_nodes)) + self.cutoff = cutoff + neighbours = int(num_neighbors) if num_neighbors >= 0 else int(data.num_nodes) + super().__init__(data, depth=depth, num_neighbors=neighbours, node_idx=roots, **loader_kwargs) + + def __collate__(self, n_id): + batch = super().__collate__(n_id) + n = int(batch.ptr.numel()) - 1 + batch.root_cell = batch.node_position[batch.root_n_id] + del batch.node_position + if self.cutoff is not None: + batch = _cut_each(batch, self.cutoff) + y = getattr(batch, "graph_y", None) + if isinstance(y, torch.Tensor): + batch.graph_y = y.view(1, -1).expand(n, -1).contiguous() + for key in ID_KEYS: + if isinstance(getattr(batch, key, None), str): + batch[key] = [batch[key]] * n + return batch + + +def _cut_each(batch: Batch, cutoff: EgoRadiusCutoff) -> Batch: + """Apply the radius cap to every ego-graph of a ShaDow batch and re-batch.""" + ptr = batch.ptr + n_nodes = int(batch.num_nodes) + n_edges = int(batch.edge_index.size(1)) + node_keys = [k for k, v in batch if isinstance(v, torch.Tensor) and v.dim() and v.size(0) == n_nodes and k not in ("batch",)] + edge_owner = batch.batch[batch.edge_index[0]] + order = torch.argsort(edge_owner, stable=True) + edge_counts = torch.bincount(edge_owner, minlength=int(ptr.numel()) - 1).tolist() + edge_index = batch.edge_index[:, order] + edge_attr = batch.edge_attr[order] if getattr(batch, "edge_attr", None) is not None else None + parts, roots = [], [] + start_e = 0 + for i, count in enumerate(edge_counts): + lo, hi = int(ptr[i]), int(ptr[i + 1]) + sub = Data(edge_index=edge_index[:, start_e : start_e + count] - lo, num_nodes=hi - lo) + if edge_attr is not None and edge_attr.size(0) == n_edges: + sub.edge_attr = edge_attr[start_e : start_e + count] + for key in node_keys: + sub[key] = batch[key][lo:hi] + start_e += count + sub, root = cutoff(sub, int(batch.root_n_id[i]) - lo) + parts.append(sub) + roots.append(root) + out = Batch.from_data_list(parts) + out.root_n_id = out.ptr[:-1] + torch.tensor(roots, dtype=torch.long) + out.root_cell = batch.root_cell + for key, value in batch: + if key not in out and key not in ("batch", "ptr", "edge_index", "edge_attr"): + out[key] = value + return out + + +class EgoGraphLoader: + """Iterable of ego-graph batches over a dataset of graph units. + + Args: + dataset: Units with ``entries`` (``unit_id``, ``n_cells``) and + ``__getitem__`` returning the unit graph. + cfg: Sampling settings. + mode: ``bags`` | ``chunks`` | ``merged``. + roots_per_unit: Overrides ``cfg.roots_per_unit`` (``None``: all roots). + seed: Fixes each unit's root subset (from the seed and the unit id), so + it is identical on every pass; ``None`` redraws every pass. + shuffle: Shuffle units (and roots in ``merged`` mode). + num_workers: DataLoader workers; used in ``merged`` mode, where one + loader spans the epoch. Per-unit modes collate in the main process. + """ + + def __init__( + self, + dataset, + cfg: SamplerConfig, + *, + mode: str, + roots_per_unit: int | None, + seed: int | None = None, + shuffle: bool = False, + num_workers: int = 0, + pin_memory: bool = False, + ) -> None: + if mode not in MODES: + raise ValueError(f"mode must be one of {MODES}") + if not WITH_TORCH_SPARSE: + raise ImportError( + "Ego-graph sampling needs torch-sparse; install the wheel matching your PyTorch " + "build, e.g. pip install torch-sparse -f https://data.pyg.org/whl/torch-.html" + ) + self.dataset = dataset + self.cfg = cfg + self.mode = mode + self.roots_per_unit = roots_per_unit + self.seed = seed + self.shuffle = shuffle + self.num_workers = int(num_workers) + self.pin_memory = bool(pin_memory) + self._interior = InteriorCache(cfg.interior_roots.cache_dir) + self._cutoff = ( + EgoRadiusCutoff( + radius=radius_for_depth(cfg.depth, cfg.ego_radius), + keep_root_component=cfg.ego_radius.keep_root_component, + ) + if cfg.ego_radius.enabled + else None + ) + + # ------------------------------------------------------------------ roots + def _pool(self, data: Data) -> torch.Tensor: + if not self.cfg.interior_roots.enabled: + return torch.arange(int(data.num_nodes)) + idx = self._interior.get(data, self.cfg.interior_roots, n_hops=self.cfg.depth) + return torch.as_tensor(idx, dtype=torch.long) + + def _n_roots(self, pool_size: int) -> int: + k = self.roots_per_unit + return pool_size if k is None else min(int(k), pool_size) + + def _generator(self, unit_id: str) -> torch.Generator | None: + if self.seed is None: + return None + return torch.Generator().manual_seed(int(self.seed) * 7919 + zlib.crc32(unit_id.encode())) + + def roots(self, data: Data) -> torch.Tensor: + """Root cells of one unit for this pass (sorted node indices).""" + pool = self._pool(data) + k = self._n_roots(int(pool.numel())) + if k >= pool.numel(): + return pool + gen = self._generator(str(data.unit_id)) + if self.cfg.root_weighting and getattr(data, "cell_type", None) is not None: + types = data.cell_type[pool] + counts = torch.bincount(types).float() + weights = 1.0 / counts[types] + if self.cfg.root_weighting == "sqrt_inverse": + weights = weights.sqrt() + chosen = torch.multinomial(weights, k, replacement=False, generator=gen) + else: + chosen = torch.randperm(int(pool.numel()), generator=gen)[:k] + return pool[chosen.sort().values] + + # ---------------------------------------------------------------- batches + def _batches(self, data: Data, roots: torch.Tensor, batch_size: int, shuffle: bool, **kw): + return _EgoGraphBatches( + data, + roots, + depth=self.cfg.depth, + num_neighbors=self.cfg.num_neighbors, + cutoff=self._cutoff, + batch_size=batch_size, + shuffle=shuffle, + **kw, + ) + + def __len__(self) -> int: + if self.mode == "bags": + return len(self.dataset) + if self.cfg.interior_roots.enabled: + sizes = [int(self._pool(self.dataset[i]).numel()) for i in range(len(self.dataset))] + else: + sizes = [int(e.n_cells) for e in self.dataset.entries] + n_roots = [self._n_roots(s) for s in sizes] + if self.mode == "merged": + return math.ceil(sum(n_roots) / self.cfg.batch_size) + return sum(math.ceil(n / self.cfg.batch_size) for n in n_roots) + + def __iter__(self) -> Iterator[Batch]: + if self.mode == "merged": + yield from self._iter_merged() + return + order = torch.randperm(len(self.dataset)) if self.shuffle else torch.arange(len(self.dataset)) + for i in order.tolist(): + data = self.dataset[i] + roots = self.roots(data) + if roots.numel() == 0: + continue + size = int(roots.numel()) if self.mode == "bags" else self.cfg.batch_size + yield from self._batches(data, roots, size, shuffle=False) + + def _iter_merged(self) -> Iterator[Batch]: + units = [self.dataset[i] for i in range(len(self.dataset))] + offsets = [0] + for unit in units: + offsets.append(offsets[-1] + int(unit.num_nodes)) + roots = torch.cat([self.roots(u) + offsets[k] for k, u in enumerate(units)]) + merged = Data( + edge_index=torch.cat([u.edge_index + offsets[k] for k, u in enumerate(units)], dim=1), + num_nodes=offsets[-1], + ) + for key in ("x", "pos", "cell_type", "edge_attr"): + values = [getattr(u, key, None) for u in units] + if all(isinstance(v, torch.Tensor) for v in values): + merged[key] = torch.cat(values) + kw = {"num_workers": self.num_workers, "pin_memory": self.pin_memory} + yield from self._batches(merged, roots, self.cfg.batch_size, shuffle=self.shuffle, **kw) diff --git a/src/grass_mil/data/spatial_omics_datamodule.py b/src/grass_mil/data/spatial_omics_datamodule.py deleted file mode 100644 index ceb2434..0000000 --- a/src/grass_mil/data/spatial_omics_datamodule.py +++ /dev/null @@ -1,355 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Dict, Optional, Tuple - -import lightning as L - -from .components.datasets import SpatialOmicsGraphDataset, TransformDataset -from .components.feature_reducers import FeatureReducerConfig -from .components.graph_builders import GraphBuilderConfig -from .components.loaders import CsvConfig, H5adConfig, SceConfig -from .components.patching import TileConfig -from .components.splits import ExplicitSplit -from .components.precompute import ( - CategoricalFeatureConfig, - GraphLabelConfig, - LoocvConfig, - ManifestConfig, - PrecomputeConfig, - SpatialOmicsPreprocessor, -) -from .components.samplers import ( - BaseSamplerStrategy, - SamplerConfig, - get_sampler_strategy, -) -from .components.transforms import instantiate_transforms - - -@dataclass -class SplitConfig: - train_val_test_split: Tuple[float, float, float] - split_by: str - loocv: LoocvConfig - - -class SpatialOmicsDataModule(L.LightningDataModule): - def __init__( - self, - data_dir: str, - raw_manifest_path: str, - processed_dir: str, - coord_scale_um: float, - sample_unit: str, - batch_size: int, - num_workers: int, - pin_memory: bool, - reducer_scope: str, - keep_raw_molecular: bool, - force_precompute: bool, - min_cells: int, - use_molecular_features: bool, - categorical_features: Dict[str, Any], - split: Dict[str, Any], - csv: Dict[str, Any], - h5ad: Dict[str, Any], - sce: Dict[str, Any], - graph_builder: Dict[str, Any], - feature_reducer: Dict[str, Any], - tiling: Dict[str, Any], - sampler: Optional[Dict[str, Any]] = None, - val_sampler: Optional[Dict[str, Any]] = None, - test_sampler: Optional[Dict[str, Any]] = None, - manifest: Optional[Dict[str, Any]] = None, - graph_labels: Optional[Dict[str, Any]] = None, - transforms: Optional[list[Dict[str, Any]]] = None, - split_seed: int = 42, - ) -> None: - super().__init__() - self.save_hyperparameters(logger=False) - - self.data_dir = Path(data_dir) - self.raw_manifest_path = raw_manifest_path - self.processed_dir = processed_dir - - self.dataset_train: Optional[SpatialOmicsGraphDataset] = None - self.dataset_val: Optional[SpatialOmicsGraphDataset] = None - self.dataset_test: Optional[SpatialOmicsGraphDataset] = None - self.split_seed = int(split_seed) - - self.split_config = SplitConfig( - train_val_test_split=tuple(split["train_val_test_split"]), - split_by=split["split_by"], - loocv=LoocvConfig.from_dict(split.get("loocv")), - ) - # Explicit id-list splits (nested CV): applied in setup() on top of the - # persisted precompute, without re-processing the graphs. - self.explicit_split = ExplicitSplit.from_dict(split.get("explicit")) - if self.explicit_split is not None: - reducer_name = str(feature_reducer.get("name", "identity")).strip().lower() - fit_mode = str(feature_reducer.get("fit_mode", "train_only")).strip().lower() - if reducer_name != "identity" and fit_mode == "train_only": - raise ValueError( - "data.split.explicit re-partitions an existing precompute, so a feature " - "reducer fitted on the persisted training split would leak into the new " - "test split. Use feature_reducer.name=identity or fit_mode=global." - ) - self.csv_config = CsvConfig.from_dict(csv, use_molecular_features) - self.h5ad_config = H5adConfig.from_dict(h5ad, use_molecular_features) - self.sce_config = SceConfig.from_dict(sce, use_molecular_features) - self.graph_builder_config = GraphBuilderConfig.from_dict(graph_builder) - self.feature_reducer_config = FeatureReducerConfig.from_dict(feature_reducer) - self.tile_config = TileConfig.from_dict(tiling) - self.sampler_config = SamplerConfig.from_dict(sampler) - self.val_sampler_config = ( - SamplerConfig.from_dict(val_sampler) if val_sampler is not None else None - ) - # Test/predict loader. `null` derives it from the training sampler with - # every root kept once (no subsampling, no proportional weighting, no - # unit-as-bag), so a held-out region is scored on all of its cells. - self.test_sampler_config = ( - SamplerConfig.from_dict(test_sampler) - if test_sampler is not None - else _all_roots_sampler_config(self.sampler_config) - ) - self.sampler_strategy: Optional[BaseSamplerStrategy] = None - self.val_sampler_strategy: Optional[BaseSamplerStrategy] = None - self.test_sampler_strategy: Optional[BaseSamplerStrategy] = None - self.manifest_config = ManifestConfig.from_dict(manifest) - self.graph_label_config = GraphLabelConfig.from_dict(graph_labels) - self.transforms = instantiate_transforms(transforms) - - self.precompute_config = PrecomputeConfig.from_args( - raw_manifest_path=self.raw_manifest_path, - processed_dir=self.processed_dir, - coord_scale_um=coord_scale_um, - sample_unit=sample_unit, - reducer_scope=reducer_scope, - reducer_fit_mode=str(feature_reducer.get("fit_mode", "train_only")), - keep_raw_molecular=keep_raw_molecular, - force=force_precompute, - min_cells=min_cells, - use_molecular_features=use_molecular_features, - split_by=self.split_config.split_by, - split_ratios=self.split_config.train_val_test_split, - split_seed=self.split_seed, - loocv=split.get("loocv"), - ) - self.categorical_feature_config = CategoricalFeatureConfig.from_dict( - categorical_features - ) - - def prepare_data(self) -> None: - if ( - self.split_config.loocv.enabled - and self.hparams.reducer_scope == "dataset" - and str(self.hparams.feature_reducer.get("fit_mode", "train_only")) - .strip() - .lower() - == "train_only" - ): - processed_index_path = Path(self.processed_dir) / "processed_index.json" - if ( - not bool(self.hparams.force_precompute) - and processed_index_path.exists() - ): - raise ValueError( - "LOOCV with reducer_scope='dataset' and " - "feature_reducer.fit_mode='train_only' cannot reuse " - "an existing " - "processed_index when data.force_precompute=false. Set " - "data.force_precompute=true or use a fresh fold-specific " - "processed_dir." - ) - preprocessor = self._build_preprocessor() - preprocessor.precompute() - - def _build_preprocessor(self) -> SpatialOmicsPreprocessor: - return SpatialOmicsPreprocessor( - manifest_config=self.manifest_config, - precompute_config=self.precompute_config, - csv_config=self.csv_config, - h5ad_config=self.h5ad_config, - sce_config=self.sce_config, - graph_builder_config=self.graph_builder_config, - feature_reducer_config=self.feature_reducer_config, - categorical_feature_config=self.categorical_feature_config, - tile_config=self.tile_config, - graph_label_config=self.graph_label_config, - ) - - def setup(self, stage: Optional[str] = None) -> None: - index_path = Path(self.processed_dir) / "processed_index.json" - base_dataset = SpatialOmicsGraphDataset(index_path) - - if self.explicit_split is not None: - members = _apply_explicit_split(base_dataset, self.explicit_split) - else: - if not _entries_have_persisted_splits(base_dataset.entries): - raise ValueError( - "processed_index.json must contain valid per-entry split labels " - "('train', 'val', 'test'). Run data precompute/prepare_data() " - "to regenerate processed artifacts." - ) - members = { - name: [entry for entry in base_dataset.entries if entry.split == name] - for name in ("train", "val", "test") - } - - self.dataset_train = _subset_dataset_from_entries( - base_dataset, members["train"], self.transforms - ) - self.dataset_val = _subset_dataset_from_entries(base_dataset, members["val"], self.transforms) - self.dataset_test = _subset_dataset_from_entries( - base_dataset, members["test"], self.transforms - ) - self._ensure_sampler_strategy() - self._ensure_val_sampler_strategy() - - def _ensure_sampler_strategy(self) -> BaseSamplerStrategy: - if self.sampler_strategy is None: - self.sampler_strategy = get_sampler_strategy(self.sampler_config) - return self.sampler_strategy - - def _ensure_val_sampler_strategy(self) -> BaseSamplerStrategy: - if self.val_sampler_config is None: - return self._ensure_sampler_strategy() - if self.val_sampler_strategy is None: - self.val_sampler_strategy = get_sampler_strategy(self.val_sampler_config) - return self.val_sampler_strategy - - def _ensure_test_sampler_strategy(self) -> BaseSamplerStrategy: - if self.test_sampler_config is None: - return self._ensure_sampler_strategy() - if self.test_sampler_strategy is None: - self.test_sampler_strategy = get_sampler_strategy(self.test_sampler_config) - return self.test_sampler_strategy - - def train_dataloader(self) -> Any: - sampler = self._ensure_sampler_strategy() - return sampler.build_dataset_loader( - dataset=self.dataset_train, - batch_size=self.hparams.batch_size, - num_workers=self.hparams.num_workers, - pin_memory=self.hparams.pin_memory, - shuffle=True, - ) - - def val_dataloader(self) -> Any: - sampler = self._ensure_val_sampler_strategy() - return sampler.build_dataset_loader( - dataset=self.dataset_val, - batch_size=self.hparams.batch_size, - num_workers=self.hparams.num_workers, - pin_memory=self.hparams.pin_memory, - shuffle=False, - ) - - def test_dataloader(self) -> Any: - sampler = self._ensure_test_sampler_strategy() - return sampler.build_dataset_loader( - dataset=self.dataset_test, - batch_size=self.hparams.batch_size, - num_workers=self.hparams.num_workers, - pin_memory=self.hparams.pin_memory, - shuffle=False, - ) - - def predict_dataloader(self) -> Any: - """Inference loader, over the test split. - - Required by `grass-mil-predict`: `Trainer.predict` raises - `MisconfigurationException` without it. Predicting on the held-out split - mirrors `test_dataloader` so exported tables line up with test metrics. - To score a different set of graphs, point `data.processed_dir` at a - precompute whose entries carry the split you want. - """ - return self.test_dataloader() - - -def unit_key(entry, id_kind: str) -> str: - """The id a processed entry is addressed by: region (default), sample or patch. - - ``patch`` makes every tile its own unit, which is how a cohort of a few - large images is cut into many bags for cross-validation. - """ - if id_kind == "patch": - return str(entry.patch_id) - if id_kind == "sample": - return str(entry.sample_id) - return str(entry.region_id if entry.region_id is not None else entry.sample_id) - - -def _apply_explicit_split( - dataset: SpatialOmicsGraphDataset, split: ExplicitSplit -) -> Dict[str, list]: - """Entries per split from explicit id lists; ids not listed are dropped. - - The same id may sit in ``val`` and ``test`` (refit runs log the held-out fold - as validation and score it as test), so membership is resolved per split - rather than by a single label on the entry. - """ - import dataclasses - - by_key: Dict[str, Any] = {} - for entry in dataset.entries: - key = str(unit_key(entry, split.id_kind)) - by_key.setdefault(key, []).append(entry) - wanted = {name: [str(r) for r in getattr(split, name)] for name in ("train", "val", "test")} - missing = sorted({r for ids in wanted.values() for r in ids} - set(by_key)) - if missing: - raise ValueError( - f"split.explicit names {len(missing)} {split.id_kind} id(s) with no processed " - f"graph (first: {missing[:10]}). Check the ids against processed_index.json." - ) - members: Dict[str, list] = {} - for name, ids in wanted.items(): - members[name] = [ - dataclasses.replace(entry, split=name) for rid in ids for entry in by_key[rid] - ] - dataset.entries = [e for group in members.values() for e in group] - return members - - -def _all_roots_sampler_config(train_cfg: SamplerConfig) -> Optional[SamplerConfig]: - """Evaluation variant of a sampler config: every root once, uniformly, no bagging.""" - if train_cfg is None or train_cfg.name == "identity": - return None - runtime = dict(train_cfg.runtime or {}) - runtime.update( - { - "subsample_fraction": 1.0, - "subsample_count": None, - "subsample_seed": None, - "fixed_per_unit_seed": None, - "proportional_root_sampling": False, - "unit_as_bag": False, - "merge_units": False, - "shuffle_subgraphs": False, - } - ) - return SamplerConfig(name=train_cfg.name, kwargs=dict(train_cfg.kwargs or {}), runtime=runtime) - - -def _subset_dataset_from_entries( - dataset: SpatialOmicsGraphDataset, - selected_entries, - transforms, -) -> SpatialOmicsGraphDataset: - subset = SpatialOmicsGraphDataset(dataset.index_path) - subset.entries = list(selected_entries) - subset.label_maps = dataset.label_maps - subset.reducer_state = dataset.reducer_state - subset.graph_label_maps = dataset.graph_label_maps - if transforms: - return TransformDataset(subset, transforms) - return subset - - -def _entries_have_persisted_splits(entries) -> bool: - if not entries: - return True - valid = {"train", "val", "test"} - return all(getattr(entry, "split", None) in valid for entry in entries) diff --git a/src/grass_mil/data/splits.py b/src/grass_mil/data/splits.py new file mode 100644 index 0000000..88ea990 --- /dev/null +++ b/src/grass_mil/data/splits.py @@ -0,0 +1,125 @@ +"""Train/validation/test splits, outer fold files and inner cross-validation folds. + +All splits are over groups of units: units of the same patient (or, without a +patient column, of the same region) always fall on the same side. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Mapping, Sequence + +import numpy as np + +__all__ = ["read_ids", "load_folds", "random_split", "inner_folds"] + + +def read_ids(value: Any) -> list[str] | None: + """Ids from a list, or from a ``.json`` (list) or text (one per line) file.""" + if value is None: + return None + if isinstance(value, (str, Path)): + path = Path(str(value)) + if path.suffix.lower() == ".json": + return [str(v) for v in json.loads(path.read_text())] + return [line.strip() for line in path.read_text().splitlines() if line.strip()] + return [str(v) for v in value] + + +def load_folds(path: str | Path) -> list[list[str]]: + """Outer folds: a JSON list of lists of unit ids, each id in at most one fold.""" + folds = json.loads(Path(path).read_text()) + if not isinstance(folds, list) or not all(isinstance(f, list) for f in folds): + raise ValueError(f"{path} must hold a JSON list of lists of unit ids.") + folds = [[str(r) for r in fold] for fold in folds] + seen: set[str] = set() + for fold in folds: + repeated = seen & set(fold) + if repeated: + raise ValueError(f"{path}: ids in more than one fold: {sorted(repeated)[:10]}") + seen |= set(fold) + return folds + + +def random_split( + groups: Mapping[str, str], + ratios: Sequence[float], + seed: int, + strata: Mapping[str, str] | None = None, +) -> dict[str, str]: + """Assign every unit to ``train``/``val``/``test`` by shuffling its group. + + ``groups`` maps unit id to group id; ``strata`` (optional) maps group id to + a stratum (the class), and every stratum is split separately so each split + holds every class when possible. Group counts per split follow ``ratios`` + (largest remainder); a split with a positive ratio gets at least one group + of a stratum when the stratum has enough groups. + """ + ratios = np.asarray(ratios, dtype=float) + if ratios.shape != (3,) or not np.isclose(ratios.sum(), 1.0) or (ratios < 0).any(): + raise ValueError("data.split.ratios must be three non-negative numbers summing to 1.") + rng = np.random.default_rng(int(seed)) + by_stratum: dict[str, list[str]] = {} + for name in sorted(set(groups.values())): + by_stratum.setdefault(str(strata[name]) if strata else "", []).append(name) + split_of_group: dict[str, str] = {} + for stratum in sorted(by_stratum): + names = by_stratum[stratum] + order = rng.permutation(len(names)) + raw = ratios * len(names) + counts = np.floor(raw).astype(int) + for i in np.argsort(-(raw - counts))[: len(names) - int(counts.sum())]: + counts[i] += 1 + for i in np.flatnonzero((ratios > 0) & (counts == 0)): + donor = int(np.argmax(counts)) + if counts[donor] > 1: + counts[donor] -= 1 + counts[i] += 1 + start = 0 + for split, count in zip(("train", "val", "test"), counts): + for j in order[start : start + count]: + split_of_group[names[j]] = split + start += count + return {unit: split_of_group[group] for unit, group in groups.items()} + + +def inner_folds( + ids: Sequence[str], + *, + groups: Mapping[str, str], + classes: Mapping[str, int] | None, + k: int, + seed: int, + max_reseeds: int = 20, +) -> list[tuple[list[str], list[str]]]: + """``k`` group-aware inner folds over the sorted ids. + + With ``classes`` (classification) the folds are stratified, and a split + leaving a fold without one of the classes is redrawn with the next seed (up + to ``max_reseeds`` times), because AUROC is undefined on such a fold. + Sorting the ids first makes the folds a function of the id set alone. + """ + from sklearn.model_selection import GroupKFold, StratifiedGroupKFold + + ids = np.array(sorted(str(r) for r in ids)) + g = np.array([groups[r] for r in ids]) + if len(np.unique(g)) < k: + raise ValueError(f"Inner folds need at least k={k} groups; got {len(np.unique(g))}.") + if classes is None: + splitter = GroupKFold(n_splits=int(k)) + return [(ids[tr].tolist(), ids[va].tolist()) for tr, va in splitter.split(ids, groups=g)] + + y = np.array([classes[r] for r in ids]) + n_classes = len(np.unique(y)) + if n_classes < 2: + raise ValueError("Inner folds need at least two classes among the development units.") + for attempt in range(int(max_reseeds) + 1): + splitter = StratifiedGroupKFold(n_splits=int(k), shuffle=True, random_state=seed + attempt) + folds = [(tr, va) for tr, va in splitter.split(ids, y, g)] + if all(len(np.unique(y[tr])) == n_classes and len(np.unique(y[va])) > 1 for tr, va in folds): + return [(ids[tr].tolist(), ids[va].tolist()) for tr, va in folds] + raise ValueError( + f"Could not build {k} stratified, group-aware inner folds with every class in every " + f"split after {max_reseeds + 1} seeds. Use fewer inner folds or more patients." + ) diff --git a/src/grass_mil/data/units.py b/src/grass_mil/data/units.py new file mode 100644 index 0000000..f604c04 --- /dev/null +++ b/src/grass_mil/data/units.py @@ -0,0 +1,64 @@ +"""Cutting a region into graph units: the whole region, square tiles or polygons.""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import Any, Mapping, Sequence + +import numpy as np + +__all__ = ["TilingConfig", "grid_tiles", "polygon_members", "inside_any"] + + +@dataclass(frozen=True) +class TilingConfig: + """Square tiles of ``size`` with a step of ``stride`` (coordinate units).""" + + size: float = 200.0 + stride: float = 200.0 + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any] | None) -> TilingConfig: + cfg = dict(cfg or {}) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown data.tiling keys: {sorted(unknown)}") + out = cls(**{k: float(v) for k, v in cfg.items()}) + if out.size <= 0 or out.stride <= 0: + raise ValueError("data.tiling.size and stride must be positive.") + return out + + +def grid_tiles(coords: np.ndarray, cfg: TilingConfig) -> list[np.ndarray]: + """Cell indices of every tile of a regular grid anchored at the minimum corner.""" + if coords.size == 0: + return [] + lo = coords.min(axis=0) + hi = coords.max(axis=0) + tiles = [] + for x0 in np.arange(lo[0], hi[0] + cfg.stride, cfg.stride): + for y0 in np.arange(lo[1], hi[1] + cfg.stride, cfg.stride): + inside = ( + (coords[:, 0] >= x0) + & (coords[:, 0] < x0 + cfg.size) + & (coords[:, 1] >= y0) + & (coords[:, 1] < y0 + cfg.size) + ) + if inside.any(): + tiles.append(np.flatnonzero(inside)) + return tiles + + +def polygon_members(coords: np.ndarray, polygons: Sequence[Any]) -> list[np.ndarray]: + """Cell indices inside each polygon.""" + import shapely + + return [np.flatnonzero(shapely.contains_xy(p, coords[:, 0], coords[:, 1])) for p in polygons] + + +def inside_any(coords: np.ndarray, polygons: Sequence[Any]) -> np.ndarray: + """Indices of the cells inside the union of ``polygons``.""" + import shapely + + union = shapely.union_all(list(polygons)) + return np.flatnonzero(shapely.contains_xy(union, coords[:, 0], coords[:, 1])) diff --git a/src/grass_mil/eval.py b/src/grass_mil/eval.py deleted file mode 100644 index 288d5f6..0000000 --- a/src/grass_mil/eval.py +++ /dev/null @@ -1,73 +0,0 @@ -from typing import Any, Dict, List, Tuple - -import hydra -from lightning import LightningDataModule, LightningModule, Trainer -from lightning.pytorch.loggers import Logger -from omegaconf import DictConfig - -from grass_mil._paths import setup_project_root - -setup_project_root() - -from grass_mil.utils import ( - RankedLogger, - extras, - instantiate_loggers, - log_hyperparameters, - task_wrapper, -) - -log = RankedLogger(__name__, rank_zero_only=True) - - -@task_wrapper -def evaluate(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: - """Evaluate a checkpoint on the configured test set.""" - ckpt_path = cfg.get("ckpt_path") - if ckpt_path is None or not str(ckpt_path).strip(): - raise ValueError("Missing required `ckpt_path` for evaluation.") - - log.info(f"Instantiating datamodule <{cfg.data._target_}>") - datamodule: LightningDataModule = hydra.utils.instantiate(cfg.data) - - if cfg.get("model") is None: - raise ValueError("No model config provided. Set `model=...` in your Hydra config.") - log.info(f"Instantiating model <{cfg.model._target_}>") - model: LightningModule = hydra.utils.instantiate(cfg.model) - - log.info("Instantiating loggers...") - logger: List[Logger] = instantiate_loggers(cfg.get("logger")) - - log.info(f"Instantiating trainer <{cfg.trainer._target_}>") - trainer: Trainer = hydra.utils.instantiate(cfg.trainer, logger=logger) - - object_dict = { - "cfg": cfg, - "datamodule": datamodule, - "model": model, - "logger": logger, - "trainer": trainer, - } - - if logger: - log.info("Logging hyperparameters!") - log_hyperparameters(object_dict) - - log.info("Starting testing!") - trainer.test(model=model, datamodule=datamodule, ckpt_path=cfg.ckpt_path) - - metric_dict = trainer.callback_metrics - - return metric_dict, object_dict - - -@hydra.main(version_base="1.3", config_path="configs", config_name="eval.yaml") -def main(cfg: DictConfig) -> None: - """Entry point for evaluation.""" - extras(cfg) - - evaluate(cfg) - - -if __name__ == "__main__": - main() diff --git a/src/grass_mil/extract.py b/src/grass_mil/extract.py index bf9743f..18b0811 100644 --- a/src/grass_mil/extract.py +++ b/src/grass_mil/extract.py @@ -1,142 +1,87 @@ -"""``grass-mil-extract``: per-cell cache of trained models (see inference/cell_cache.py).""" +"""``grass-mil-extract``: per-cell outputs of trained models (see inference/cell_cache.py).""" from __future__ import annotations import json +import logging from pathlib import Path -from typing import Any, Dict, List, Optional +from typing import Any import hydra +import pandas as pd import torch -from omegaconf import DictConfig, open_dict +from omegaconf import DictConfig -from grass_mil._paths import setup_project_root - -setup_project_root() - -from grass_mil.data.components.datasets import SpatialOmicsGraphDataset # noqa: E402 -from grass_mil.inference.cell_cache import extract_cell_cache # noqa: E402 -from grass_mil.inference.model_loading import load_trained_module # noqa: E402 -from grass_mil.models.training.probe import probe_sampler_config # noqa: E402 -from grass_mil.utils import RankedLogger, extras # noqa: E402 - -log = RankedLogger(__name__, rank_zero_only=True) +from grass_mil.data.splits import read_ids +from grass_mil.inference.cell_cache import extract_cells, sampler_of_checkpoint +from grass_mil.models import SupervisedModule +from grass_mil.runtime import resolve_device, start_run __all__ = ["run_extract", "resolve_models", "main"] +log = logging.getLogger(__name__) -def resolve_models(cfg: DictConfig) -> List[Dict[str, Any]]: - """[{name, ckpt_path, held_out}] from extract.models or a nested-CV summary.""" - models: List[Dict[str, Any]] = [] - summary_path = cfg.extract.get("nested_cv_summary") - if summary_path: - summary = json.loads(Path(str(summary_path)).read_text()) - oof = Path(str(summary_path)).parent / "oof_predictions.csv" - held: Dict[int, List[str]] = {} - if oof.exists(): - import pandas as pd - df = pd.read_csv(oof) - for outer, grp in df.groupby("outer"): - held[int(outer)] = sorted(grp["region_id"].astype(str).unique()) +def resolve_models(ecfg: DictConfig) -> list[dict[str, Any]]: + """``[{name, ckpt_path, held_out}]`` from a cross-validation summary and/or explicit models.""" + models = [] + if ecfg.get("nested_cv_summary"): + summary_path = Path(str(ecfg.nested_cv_summary)) + summary = json.loads(summary_path.read_text()) + oof = pd.read_csv(summary_path.parent / "oof_predictions.csv") for fold in summary["folds"]: - run_dir = Path(str(fold["run_dir"])) - ckpt = run_dir / "checkpoints" / "last.ckpt" - if not ckpt.exists(): - candidates = sorted((run_dir / "checkpoints").glob("*.ckpt")) - if not candidates: - raise FileNotFoundError(f"No checkpoint under {run_dir / 'checkpoints'}") - ckpt = candidates[-1] - models.append( - { - "name": f"fold{int(fold['outer'])}", - "ckpt_path": str(ckpt), - "held_out": held.get(int(fold["outer"]), []), - } - ) - held_cfg = dict(cfg.extract.get("held_out") or {}) - for name, path in dict(cfg.extract.get("models") or {}).items(): - held = held_cfg.get(str(name)) or [] - models.append( - {"name": str(name), "ckpt_path": str(path), "held_out": [str(r) for r in held]} - ) + ckpt = Path(fold["run_dir"]) / "checkpoints" / "last.ckpt" + held_out = oof.loc[oof["outer"] == fold["outer"], "unit_id"].astype(str).tolist() + models.append({"name": f"fold{fold['outer']}", "ckpt_path": str(ckpt), "held_out": held_out}) + for name, ckpt in dict(ecfg.get("models") or {}).items(): + models.append({"name": str(name), "ckpt_path": str(ckpt), "held_out": []}) if not models: - raise ValueError("Set extract.nested_cv_summary or extract.models={name: ckpt}.") + raise ValueError("Set extract.nested_cv_summary or extract.models={name: checkpoint}.") return models -def _regions_filter(value) -> Optional[List[str]]: - if value is None: - return None - if isinstance(value, str): - p = Path(value) - if p.suffix.lower() == ".json": - return [str(v) for v in json.loads(p.read_text())] - return [line.strip() for line in p.read_text().splitlines() if line.strip()] - return [str(v) for v in value] - - -def run_extract(cfg: DictConfig) -> Dict[str, Any]: +def run_extract(cfg: DictConfig) -> list[str]: + start_run(cfg) ecfg = cfg.extract - device = str(ecfg.get("device", "auto")).lower() - device = ( - torch.device("cuda" if torch.cuda.is_available() else "cpu") - if device == "auto" - else torch.device(device) - ) - - # All regions of the processed cohort, regardless of persisted splits. - with open_dict(cfg): - cfg.data.split.explicit = {"train": None, "val": None, "test": None, "id_kind": "region"} - datamodule = hydra.utils.instantiate(cfg.data) + device = resolve_device(ecfg.device) + datamodule = hydra.utils.instantiate(cfg.data, _recursive_=False) datamodule.prepare_data() - datamodule.setup("fit") - full = SpatialOmicsGraphDataset(Path(str(cfg.data.processed_dir)) / "processed_index.json") - datamodule.dataset_train = full # width inference for the model build - sampler_cfg = probe_sampler_config( - datamodule.sampler_config, - roots_per_unit=None, - seed=0, - batch_size=int(ecfg.get("batch_size", 512)), - ) - ct_names = None - label_map = full.label_maps.get("cell_type") if full.label_maps else None - if label_map: - ct_names = [name for name, _ in sorted(label_map.items(), key=lambda kv: kv[1])] - out_dir = Path(str(ecfg.output_dir)) + units = datamodule.all_units() + meta = datamodule.unit_table(units) + unit_meta = {} + for entry in units.entries: + info = {k: str(v) for k, v in meta.loc[entry.unit_id].items() if k == "patient_id"} + label_id = units.label_id(entry) + if label_id is not None: + info["targets"] = units.labels.target(label_id).tolist() + unit_meta[entry.unit_id] = info written = [] - for spec in resolve_models(cfg): - log.info(f"[{spec['name']}] loading {spec['ckpt_path']}") - model = load_trained_module( - cfg, ckpt_path=spec["ckpt_path"], datamodule=datamodule, device=device - ) - model_dir = extract_cell_cache( + for spec in resolve_models(ecfg): + log.info(f"[{spec['name']}] {spec['ckpt_path']}") + ckpt = torch.load(spec["ckpt_path"], map_location="cpu", weights_only=False) + model = SupervisedModule.from_checkpoint(spec["ckpt_path"]) + if tuple(model.input_spec.cell_types) != tuple(datamodule.input_spec.cell_types): + raise ValueError(f"{spec['ckpt_path']} was trained on other cell types than this cohort.") + model_dir = extract_cells( model, - full, - out_dir=out_dir, + units, + sampler=sampler_of_checkpoint(ckpt, int(ecfg.batch_size)), + out_dir=str(ecfg.output_dir), model_name=spec["name"], - sampler_config=sampler_cfg, ckpt_path=spec["ckpt_path"], - held_out_regions=spec.get("held_out"), - regions=_regions_filter(ecfg.get("regions")), - cell_type_names=ct_names, - autocast=bool(ecfg.get("autocast", True)), + held_out=spec["held_out"], + unit_meta=unit_meta, + units=read_ids(ecfg.get("units")), device=device, - resume=bool(ecfg.get("resume", True)), - unit_id=str(ecfg.get("unit_id", "auto")), - log=log, + resume=bool(ecfg.resume), ) written.append(str(model_dir)) - log.info(f"[{spec['name']}] cache written to {model_dir}") - del model - if device.type == "cuda": - torch.cuda.empty_cache() - return {"output_dir": str(out_dir), "models": written} + log.info(f"[{spec['name']}] cells written to {model_dir}") + return written @hydra.main(version_base="1.3", config_path="configs", config_name="extract.yaml") def main(cfg: DictConfig) -> None: - extras(cfg) run_extract(cfg) diff --git a/src/grass_mil/inference/__init__.py b/src/grass_mil/inference/__init__.py index d7df6ee..98c832c 100644 --- a/src/grass_mil/inference/__init__.py +++ b/src/grass_mil/inference/__init__.py @@ -1,17 +1 @@ -from grass_mil.inference.aggregation import aggregate_group_logits -from grass_mil.inference.cell_cache import CellCache, extract_cell_cache -from grass_mil.inference.collectors import collect_inference_payload -from grass_mil.inference.io import embeddings_to_dataframe, node_embeddings_to_dataframe -from grass_mil.inference.metrics import compute_task_metrics -from grass_mil.inference.model_loading import load_trained_module - -__all__ = [ - "CellCache", - "aggregate_group_logits", - "collect_inference_payload", - "compute_task_metrics", - "embeddings_to_dataframe", - "extract_cell_cache", - "load_trained_module", - "node_embeddings_to_dataframe", -] +"""Scoring trained models: unit predictions and per-cell outputs.""" diff --git a/src/grass_mil/inference/aggregation.py b/src/grass_mil/inference/aggregation.py deleted file mode 100644 index a8d488d..0000000 --- a/src/grass_mil/inference/aggregation.py +++ /dev/null @@ -1,427 +0,0 @@ -from __future__ import annotations - -import warnings -from typing import Dict, List, Optional - -import torch - -from grass_mil.inference.schemas import AggregatedPredictionPayload, BatchPredictionPayload -from grass_mil.contracts import validate_attention_width - - -def _aggregate_rows(values: torch.Tensor, mode: str) -> torch.Tensor: - if mode == "mean_probability": - # Mean of the per-row class probabilities, returned as log-probabilities - # so that softmax of the result is exactly the mean probability (patient - # score = mean of core probabilities, the preprint's default). - probs = torch.softmax(values.float(), dim=-1).mean(dim=0, keepdim=True) - return torch.log(probs.clamp_min(1e-12)) - if mode == "mean": - return values.mean(dim=0, keepdim=True) - if mode == "max": - return values.max(dim=0, keepdim=True).values - raise ValueError(f"Unsupported aggregation mode '{mode}'.") - - -def _group_indices(ids: List[str]) -> Dict[str, List[int]]: - grouped: Dict[str, List[int]] = {} - for idx, bag_id in enumerate(ids): - grouped.setdefault(str(bag_id), []).append(idx) - return grouped - - -def _is_missing_group_value(value: object) -> bool: - if value is None: - return True - if isinstance(value, str): - text = value.strip().lower() - return text in {"", "nan", "none", "null"} - if isinstance(value, float): - return value != value - return False - - -def _normalize_group_ids( - ids: Optional[List[Optional[str]]], *, expected_length: int, name: str -) -> List[str]: - if ids is None: - raise ValueError( - f"Aggregation requires '{name}' to regroup predictions at the requested bag scope." - ) - if len(ids) != expected_length: - raise ValueError( - f"Mismatch between rows and '{name}': {expected_length} rows vs {len(ids)} ids." - ) - normalized: List[str] = [] - for idx, value in enumerate(ids): - if _is_missing_group_value(value): - raise ValueError( - f"Missing '{name}' at row {idx}; cannot regroup predictions for this scope." - ) - normalized.append(str(value)) - return normalized - - -def _normalize_group_ids_with_fallback( - primary_ids: Optional[List[Optional[str]]], - fallback_ids: Optional[List[Optional[str]]], - *, - expected_length: int, - primary_name: str, - fallback_name: str, -) -> List[str]: - if primary_ids is None and fallback_ids is None: - raise ValueError( - f"Aggregation requires '{primary_name}' or fallback '{fallback_name}' " - "to regroup predictions at the requested bag scope." - ) - if primary_ids is not None and len(primary_ids) != expected_length: - raise ValueError( - f"Mismatch between rows and '{primary_name}': {expected_length} rows vs " - f"{len(primary_ids)} ids." - ) - if fallback_ids is not None and len(fallback_ids) != expected_length: - raise ValueError( - f"Mismatch between rows and '{fallback_name}': {expected_length} rows vs " - f"{len(fallback_ids)} ids." - ) - - resolved: List[str] = [] - fallback_count = 0 - for idx in range(expected_length): - primary_value = primary_ids[idx] if primary_ids is not None else None - fallback_value = fallback_ids[idx] if fallback_ids is not None else None - use_fallback = _is_missing_group_value(primary_value) - if use_fallback and _is_missing_group_value(fallback_value): - raise ValueError( - f"Missing both '{primary_name}' and fallback '{fallback_name}' at row {idx}; " - "cannot regroup predictions for this scope." - ) - if use_fallback: - fallback_count += 1 - resolved.append(str(fallback_value if use_fallback else primary_value)) - if fallback_count > 0: - warnings.warn( - f"Fell back from '{primary_name}' to '{fallback_name}' for {fallback_count} " - f"of {expected_length} rows while regrouping predictions.", - RuntimeWarning, - stacklevel=3, - ) - return resolved - - -def _map_to_patient(ids: List[str], group_map: Optional[Dict[str, str]]) -> List[str]: - if group_map is None: - raise ValueError( - "bag_scope='patient' needs a region -> patient mapping (group_map); set " - "aggregation.patient_map in the predict config." - ) - out: List[str] = [] - missing: set[str] = set() - for rid in ids: - pid = group_map.get(str(rid)) - if pid is None: - missing.add(str(rid)) - pid = str(rid) - out.append(str(pid)) - if missing: - warnings.warn( - f"{len(missing)} region(s) have no patient in the mapping and were kept as their " - f"own group (first: {sorted(missing)[:5]}).", - UserWarning, - stacklevel=3, - ) - return out - - -def _row_group_ids( - payload: BatchPredictionPayload, - *, - bag_scope: str, - group_map: Optional[Dict[str, str]] = None, -) -> List[str]: - if bag_scope == "patient": - return _map_to_patient(_row_group_ids(payload, bag_scope="region"), group_map) - if bag_scope == "patch": - return [str(v) for v in payload.bag_ids] - if bag_scope == "region": - return _normalize_group_ids_with_fallback( - payload.row_region_ids, - payload.row_sample_ids, - expected_length=len(payload.bag_ids), - primary_name="row_region_ids", - fallback_name="row_sample_ids", - ) - if bag_scope == "sample": - return _normalize_group_ids( - payload.row_sample_ids, - expected_length=len(payload.bag_ids), - name="row_sample_ids", - ) - raise ValueError( - f"Unsupported bag_scope '{bag_scope}'. Expected one of " - "['patch', 'region', 'sample', 'patient']." - ) - - -def _instance_group_ids( - payload: BatchPredictionPayload, - *, - bag_scope: str, - group_map: Optional[Dict[str, str]] = None, -) -> List[str]: - if payload.instance_logits is None: - raise ValueError( - "attention_weighted aggregation requires instance_logits in prediction payload." - ) - n_instances = int(payload.instance_logits.shape[0]) - if bag_scope == "patient": - # Merged bag: every instance of every core of the patient in ONE bag, so - # the attention softmax runs jointly over all of them. - return _map_to_patient(_instance_group_ids(payload, bag_scope="region"), group_map) - if bag_scope == "patch": - patch_ids = payload.instance_patch_ids - if patch_ids is None: - raise ValueError( - "attention_weighted aggregation with bag_scope='patch' requires " - "instance_patch_ids." - ) - if len(patch_ids) != n_instances: - raise ValueError( - "Mismatch between instance_logits rows and instance_patch_ids: " - f"{n_instances} vs {len(patch_ids)}." - ) - return [str(v) for v in patch_ids] - if bag_scope == "region": - return _normalize_group_ids_with_fallback( - payload.instance_region_ids, - payload.instance_sample_ids, - expected_length=n_instances, - primary_name="instance_region_ids", - fallback_name="instance_sample_ids", - ) - if bag_scope == "sample": - return _normalize_group_ids( - payload.instance_sample_ids, - expected_length=n_instances, - name="instance_sample_ids", - ) - raise ValueError( - f"Unsupported bag_scope '{bag_scope}'. Expected one of " - "['patch', 'region', 'sample', 'patient']." - ) - - -def _group_row_indices( - payload: BatchPredictionPayload, - *, - bag_scope: str, - group_map: Optional[Dict[str, str]] = None, -) -> Dict[str, List[int]]: - row_group_ids = _row_group_ids(payload, bag_scope=bag_scope, group_map=group_map) - return _group_indices(row_group_ids) - - -def _sample_indices( - indices: List[int], *, subsample_fraction: float, generator: Optional[torch.Generator] -) -> List[int]: - if subsample_fraction == 1.0 or len(indices) <= 1: - return indices - n_keep = int(len(indices) * subsample_fraction) - if n_keep <= 0: - n_keep = 1 - n_keep = min(n_keep, len(indices)) - if n_keep == len(indices): - return indices - perm = torch.randperm(len(indices), generator=generator).tolist() - chosen = sorted(perm[:n_keep]) - return [indices[i] for i in chosen] - - -def _validate_subsample_fraction(subsample_fraction: float) -> None: - if subsample_fraction == 0.0: - raise ValueError("aggregation.subsample_fraction=0.0 selects no instances and is invalid.") - if not (0.0 <= subsample_fraction <= 1.0): - raise ValueError( - "aggregation.subsample_fraction must be within [0, 1]. " f"Got {subsample_fraction}." - ) - - -def aggregate_group_logits( - payload: BatchPredictionPayload, - *, - mode: str, - bag_scope: str = "patch", - subsample_fraction: float = 1.0, - subsample_seed: Optional[int] = None, - group_map: Optional[Dict[str, str]] = None, -) -> AggregatedPredictionPayload: - """Aggregate duplicated bag rows into one logit row per bag. - - ``bag_scope='patient'`` groups regions through ``group_map`` (region id -> - patient id). With ``mode='mean_probability'`` the patient score is the mean - of its regions' class probabilities; with ``mode='attention_weighted'`` all - instances of the patient's regions form one merged bag with a joint - attention softmax. - """ - _validate_subsample_fraction(subsample_fraction) - if mode == "none": - return AggregatedPredictionPayload( - bag_ids=payload.bag_ids, - bag_logits=payload.bag_logits, - bag_targets=payload.bag_targets, - bag_attention=payload.bag_attention, - metadata={"mode": "none", "bag_scope": bag_scope}, - row_region_ids=payload.row_region_ids, - row_sample_ids=payload.row_sample_ids, - ) - - aggregated_ids: List[str] = [] - aggregated_logits: List[torch.Tensor] = [] - aggregated_targets: List[torch.Tensor] = [] - aggregated_attention_rows: List[Optional[torch.Tensor]] = [] - aggregated_row_region_ids: List[Optional[str]] = [] - aggregated_row_sample_ids: List[Optional[str]] = [] - used_attention = mode == "attention_weighted" - generator = None - if subsample_seed is not None: - generator = torch.Generator() - generator.manual_seed(int(subsample_seed)) - if payload.bag_attention is not None and len(payload.bag_attention) != len(payload.bag_ids): - raise ValueError( - "Mismatch between bag_ids and attention rows for aggregation: " - f"{len(payload.bag_ids)} ids vs {len(payload.bag_attention)} attention rows." - ) - grouped_row_indices = _group_row_indices(payload, bag_scope=bag_scope, group_map=group_map) - - if mode == "attention_weighted": - if payload.instance_attention_logits is None: - raise ValueError( - "attention_weighted aggregation requires instance_attention_logits " - "in prediction payload." - ) - instance_group_ids = _instance_group_ids(payload, bag_scope=bag_scope, group_map=group_map) - grouped_instance_indices = _group_indices(instance_group_ids) - if payload.instance_attention_logits.shape[0] != payload.instance_logits.shape[0]: - raise ValueError( - "Mismatch between instance_attention_logits and instance_logits rows: " - f"{payload.instance_attention_logits.shape[0]} vs " - f"{payload.instance_logits.shape[0]}." - ) - for bag_id in sorted(grouped_instance_indices): - indices = _sample_indices( - grouped_instance_indices[bag_id], - subsample_fraction=subsample_fraction, - generator=generator, - ) - idx = torch.tensor(indices, dtype=torch.long) - instance_logits = payload.instance_logits.index_select(0, idx) - attn_logits = payload.instance_attention_logits.index_select(0, idx) - if attn_logits.dim() == 1: - attn_logits = attn_logits.unsqueeze(-1) - validate_attention_width( - attention_width=int(attn_logits.size(-1)), - num_classes=int(instance_logits.size(-1)), - ) - # Softmax down the instance axis, per class column, matching the - # training-time bag pooling. - attn_scores = torch.softmax(attn_logits, dim=0) - agg_logits = (instance_logits * attn_scores).sum(dim=0, keepdim=True) - aggregated_ids.append(bag_id) - aggregated_logits.append(agg_logits) - aggregated_attention_rows.append(attn_scores.detach().cpu()) - - row_indices = grouped_row_indices.get(bag_id, []) - if row_indices: - aggregated_row_region_ids.append( - payload.row_region_ids[row_indices[0]] - if payload.row_region_ids is not None - else None - ) - aggregated_row_sample_ids.append( - payload.row_sample_ids[row_indices[0]] - if payload.row_sample_ids is not None - else None - ) - if payload.bag_targets is not None: - target_idx = torch.tensor(row_indices, dtype=torch.long) - bag_targets = payload.bag_targets.index_select(0, target_idx) - aggregated_targets.append(bag_targets[:1]) - else: - aggregated_row_region_ids.append(bag_id if bag_scope == "region" else None) - aggregated_row_sample_ids.append(bag_id if bag_scope == "sample" else None) - elif payload.instance_logits is not None and mode != "mean_probability": - instance_group_ids = _instance_group_ids(payload, bag_scope=bag_scope, group_map=group_map) - grouped_instance_indices = _group_indices(instance_group_ids) - for bag_id in sorted(grouped_instance_indices): - indices = grouped_instance_indices[bag_id] - idx = torch.tensor(indices, dtype=torch.long) - instance_logits = payload.instance_logits.index_select(0, idx) - agg_logits = _aggregate_rows(instance_logits, mode=mode) - aggregated_ids.append(bag_id) - aggregated_logits.append(agg_logits) - - row_indices = grouped_row_indices.get(bag_id, []) - if row_indices: - aggregated_row_region_ids.append( - payload.row_region_ids[row_indices[0]] - if payload.row_region_ids is not None - else None - ) - aggregated_row_sample_ids.append( - payload.row_sample_ids[row_indices[0]] - if payload.row_sample_ids is not None - else None - ) - if payload.bag_targets is not None: - target_idx = torch.tensor(row_indices, dtype=torch.long) - bag_targets = payload.bag_targets.index_select(0, target_idx) - aggregated_targets.append(bag_targets[:1]) - else: - aggregated_row_region_ids.append(bag_id if bag_scope == "region" else None) - aggregated_row_sample_ids.append(bag_id if bag_scope == "sample" else None) - else: - # Row-level aggregation (mean_probability always works on the per-region - # rows: a patient's score is the mean of its regions' probabilities). - row_group_ids = _row_group_ids(payload, bag_scope=bag_scope, group_map=group_map) - grouped_row_indices = _group_indices(row_group_ids) - for bag_id in sorted(grouped_row_indices): - indices = grouped_row_indices[bag_id] - idx = torch.tensor(indices, dtype=torch.long) - bag_logits = payload.bag_logits.index_select(0, idx) - agg_logits = _aggregate_rows(bag_logits, mode=mode) - aggregated_ids.append(bag_id) - aggregated_logits.append(agg_logits) - aggregated_row_region_ids.append( - payload.row_region_ids[indices[0]] if payload.row_region_ids is not None else None - ) - aggregated_row_sample_ids.append( - payload.row_sample_ids[indices[0]] if payload.row_sample_ids is not None else None - ) - if payload.bag_targets is not None: - bag_targets = payload.bag_targets.index_select(0, idx) - aggregated_targets.append(bag_targets[:1]) - - target_tensor = None - if aggregated_targets and len(aggregated_targets) == len(aggregated_ids): - target_tensor = torch.cat(aggregated_targets, dim=0) - aggregated_attention: Optional[List[Optional[torch.Tensor]]] = None - if aggregated_attention_rows: - aggregated_attention = aggregated_attention_rows - - return AggregatedPredictionPayload( - bag_ids=aggregated_ids, - bag_logits=torch.cat(aggregated_logits, dim=0), - bag_targets=target_tensor, - bag_attention=aggregated_attention, - metadata={ - "mode": mode, - "used_attention": used_attention, - "bag_attention_retained": used_attention, - "bag_scope": bag_scope, - "subsample_fraction": float(subsample_fraction), - "subsample_seed": subsample_seed, - }, - row_region_ids=aggregated_row_region_ids, - row_sample_ids=aggregated_row_sample_ids, - ) diff --git a/src/grass_mil/inference/cell_cache.py b/src/grass_mil/inference/cell_cache.py index 6eea7a6..521cff0 100644 --- a/src/grass_mil/inference/cell_cache.py +++ b/src/grass_mil/inference/cell_cache.py @@ -1,266 +1,187 @@ -"""Per-cell cache of a trained GRASS-MIL model: the input of niche discovery. +"""Per-cell outputs of a trained model: the input of niche discovery. -For every region and every cell (the ego-graph rooted at that cell) the cache -stores the encoder embedding ``E`` (fp16), the per-class instance logits -``LG``, the raw per-class attention logits ``AR`` and, once per region, the -head bias, cell types, positions, the cell graph and the region label. One -directory per model:: +For every unit and every cell (the ego-graph rooted at it) the cache stores + +* ``embedding``: the encoder embedding (fp16), +* ``attention_space``: the attention head's low-dimensional pre-tanh activation, +* ``instance_logits`` and ``attention_logits``: per-class logits and raw + attention scores, + +and once per unit the cell types, positions, cell ids and the cell graph:: //index.json - //.npz + //.npz -``index.json`` records the model checkpoint, the region list with cell counts, -the head bias, the cell-type vocabulary and, when known, which regions were -held out for that model (nested CV). ``CellCache`` reads it back. +``index.json`` records the checkpoint, the units with their labels and +patients, the units held out from the model's training (cross-validation), the +head bias and the cell-type names. """ from __future__ import annotations import json -import time +import logging from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Dict, Iterable, List, Optional, Sequence +from typing import Any, Iterator, Sequence import numpy as np import torch -__all__ = ["CellCache", "extract_cell_cache", "region_npz_path"] +from grass_mil.data.samplers import EgoGraphLoader, SamplerConfig + +__all__ = ["CellCache", "extract_cells", "sampler_of_checkpoint"] -_INDEX = "index.json" +log = logging.getLogger(__name__) +INDEX_NAME = "index.json" +ARRAYS = ("embedding", "attention_space", "instance_logits", "attention_logits") -def region_npz_path(model_dir: Path, region_id: str) -> Path: - safe = "".join(c if c.isalnum() or c in "._-" else "_" for c in str(region_id)) - return model_dir / f"{safe}.npz" + +def _npz_path(model_dir: Path, unit_id: str) -> Path: + return model_dir / ("".join(c if c.isalnum() or c in "._-" else "_" for c in unit_id) + ".npz") @dataclass class CellCache: - """Reader for one model's cache directory.""" + """Reader of one model's cache directory.""" model_dir: Path - index: Dict[str, Any] = field(default_factory=dict) + index: dict[str, Any] = field(default_factory=dict) @classmethod - def open(cls, model_dir: str | Path) -> "CellCache": + def open(cls, model_dir: str | Path) -> CellCache: model_dir = Path(model_dir) - index_path = model_dir / _INDEX - if not index_path.exists(): - raise FileNotFoundError(f"No cell cache index at {index_path}") - return cls(model_dir=model_dir, index=json.loads(index_path.read_text())) + path = model_dir / INDEX_NAME + if not path.exists(): + raise FileNotFoundError(f"No cell cache at {model_dir} (missing {INDEX_NAME}).") + return cls(model_dir=model_dir, index=json.loads(path.read_text())) @property - def regions(self) -> List[str]: - return [r["region_id"] for r in self.index["regions"]] + def name(self) -> str: + return str(self.index["model_name"]) @property - def held_out(self) -> List[str]: - return list(self.index.get("held_out_regions") or []) + def units(self) -> list[str]: + return [u["unit_id"] for u in self.index["units"]] @property - def bias(self) -> np.ndarray: - return np.asarray(self.index["head_bias"], dtype=np.float64) + def held_out(self) -> list[str]: + return list(self.index.get("held_out_units") or []) + + @property + def cell_types(self) -> list[str]: + return list(self.index["cell_types"]) @property - def cell_type_names(self) -> Optional[List[str]]: - names = self.index.get("cell_type_names") - return list(names) if names else None - - def label_of(self, region_id: str) -> Optional[float]: - for r in self.index["regions"]: - if r["region_id"] == region_id: - return r.get("label") - return None - - def n_cells(self, region_id: str) -> int: - for r in self.index["regions"]: - if r["region_id"] == region_id: - return int(r["n_cells"]) - raise KeyError(region_id) - - def load(self, region_id: str, keys: Optional[Iterable[str]] = None) -> Dict[str, np.ndarray]: - path = region_npz_path(self.model_dir, region_id) - with np.load(path, allow_pickle=False) as z: - if keys is None: - return {k: z[k] for k in z.files} - return {k: z[k] for k in keys if k in z.files} - - def iter_regions(self, regions: Optional[Sequence[str]] = None, keys=None): - for rid in regions if regions is not None else self.regions: - yield rid, self.load(rid, keys) - - -def _cell_type_codes(data) -> Optional[np.ndarray]: - codes = getattr(data, "categorical_codes", None) - if codes is None: - return None - codes = codes.detach().cpu() - if codes.dim() == 2: - codes = codes[:, 0] - return codes.numpy().astype(np.int32) - - -def _roots_of(batch) -> np.ndarray: - roots = getattr(batch, "root_n_id", None) - if roots is None: - raise ValueError("ShaDow batch without root_n_id; cannot align cells.") - roots = roots.detach().cpu().long().view(-1) - n_id = getattr(batch, "n_id", None) - flag = getattr(batch, "root_n_id_is_global", None) - is_global = False - if flag is not None: - f = torch.as_tensor(flag).view(-1) - is_global = bool(f[0]) if f.numel() else False - if not is_global and n_id is not None: - roots = n_id.detach().cpu().long().view(-1)[roots] - return roots.numpy() + def head_bias(self) -> np.ndarray: + return np.asarray(self.index["head_bias"], dtype=np.float64) + + def unit_info(self, unit_id: str) -> dict[str, Any]: + for u in self.index["units"]: + if u["unit_id"] == unit_id: + return u + raise KeyError(unit_id) + + def load(self, unit_id: str) -> dict[str, np.ndarray]: + with np.load(_npz_path(self.model_dir, unit_id), allow_pickle=False) as z: + return {k: z[k] for k in z.files} + + def __iter__(self) -> Iterator[tuple[str, dict[str, np.ndarray]]]: + for unit in self.units: + yield unit, self.load(unit) + + +def sampler_of_checkpoint(ckpt: dict[str, Any], batch_size: int) -> SamplerConfig: + """Ego-graph settings the model was trained with, scoring every cell.""" + cfg = dict(ckpt["datamodule_hyper_parameters"]["sampler"]) + cfg.update(roots_per_unit=None, batch_size=int(batch_size), merge_units=False, root_weighting=None) + cfg["interior_roots"] = {**dict(cfg.get("interior_roots") or {}), "enabled": False} + return SamplerConfig.from_dict(cfg) + + +class _OneUnit: + def __init__(self, data, entry) -> None: + self.data, self.entries = data, [entry] + + def __len__(self) -> int: + return 1 + + def __getitem__(self, idx: int): + return self.data @torch.no_grad() -def extract_cell_cache( +def extract_cells( model, dataset, *, + sampler: SamplerConfig, out_dir: str | Path, model_name: str, - sampler_config, - ckpt_path: Optional[str] = None, - held_out_regions: Optional[Sequence[str]] = None, - regions: Optional[Sequence[str]] = None, - cell_type_names: Optional[Sequence[str]] = None, - autocast: bool = True, - device: Optional[torch.device | str] = None, + ckpt_path: str, + held_out: Sequence[str] = (), + unit_meta: dict[str, dict[str, Any]] | None = None, + units: Sequence[str] | None = None, + device: torch.device, resume: bool = True, - unit_id: str = "auto", - log=None, ) -> Path: - """Run the model over every cell of every unit (region, sample or patch) and write the cache. - - ``unit_id='auto'`` addresses units by patch id when the dataset holds several - tiles per region and by region id otherwise. - """ - from grass_mil.data.components.samplers import get_sampler_strategy - from grass_mil.data.spatial_omics_datamodule import unit_key - - kind = str(unit_id).lower() - if kind == "auto": - regions_seen = {} - for e in dataset.entries: - regions_seen.setdefault(unit_key(e, "region"), set()).add(str(e.patch_id)) - kind = "patch" if any(len(v) > 1 for v in regions_seen.values()) else "region" - + """Score every cell of every unit of ``dataset`` and write the cache.""" model_dir = Path(out_dir) / model_name model_dir.mkdir(parents=True, exist_ok=True) - device = torch.device(device) if device is not None else next(model.parameters()).device - strategy = get_sampler_strategy(sampler_config) - head_bias = _head_bias(model) - wanted = set(str(r) for r in regions) if regions is not None else None - index_regions: List[Dict[str, Any]] = [] - use_autocast = bool(autocast) and device.type == "cuda" - t0 = time.time() - for idx in range(len(dataset)): - entry = dataset.entries[idx] - rid = unit_key(entry, kind) - if wanted is not None and rid not in wanted: + model.to(device).eval() + wanted = set(units) if units is not None else None + index_units = [] + autocast = device.type == "cuda" + for i, entry in enumerate(dataset.entries): + if wanted is not None and entry.unit_id not in wanted: continue - npz = region_npz_path(model_dir, rid) - data = dataset[idx] - n = int(data.num_nodes) - label = None - if getattr(data, "graph_y", None) is not None: - label = float(data.graph_y.view(-1)[0]) - if resume and npz.exists(): - index_regions.append({"region_id": rid, "n_cells": n, "label": label}) + info = {"unit_id": entry.unit_id, "region_id": entry.region_id, "sample_id": entry.sample_id, "n_cells": entry.n_cells} + info.update((unit_meta or {}).get(entry.unit_id, {})) + index_units.append(info) + path = _npz_path(model_dir, entry.unit_id) + if resume and path.exists(): continue - E = np.zeros((n, int(model.encoder.output_dim)), dtype=np.float16) - LG = None - AR = None + data = dataset[i] + n = int(data.num_nodes) + arrays: dict[str, np.ndarray] = {} seen = np.zeros(n, dtype=bool) - loader = strategy.build_dataset_loader( - dataset=[data], batch_size=1, num_workers=0, pin_memory=False, shuffle=False - ) - for batch in loader: - roots = _roots_of(batch) + for batch in EgoGraphLoader(_OneUnit(data, entry), sampler, mode="chunks", roots_per_unit=None): + roots = batch.root_cell.numpy() batch = batch.to(device) - with torch.autocast(device.type, dtype=torch.float16, enabled=use_autocast): - logits, emb, _ = model._forward_instances(batch) - attn = model.attention(emb)[0] if model.attention is not None else None - emb = emb.float() - if LG is None: - LG = np.zeros((n, int(logits.shape[1])), dtype=np.float32) - AR = ( - np.zeros((n, int(attn.shape[1])), dtype=np.float32) - if attn is not None - else None - ) - E[roots] = emb.half().cpu().numpy() - LG[roots] = logits.float().cpu().numpy() - if AR is not None: - AR[roots] = attn.float().cpu().numpy() + with torch.autocast(device.type, dtype=torch.float16, enabled=autocast): + emb = model.encoder(batch) + outputs = { + "embedding": emb, + "attention_space": model.attention.space(emb), + "instance_logits": model.head(emb), + "attention_logits": model.attention(emb), + } + for key, value in outputs.items(): + if key not in arrays: + dtype = np.float16 if key == "embedding" else np.float32 + arrays[key] = np.zeros((n, int(value.shape[1])), dtype=dtype) + arrays[key][roots] = value.float().cpu().numpy() seen[roots] = True if not seen.all(): - raise RuntimeError( - f"Region {rid}: {int((~seen).sum())} of {n} cells were not visited; " - "the extraction sampler must keep every root (all-roots mode)." - ) - payload: Dict[str, Any] = { - "E": E, - "LG": LG, - "pos": data.pos.detach().cpu().numpy().astype(np.float32), - "edge_index": data.edge_index.detach().cpu().numpy().astype(np.int64), - } - if AR is not None: - payload["AR"] = AR - ct = _cell_type_codes(data) - if ct is not None: - payload["ct"] = ct - if getattr(data, "edge_attr", None) is not None: - payload["edge_attr"] = data.edge_attr.detach().cpu().numpy().astype(np.float32) - np.savez(npz, **payload) - index_regions.append({"region_id": rid, "n_cells": n, "label": label}) - if log is not None and len(index_regions) % 25 == 0: - log.info( - f"[{model_name}] {len(index_regions)} regions ({(time.time() - t0) / 60:.1f} min)" - ) + raise RuntimeError(f"Unit {entry.unit_id}: {int((~seen).sum())} cells were not scored.") + arrays["pos"] = data.pos.numpy().astype(np.float32) + arrays["edge_index"] = data.edge_index.numpy().astype(np.int64) + arrays["cell_ids"] = np.asarray(dataset.load(i)["cell_ids"], dtype=str) + if getattr(data, "cell_type", None) is not None: + arrays["cell_type"] = data.cell_type.numpy().astype(np.int32) + np.savez(path, **arrays) + if len(index_units) % 25 == 0: + log.info(f"[{model_name}] {len(index_units)} units") index = { "model_name": model_name, - "ckpt_path": str(ckpt_path) if ckpt_path else None, - "regions": index_regions, - "held_out_regions": [str(r) for r in held_out_regions] if held_out_regions else [], - "head_bias": head_bias, - "unit_id": kind, - "cell_type_names": list(cell_type_names) if cell_type_names else None, - "embedding_dim": int(model.encoder.output_dim), - "sampler_runtime": _plain(getattr(sampler_config, "runtime", {}) or {}), + "ckpt_path": str(ckpt_path), + "target_type": model.task["target_type"], + "units": index_units, + "held_out_units": sorted(map(str, held_out)), + "head_bias": model.head.bias.detach().cpu().tolist(), + "cell_types": list(model.input_spec.cell_types), + "sampler": sampler.as_dict(), } - (model_dir / _INDEX).write_text(json.dumps(index, indent=1, default=str)) + (model_dir / INDEX_NAME).write_text(json.dumps(index, indent=1, default=str)) return model_dir - - -def _plain(obj: Any) -> Any: - try: - from omegaconf import DictConfig, ListConfig, OmegaConf - - if isinstance(obj, (DictConfig, ListConfig)): - return OmegaConf.to_container(obj, resolve=True) - except ImportError: # pragma: no cover - pass - if isinstance(obj, dict): - return {str(k): _plain(v) for k, v in obj.items()} - if isinstance(obj, (list, tuple)): - return [_plain(v) for v in obj] - return obj - - -def _head_bias(model) -> List[float]: - head = getattr(model, "graph_head", None) - if head is None: - return [] - last = None - for module in head.modules(): - if isinstance(module, torch.nn.Linear): - last = module - if last is None or last.bias is None: - return [] - return [float(v) for v in last.bias.detach().cpu().view(-1)] diff --git a/src/grass_mil/inference/collectors.py b/src/grass_mil/inference/collectors.py deleted file mode 100644 index 32941dc..0000000 --- a/src/grass_mil/inference/collectors.py +++ /dev/null @@ -1,483 +0,0 @@ -from __future__ import annotations - -from collections import deque -from typing import Any, Dict, Iterable, List, Optional - -import torch -from lightning import LightningDataModule, LightningModule, Trainer - -from grass_mil.inference.schemas import ( - BatchPredictionPayload, - CollectedInferencePayload, - EmbeddingPayload, -) - - -def _as_str_list(values: Iterable[Any]) -> List[str]: - return [str(v) for v in values] - - -def _ensure_2d(values: torch.Tensor) -> torch.Tensor: - if values.ndim == 1: - return values.unsqueeze(-1) - return values - - -def _stack_optional_tensors(values: List[torch.Tensor]) -> Optional[torch.Tensor]: - if not values: - return None - return torch.cat([_ensure_2d(v) for v in values], dim=0) - - -def _stack_optional_tensors_with_padding( - values: List[torch.Tensor], *, fill_value: float = 0.0 -) -> Optional[torch.Tensor]: - if not values: - return None - prepared = [_ensure_2d(v) for v in values] - max_width = max(int(v.shape[1]) for v in prepared) - if max_width <= 0: - return torch.cat(prepared, dim=0) - padded: List[torch.Tensor] = [] - for value in prepared: - if int(value.shape[1]) == max_width: - padded.append(value) - continue - pad = torch.full( - (int(value.shape[0]), max_width - int(value.shape[1])), - fill_value=fill_value, - dtype=value.dtype, - device=value.device, - ) - padded.append(torch.cat([value, pad], dim=1)) - return torch.cat(padded, dim=0) - - -def _as_positive_count_tensor(values: Any, *, expected_length: int) -> torch.Tensor: - if values is None: - raise ValueError("Missing required 'embedding_bag_counts' while collecting embeddings.") - if isinstance(values, torch.Tensor): - counts = values.detach().cpu().reshape(-1).float() - else: - counts = torch.tensor(list(values), dtype=torch.float32).reshape(-1) - if counts.numel() != expected_length: - raise ValueError( - "Mismatch between embedding_bag_ids and embedding_bag_counts rows while " - "collecting embeddings: " - f"{expected_length} ids vs {counts.numel()} counts." - ) - if not torch.all(counts > 0): - raise ValueError("embedding_bag_counts must be strictly positive for aggregation.") - return counts - - -def _align_chunk_attention_to_rows( - chunk_bag_ids: List[str], chunk_attention: Optional[Dict[str, torch.Tensor]] -) -> List[Optional[torch.Tensor]]: - if not chunk_attention: - return [None] * len(chunk_bag_ids) - return [ - chunk_attention[bag_id].detach().cpu() if bag_id in chunk_attention else None - for bag_id in chunk_bag_ids - ] - - -def collect_inference_payload( - *, - trainer: Trainer, - model: LightningModule, - datamodule: LightningDataModule, - ckpt_path: Optional[str], - include_instance_payload: bool = True, - include_instance_embeddings: bool = False, - include_embeddings: bool = False, - include_node_embeddings: bool = False, -) -> CollectedInferencePayload: - """Collect predictions and optional embeddings from a single predict pass.""" - outputs = deque(trainer.predict(model=model, datamodule=datamodule, ckpt_path=ckpt_path)) - - bag_ids: List[str] = [] - logits_chunks: List[torch.Tensor] = [] - target_chunks: List[torch.Tensor] = [] - attention_rows: List[Optional[torch.Tensor]] = [] - row_region_ids: List[Optional[str]] = [] - row_sample_ids: List[Optional[str]] = [] - instance_logits_chunks: List[torch.Tensor] = [] - instance_attention_logits_chunks: List[torch.Tensor] = [] - instance_patch_ids: List[str] = [] - instance_bag_ids: List[str] = [] - instance_region_ids: List[Optional[str]] = [] - instance_sample_ids: List[Optional[str]] = [] - instance_embedding_chunks: List[torch.Tensor] = [] - instance_composition_chunks: List[torch.Tensor] = [] - instance_centroid_chunks: List[torch.Tensor] = [] - instance_graphs: List[Any] = [] - expect_targets: Optional[bool] = None - - bag_emb_sums: Dict[str, torch.Tensor] = {} - bag_emb_weights: Dict[str, float] = {} - node_embs: List[torch.Tensor] = [] - node_bag_ids: List[str] = [] - - while outputs: - chunk = outputs.popleft() - if not isinstance(chunk, dict): - raise TypeError("Expected `predict_step` output to be a dict.") - if "bag_ids" not in chunk or "bag_logits" not in chunk: - raise KeyError("`predict_step` must include bag_ids and bag_logits.") - - chunk_bag_ids = _as_str_list(chunk["bag_ids"]) - chunk_logits = _ensure_2d(chunk["bag_logits"]).detach().cpu() - if len(chunk_bag_ids) != chunk_logits.shape[0]: - raise ValueError( - "Mismatch within predict chunk between bag_ids and bag_logits rows: " - f"{len(chunk_bag_ids)} ids vs {chunk_logits.shape[0]} logits." - ) - - bag_ids.extend(chunk_bag_ids) - logits_chunks.append(chunk_logits) - - has_targets = "bag_targets" in chunk and chunk["bag_targets"] is not None - if expect_targets is None: - expect_targets = has_targets - elif has_targets != expect_targets: - raise ValueError( - "Inconsistent predict_step target payloads across chunks: either all " - "chunks must include bag_targets or none." - ) - if has_targets: - chunk_targets = _ensure_2d(chunk["bag_targets"]).detach().cpu() - if chunk_targets.shape[0] != len(chunk_bag_ids): - raise ValueError( - "Mismatch within predict chunk between bag_ids and bag_targets rows: " - f"{len(chunk_bag_ids)} ids vs {chunk_targets.shape[0]} targets." - ) - target_chunks.append(chunk_targets) - - chunk_row_region_ids = chunk.get("row_region_ids") - if chunk_row_region_ids is None: - row_region_ids.extend([None] * len(chunk_bag_ids)) - else: - if len(chunk_row_region_ids) != len(chunk_bag_ids): - raise ValueError( - "Mismatch between bag_ids and row_region_ids within predict chunk: " - f"{len(chunk_bag_ids)} ids vs {len(chunk_row_region_ids)} row ids." - ) - row_region_ids.extend([None if v is None else str(v) for v in chunk_row_region_ids]) - - chunk_row_sample_ids = chunk.get("row_sample_ids") - if chunk_row_sample_ids is None: - row_sample_ids.extend([None] * len(chunk_bag_ids)) - else: - if len(chunk_row_sample_ids) != len(chunk_bag_ids): - raise ValueError( - "Mismatch between bag_ids and row_sample_ids within predict chunk: " - f"{len(chunk_bag_ids)} ids vs {len(chunk_row_sample_ids)} row ids." - ) - row_sample_ids.extend([None if v is None else str(v) for v in chunk_row_sample_ids]) - - attention_rows.extend( - _align_chunk_attention_to_rows(chunk_bag_ids, chunk.get("bag_attention")) - ) - - if ( - include_instance_payload - and "instance_logits" in chunk - and chunk["instance_logits"] is not None - ): - for field in ("instance_patch_ids", "instance_region_ids", "instance_sample_ids"): - if field not in chunk: - raise KeyError( - f"predict_step provided instance_logits without required field '{field}'." - ) - chunk_instance_logits = _ensure_2d(chunk["instance_logits"]).detach().cpu() - if len(chunk["instance_patch_ids"]) != chunk_instance_logits.shape[0]: - raise ValueError( - "Mismatch between instance_logits and instance_patch_ids rows: " - f"{chunk_instance_logits.shape[0]} vs {len(chunk['instance_patch_ids'])}." - ) - if len(chunk["instance_region_ids"]) != chunk_instance_logits.shape[0]: - raise ValueError( - "Mismatch between instance_logits and instance_region_ids rows: " - f"{chunk_instance_logits.shape[0]} vs {len(chunk['instance_region_ids'])}." - ) - if len(chunk["instance_sample_ids"]) != chunk_instance_logits.shape[0]: - raise ValueError( - "Mismatch between instance_logits and instance_sample_ids rows: " - f"{chunk_instance_logits.shape[0]} vs {len(chunk['instance_sample_ids'])}." - ) - chunk_instance_bag_ids = chunk.get("instance_bag_ids") - if chunk_instance_bag_ids is not None: - if len(chunk_instance_bag_ids) != chunk_instance_logits.shape[0]: - raise ValueError( - "Mismatch between instance_logits and instance_bag_ids rows: " - f"{chunk_instance_logits.shape[0]} vs {len(chunk_instance_bag_ids)}." - ) - instance_bag_ids.extend([str(v) for v in chunk_instance_bag_ids]) - elif include_instance_embeddings: - raise KeyError( - "predict_step must include 'instance_bag_ids' when interpretability " - "instance export is enabled." - ) - instance_logits_chunks.append(chunk_instance_logits) - chunk_instance_attn = chunk.get("instance_attention_logits") - if chunk_instance_attn is not None: - chunk_instance_attn = _ensure_2d(chunk_instance_attn).detach().cpu() - if chunk_instance_logits.shape[0] != chunk_instance_attn.shape[0]: - raise ValueError( - "Mismatch between instance_logits and instance_attention_logits rows: " - f"{chunk_instance_logits.shape[0]} vs {chunk_instance_attn.shape[0]}." - ) - instance_attention_logits_chunks.append(chunk_instance_attn) - instance_patch_ids.extend([str(v) for v in chunk["instance_patch_ids"]]) - instance_region_ids.extend( - [None if v is None else str(v) for v in chunk["instance_region_ids"]] - ) - instance_sample_ids.extend( - [None if v is None else str(v) for v in chunk["instance_sample_ids"]] - ) - chunk_instance_embeddings = chunk.get("instance_embeddings") - if include_instance_embeddings: - if chunk_instance_embeddings is None: - raise KeyError( - "predict_step must include 'instance_embeddings' when " - "interpretability instance export is enabled." - ) - chunk_instance_embeddings = _ensure_2d(chunk_instance_embeddings).detach().cpu() - if chunk_instance_embeddings.shape[0] != chunk_instance_logits.shape[0]: - raise ValueError( - "Mismatch between instance_logits and instance_embeddings rows: " - f"{chunk_instance_logits.shape[0]} vs " - f"{chunk_instance_embeddings.shape[0]}." - ) - instance_embedding_chunks.append(chunk_instance_embeddings) - elif chunk_instance_embeddings is not None: - chunk_instance_embeddings = _ensure_2d(chunk_instance_embeddings).detach().cpu() - if chunk_instance_embeddings.shape[0] != chunk_instance_logits.shape[0]: - raise ValueError( - "Mismatch between instance_logits and instance_embeddings rows: " - f"{chunk_instance_logits.shape[0]} vs " - f"{chunk_instance_embeddings.shape[0]}." - ) - instance_embedding_chunks.append(chunk_instance_embeddings) - chunk_instance_composition = chunk.get("instance_composition") - if chunk_instance_composition is not None: - chunk_instance_composition = _ensure_2d(chunk_instance_composition).detach().cpu() - if chunk_instance_composition.shape[0] != chunk_instance_logits.shape[0]: - raise ValueError( - "Mismatch between instance_logits and instance_composition rows: " - f"{chunk_instance_logits.shape[0]} vs " - f"{chunk_instance_composition.shape[0]}." - ) - instance_composition_chunks.append(chunk_instance_composition) - chunk_instance_centroids = chunk.get("instance_centroids") - if chunk_instance_centroids is not None: - chunk_instance_centroids = _ensure_2d(chunk_instance_centroids).detach().cpu() - if chunk_instance_centroids.shape[0] != chunk_instance_logits.shape[0]: - raise ValueError( - "Mismatch between instance_logits and instance_centroids rows: " - f"{chunk_instance_logits.shape[0]} vs " - f"{chunk_instance_centroids.shape[0]}." - ) - instance_centroid_chunks.append(chunk_instance_centroids) - chunk_instance_graphs = chunk.get("instance_graphs") - if chunk_instance_graphs is not None: - try: - chunk_instance_graphs = list(chunk_instance_graphs) - except TypeError as exc: - raise TypeError("predict_step instance_graphs must be iterable.") from exc - if len(chunk_instance_graphs) != int(chunk_instance_logits.shape[0]): - raise ValueError( - "Mismatch between instance_logits and instance_graphs rows: " - f"{chunk_instance_logits.shape[0]} vs {len(chunk_instance_graphs)}." - ) - instance_graphs.extend(chunk_instance_graphs) - - if include_embeddings: - missing_embed_fields = [ - field - for field in ("embedding_bag_ids", "graph_embeddings", "embedding_bag_counts") - if field not in chunk or chunk[field] is None - ] - if missing_embed_fields: - raise ValueError( - "predict_step is missing required embedding fields while embeddings " - f"collection is enabled: {missing_embed_fields}." - ) - chunk_embedding_bag_ids = _as_str_list(chunk["embedding_bag_ids"]) - chunk_graph_embeddings = _ensure_2d(chunk["graph_embeddings"]).detach().cpu() - if len(chunk_embedding_bag_ids) != chunk_graph_embeddings.shape[0]: - raise ValueError( - "Mismatch between embedding_bag_ids and graph_embeddings rows while " - "collecting embeddings: " - f"{len(chunk_embedding_bag_ids)} ids vs " - f"{chunk_graph_embeddings.shape[0]} rows." - ) - chunk_counts = _as_positive_count_tensor( - chunk["embedding_bag_counts"], expected_length=len(chunk_embedding_bag_ids) - ) - for idx, emb_bag_id in enumerate(chunk_embedding_bag_ids): - weight = float(chunk_counts[idx].item()) - row = chunk_graph_embeddings[idx] - if emb_bag_id not in bag_emb_sums: - bag_emb_sums[emb_bag_id] = row * weight - bag_emb_weights[emb_bag_id] = weight - else: - bag_emb_sums[emb_bag_id] = bag_emb_sums[emb_bag_id] + (row * weight) - bag_emb_weights[emb_bag_id] += weight - - chunk_node_embeddings = chunk.get("node_embeddings") - chunk_node_bag_ids = chunk.get("node_bag_ids") - if include_node_embeddings: - if chunk_node_embeddings is None or chunk_node_bag_ids is None: - raise ValueError( - "predict_step must include both node_embeddings and node_bag_ids " - "when node embedding extraction is enabled." - ) - if chunk_node_embeddings is not None or chunk_node_bag_ids is not None: - if chunk_node_embeddings is None or chunk_node_bag_ids is None: - raise ValueError( - "predict_step returned partial node embedding payload; both " - "node_embeddings and node_bag_ids are required together." - ) - node_tensor = _ensure_2d(chunk_node_embeddings).detach().cpu() - node_id_list = _as_str_list(chunk_node_bag_ids) - if len(node_id_list) != node_tensor.shape[0]: - raise ValueError( - "Mismatch between node_embeddings rows and node_bag_ids while " - "collecting embeddings: " - f"{node_tensor.shape[0]} rows vs {len(node_id_list)} ids." - ) - node_embs.append(node_tensor) - node_bag_ids.extend(node_id_list) - - bag_logits = _stack_optional_tensors(logits_chunks) - if bag_logits is None: - raise ValueError("No logits were collected from predict outputs.") - bag_targets = _stack_optional_tensors(target_chunks) if expect_targets else None - bag_attention = attention_rows if any(attn is not None for attn in attention_rows) else None - instance_logits = _stack_optional_tensors(instance_logits_chunks) - instance_attention_logits = _stack_optional_tensors(instance_attention_logits_chunks) - instance_embeddings = _stack_optional_tensors(instance_embedding_chunks) - instance_composition = _stack_optional_tensors_with_padding(instance_composition_chunks) - instance_centroids = _stack_optional_tensors(instance_centroid_chunks) - - if len(bag_ids) != bag_logits.shape[0]: - raise ValueError( - "Mismatch between collected bag_ids and logits rows: " - f"{len(bag_ids)} ids vs {bag_logits.shape[0]} logits." - ) - if bag_attention is not None and len(bag_attention) != len(bag_ids): - raise ValueError( - "Mismatch between collected bag_ids and attention rows: " - f"{len(bag_ids)} ids vs {len(bag_attention)} attention rows." - ) - if len(row_region_ids) != len(bag_ids): - raise ValueError( - "Mismatch between collected bag_ids and row_region_ids: " - f"{len(bag_ids)} ids vs {len(row_region_ids)} row ids." - ) - if len(row_sample_ids) != len(bag_ids): - raise ValueError( - "Mismatch between collected bag_ids and row_sample_ids: " - f"{len(bag_ids)} ids vs {len(row_sample_ids)} row ids." - ) - if bag_targets is not None and bag_targets.shape[0] != len(bag_ids): - raise ValueError( - "Mismatch between collected bag_ids and bag_targets rows: " - f"{len(bag_ids)} ids vs {bag_targets.shape[0]} targets." - ) - if instance_logits is not None: - if len(instance_patch_ids) != instance_logits.shape[0]: - raise ValueError( - "Mismatch between collected instance_logits and instance_patch_ids: " - f"{instance_logits.shape[0]} vs {len(instance_patch_ids)}." - ) - if len(instance_region_ids) != instance_logits.shape[0]: - raise ValueError( - "Mismatch between collected instance_logits and instance_region_ids: " - f"{instance_logits.shape[0]} vs {len(instance_region_ids)}." - ) - if len(instance_sample_ids) != instance_logits.shape[0]: - raise ValueError( - "Mismatch between collected instance_logits and instance_sample_ids: " - f"{instance_logits.shape[0]} vs {len(instance_sample_ids)}." - ) - if instance_bag_ids and len(instance_bag_ids) != instance_logits.shape[0]: - raise ValueError( - "Mismatch between collected instance_logits and instance_bag_ids: " - f"{instance_logits.shape[0]} vs {len(instance_bag_ids)}." - ) - if ( - instance_embeddings is not None - and instance_embeddings.shape[0] != instance_logits.shape[0] - ): - raise ValueError( - "Mismatch between collected instance_logits and instance_embeddings rows: " - f"{instance_logits.shape[0]} vs {instance_embeddings.shape[0]}." - ) - if ( - instance_composition is not None - and instance_composition.shape[0] != instance_logits.shape[0] - ): - raise ValueError( - "Mismatch between collected instance_logits and instance_composition rows: " - f"{instance_logits.shape[0]} vs {instance_composition.shape[0]}." - ) - if ( - instance_centroids is not None - and instance_centroids.shape[0] != instance_logits.shape[0] - ): - raise ValueError( - "Mismatch between collected instance_logits and instance_centroids rows: " - f"{instance_logits.shape[0]} vs {instance_centroids.shape[0]}." - ) - if instance_graphs and len(instance_graphs) != int(instance_logits.shape[0]): - raise ValueError( - "Mismatch between collected instance_logits and instance_graphs rows: " - f"{instance_logits.shape[0]} vs {len(instance_graphs)}." - ) - if include_instance_embeddings and instance_logits is not None and instance_embeddings is None: - raise ValueError( - "Interpretability export requested instance embeddings, but none were collected." - ) - - prediction_payload = BatchPredictionPayload( - bag_ids=bag_ids, - bag_logits=bag_logits, - bag_targets=bag_targets, - bag_attention=bag_attention, - row_region_ids=row_region_ids, - row_sample_ids=row_sample_ids, - instance_logits=instance_logits, - instance_attention_logits=instance_attention_logits, - instance_patch_ids=instance_patch_ids if instance_logits is not None else None, - instance_bag_ids=instance_bag_ids if instance_bag_ids else None, - instance_region_ids=instance_region_ids if instance_logits is not None else None, - instance_sample_ids=instance_sample_ids if instance_logits is not None else None, - instance_embeddings=instance_embeddings, - instance_composition=instance_composition, - instance_centroids=instance_centroids, - instance_graphs=instance_graphs if instance_graphs else None, - ) - - embedding_payload: Optional[EmbeddingPayload] = None - if include_embeddings: - if not bag_emb_sums: - raise ValueError("No embeddings were collected from predict outputs.") - ordered_bag_ids = sorted(bag_emb_sums.keys()) - graph_embs = torch.stack( - [bag_emb_sums[bag_id] / bag_emb_weights[bag_id] for bag_id in ordered_bag_ids], - dim=0, - ) - embedding_payload = EmbeddingPayload( - bag_ids=ordered_bag_ids, - graph_embeddings=graph_embs, - node_embeddings=torch.cat(node_embs, dim=0) if node_embs else None, - node_bag_ids=node_bag_ids if node_bag_ids else None, - ) - - return CollectedInferencePayload( - prediction_payload=prediction_payload, - embedding_payload=embedding_payload, - ) diff --git a/src/grass_mil/inference/interpretability_export.py b/src/grass_mil/inference/interpretability_export.py deleted file mode 100644 index dc36876..0000000 --- a/src/grass_mil/inference/interpretability_export.py +++ /dev/null @@ -1,638 +0,0 @@ -from __future__ import annotations - -import warnings -from collections import defaultdict -from typing import Any, Dict, List, Optional, Sequence - -import numpy as np -import pandas as pd -import torch - -from grass_mil.inference.schemas import BatchPredictionPayload - - -def _tensor_to_2d_cpu(values: Optional[torch.Tensor]) -> Optional[torch.Tensor]: - if values is None: - return None - values = values.detach().cpu() - if values.is_floating_point(): - values = values.float() - if values.ndim == 1: - values = values.unsqueeze(-1) - return values - - -def _normalize_optional( - values: Optional[Sequence[Optional[str]]], size: int -) -> List[Optional[str]]: - if values is None: - return [None] * size - out = [None if v is None else str(v) for v in values] - if len(out) != size: - raise ValueError(f"Expected {size} values, got {len(out)}.") - return out - - -def _infer_bag_ids(payload: BatchPredictionPayload, size: int) -> List[str]: - if payload.instance_bag_ids is not None: - bag_ids = [str(v) for v in payload.instance_bag_ids] - if len(bag_ids) != size: - raise ValueError( - "Mismatch between instance logits rows and instance_bag_ids: " - f"{size} vs {len(bag_ids)}." - ) - return bag_ids - - region_ids = _normalize_optional(payload.instance_region_ids, size) - sample_ids = _normalize_optional(payload.instance_sample_ids, size) - patch_ids = [str(v) for v in (payload.instance_patch_ids or [])] - bag_ids: List[str] = [] - for idx in range(size): - region = region_ids[idx] - sample = sample_ids[idx] - if sample and region: - bag_ids.append(f"{sample}::{region}") - elif region: - bag_ids.append(region) - elif sample: - bag_ids.append(sample) - elif idx < len(patch_ids): - bag_ids.append(patch_ids[idx]) - else: - bag_ids.append(f"bag_{idx}") - return bag_ids - - -def _make_unique_ids(values: Sequence[str]) -> List[str]: - counts: Dict[str, int] = defaultdict(int) - out: List[str] = [] - for value in values: - base = str(value) if str(value).strip() else "instance" - counts[base] += 1 - if counts[base] == 1: - out.append(base) - else: - out.append(f"{base}::{counts[base] - 1}") - return out - - -def _softmax_by_group(logits: torch.Tensor, groups: Sequence[str]) -> np.ndarray: - """Normalise attention logits over the instances of each bag. - - Returns one column per attention channel: a shared channel yields - ``[n_instances, 1]``, per-class attention yields ``[n_instances, - n_classes]``. Each column sums to 1 within a bag. - """ - if logits.ndim == 1: - logits = logits.unsqueeze(-1) - if logits.ndim != 2: - raise ValueError( - "instance_attention_logits must be a 2D tensor of shape " - f"[n_instances, n_attention_classes]; got {tuple(logits.shape)}." - ) - if logits.shape[0] != len(groups): - raise ValueError( - "Mismatch between attention logits rows and group ids: " - f"{logits.shape[0]} vs {len(groups)}." - ) - attention = np.zeros((len(groups), int(logits.shape[1])), dtype=np.float64) - grouped_indices: Dict[str, List[int]] = defaultdict(list) - for idx, group in enumerate(groups): - grouped_indices[str(group)].append(idx) - for indices in grouped_indices.values(): - chunk = logits[indices, :] - scores = torch.softmax(chunk, dim=0).cpu().numpy().astype(np.float64) - attention[indices, :] = scores - return attention - - -def _uniform_attention(groups: Sequence[str]) -> np.ndarray: - attention = np.zeros((len(groups),), dtype=np.float64) - grouped_indices: Dict[str, List[int]] = defaultdict(list) - for idx, group in enumerate(groups): - grouped_indices[str(group)].append(idx) - for indices in grouped_indices.values(): - if not indices: - continue - value = 1.0 / float(len(indices)) - for global_idx in indices: - attention[global_idx] = value - return attention - - -def _resolve_scores( - logits: torch.Tensor, - *, - score_mode: str, - score_logit_index: int, -) -> List[float]: - if logits.ndim != 2 or int(logits.shape[1]) < 1: - raise ValueError("instance_logits must be a 2D tensor with >=1 column for score export.") - score_idx = int(score_logit_index) - if score_idx < 0 or score_idx >= int(logits.shape[1]): - raise ValueError( - "score_logit_index is out of bounds for instance_logits: " - f"index={score_idx}, width={int(logits.shape[1])}." - ) - - mode = str(score_mode).strip().lower() - if mode in {"sigmoid", "sigmoid_logit", "probability"}: - return torch.sigmoid(logits[:, score_idx]).tolist() - if mode in {"identity", "raw", "logit"}: - return logits[:, score_idx].tolist() - if mode in {"softmax", "class_probability"}: - probs = torch.softmax(logits, dim=1) - return probs[:, score_idx].tolist() - raise ValueError( - "Unsupported score_mode for interpretability export. " - "Expected one of: 'sigmoid', 'identity', 'softmax'." - ) - - -def resolve_composition_column_names( - datamodule: Any, - *, - composition_label: str, - composition_prefix: str, - width: int, -) -> List[str]: - label_map = None - for attr in ("dataset_train", "dataset_val", "dataset_test"): - dataset = getattr(datamodule, attr, None) - maps = getattr(dataset, "label_maps", None) if dataset is not None else None - if ( - isinstance(maps, dict) - and composition_label in maps - and isinstance(maps[composition_label], dict) - ): - label_map = maps[composition_label] - break - if not isinstance(label_map, dict): - return [f"{composition_prefix}{idx}" for idx in range(width)] - - inv = {int(v): str(k) for k, v in label_map.items()} - names: List[str] = [] - for idx in range(width): - label = inv.get(idx, str(idx)) - names.append(f"{composition_prefix}{label}") - return names - - -def _resolve_instance_root_cell_types( - instance_graphs: Sequence[Any], - *, - labels: Sequence[str], - composition_label: str, -) -> Optional[List[str]]: - """Cell type of each instance's root cell, as a readable name. - - `root_n_id` holds a global node id, so the root's row within the subgraph is - found by matching it against `n_id`. Returns None if any instance cannot be - resolved, so a partially-correct column is never emitted. - """ - out: List[str] = [] - for graph in instance_graphs: - codes = getattr(graph, "categorical_codes", None) - n_id = getattr(graph, "n_id", None) - if not isinstance(codes, torch.Tensor) or not isinstance(n_id, torch.Tensor): - return None - try: - root_global = _resolve_instance_root_node_id(graph) - except ValueError: - return None - - matches = (n_id.detach().cpu().long().view(-1) == int(root_global)).nonzero().view(-1) - if matches.numel() == 0: - return None - local_idx = int(matches[0]) - - column = 0 - slices = getattr(graph, "categorical_slices", None) - if isinstance(slices, dict) and composition_label in slices: - column = int(slices[composition_label]) - codes_2d = codes if codes.ndim == 2 else codes.view(-1, 1) - if local_idx >= int(codes_2d.shape[0]) or column >= int(codes_2d.shape[1]): - return None - - code = int(codes_2d[local_idx, column]) - out.append(labels[code] if 0 <= code < len(labels) else str(code)) - return out - - -def build_instance_table( - payload: BatchPredictionPayload, - *, - composition_prefix: str = "comp_", - embedding_prefix: str = "inst_emb", - composition_column_names: Optional[Sequence[str]] = None, - require_composition: bool = True, - id_column: str = "instance_id", - bag_id_column: str = "bag_id", - score_column: str = "score", - score_mode: str = "sigmoid", - score_logit_index: int = 0, - attention_column: str = "attention", - attention_class_index: Optional[int] = None, - cell_type_column: Optional[str] = "cell_type", -) -> pd.DataFrame: - logits = _tensor_to_2d_cpu(payload.instance_logits) - if logits is None: - raise ValueError("Interpretability export requires instance_logits in predict payload.") - - embeddings = _tensor_to_2d_cpu(payload.instance_embeddings) - if embeddings is None: - raise ValueError( - "Interpretability export requires instance_embeddings in predict payload." - ) - if embeddings.shape[0] != logits.shape[0]: - raise ValueError( - "Mismatch between instance_logits and instance_embeddings rows: " - f"{logits.shape[0]} vs {embeddings.shape[0]}." - ) - - size = int(logits.shape[0]) - patch_ids = ( - [str(v) for v in payload.instance_patch_ids] - if payload.instance_patch_ids is not None - else [f"instance_{idx}" for idx in range(size)] - ) - if len(patch_ids) != size: - raise ValueError( - "Mismatch between instance logits rows and instance_patch_ids: " - f"{size} vs {len(patch_ids)}." - ) - bag_ids = _infer_bag_ids(payload, size) - region_ids = _normalize_optional(payload.instance_region_ids, size) - sample_ids = _normalize_optional(payload.instance_sample_ids, size) - instance_ids = _make_unique_ids(patch_ids) - - data: Dict[str, List[object]] = { - id_column: instance_ids, - bag_id_column: bag_ids, - "patch_id": patch_ids, - "region_id": region_ids, - "sample_id": sample_ids, - score_column: _resolve_scores( - logits, - score_mode=score_mode, - score_logit_index=score_logit_index, - ), - } - for col in range(int(logits.shape[1])): - data[f"logit_{col}"] = logits[:, col].tolist() - - attention_logits = _tensor_to_2d_cpu(payload.instance_attention_logits) - if attention_logits is not None: - attention_matrix = _softmax_by_group(attention_logits, bag_ids) - n_channels = int(attention_matrix.shape[1]) - canonical = n_channels - 1 if attention_class_index is None else int(attention_class_index) - if canonical < 0 or canonical >= n_channels: - raise ValueError( - f"attention_class_index={canonical} is out of range for attention with " - f"{n_channels} channel(s)." - ) - data[attention_column] = attention_matrix[:, canonical].tolist() - if n_channels > 1: - # Per-class attention and the raw pre-softmax logits. The logits are - # comparable across bags, which the within-bag normalised scores are - # not, and the attribution identity needs the per-class columns. - for channel in range(n_channels): - data[f"{attention_column}_c{channel}"] = attention_matrix[:, channel].tolist() - data[f"{attention_column}_logit_c{channel}"] = attention_logits[ - :, channel - ].tolist() - else: - data[attention_column] = _uniform_attention(bag_ids).tolist() - - for col in range(int(embeddings.shape[1])): - data[f"{embedding_prefix}_{col}"] = embeddings[:, col].tolist() - - composition = _tensor_to_2d_cpu(payload.instance_composition) - if composition is not None: - if composition.shape[0] != size: - raise ValueError( - "Mismatch between instance_logits and instance_composition rows: " - f"{size} vs {composition.shape[0]}." - ) - if composition_column_names is not None and len(composition_column_names) == int( - composition.shape[1] - ): - comp_cols = list(composition_column_names) - else: - comp_cols = [f"{composition_prefix}{idx}" for idx in range(int(composition.shape[1]))] - for col_name, values in zip(comp_cols, composition.T): - data[col_name] = values.tolist() - - if cell_type_column and payload.instance_graphs is not None: - # The type of the cell each instance is rooted at. - # - # There is exactly one instance per cell -- an instance is that - # cell's k-hop ego-graph -- so the instance table is a cell table, - # and this column makes it one that the cell-type plugins - # (`filtration_curves`, `per_niche_cell_type_enrichment`) can read. - # - # It is the root cell's own type, not a summary of the - # neighbourhood: summarising here would silently turn cell-level - # statistics into statistics over neighbourhood labels. The - # neighbourhood composition is already exported separately as the - # `comp_*` columns. - prefix_len = len(composition_prefix) - labels = [ - name[prefix_len:] if name.startswith(composition_prefix) else name - for name in comp_cols - ] - root_types = _resolve_instance_root_cell_types( - payload.instance_graphs, labels=labels, composition_label=cell_type_column - ) - if root_types is not None: - data[cell_type_column] = root_types - elif require_composition: - raise ValueError( - "Composition columns are required for interpretability export but " - "instance_composition was not provided by predict payload." - ) - - centroids = _tensor_to_2d_cpu(payload.instance_centroids) - if centroids is not None: - if centroids.shape[0] != size: - raise ValueError( - "Mismatch between instance_logits and instance_centroids rows: " - f"{size} vs {centroids.shape[0]}." - ) - if centroids.shape[1] >= 1: - data["center_x"] = centroids[:, 0].tolist() - if centroids.shape[1] >= 2: - data["center_y"] = centroids[:, 1].tolist() - - return pd.DataFrame(data) - - -def build_spatial_table( - instance_table: pd.DataFrame, - payload: BatchPredictionPayload, - *, - id_column: str = "instance_id", - undirected: bool = True, -) -> pd.DataFrame: - if id_column not in instance_table.columns: - raise ValueError(f"Cannot build spatial table, missing required column '{id_column}'.") - - instance_graphs = payload.instance_graphs - if instance_graphs is None: - raise ValueError( - "Cannot build spatial table: predict payload is missing " - "instance_graphs. Spatial edges must come from original subgraph connectivity." - ) - if len(instance_graphs) != len(instance_table): - raise ValueError( - "Mismatch between instance table rows and payload.instance_graphs: " - f"{len(instance_table)} vs {len(instance_graphs)}." - ) - - instance_ids = instance_table[id_column].astype(str).tolist() - bag_ids = ( - [str(v) for v in instance_table["bag_id"].tolist()] - if "bag_id" in instance_table - else ["" for _ in instance_ids] - ) - roots: List[int] = [] - for graph in instance_graphs: - roots.append(_resolve_instance_root_node_id(graph)) - - root_to_instance: Dict[tuple[str, int], int] = {} - for idx, (bag_id, root_id) in enumerate(zip(bag_ids, roots)): - root_to_instance[(bag_id, int(root_id))] = int(idx) - - edge_rows: List[tuple[str, str, float]] = [] - for src_idx, graph in enumerate(instance_graphs): - src_bag = bag_ids[src_idx] - src_root = roots[src_idx] - edge_index = getattr(graph, "edge_index", None) - if ( - not isinstance(edge_index, torch.Tensor) - or edge_index.ndim != 2 - or edge_index.shape[0] != 2 - ): - continue - n_id = getattr(graph, "n_id", None) - if not isinstance(n_id, torch.Tensor): - continue - n_id = n_id.detach().cpu().long().view(-1) - if n_id.numel() == 0: - continue - src = edge_index[0].detach().cpu().long().view(-1) - dst = edge_index[1].detach().cpu().long().view(-1) - if src.numel() != dst.numel(): - continue - - distance_values = _resolve_graph_edge_distances(graph, int(src.numel())) - for edge_pos, (local_src, local_dst) in enumerate(zip(src.tolist(), dst.tolist())): - if local_src < 0 or local_src >= int(n_id.numel()): - continue - if int(n_id[local_src].item()) != int(src_root): - continue - if local_dst < 0 or local_dst >= int(n_id.numel()): - continue - dst_root_global = int(n_id[local_dst].item()) - dst_idx = root_to_instance.get((src_bag, dst_root_global)) - if dst_idx is None: - continue - distance = ( - float(distance_values[edge_pos].item()) - if distance_values is not None and edge_pos < int(distance_values.numel()) - else float("nan") - ) - edge_rows.append((instance_ids[src_idx], instance_ids[dst_idx], distance)) - - if undirected: - symmetric_rows = [(dst, src, d) for src, dst, d in edge_rows] - edge_rows = edge_rows + symmetric_rows - - if not edge_rows: - return pd.DataFrame(columns=["source_id", "target_id", "distance"]) - - edge_df = pd.DataFrame(edge_rows, columns=["source_id", "target_id", "distance"]) - edge_df = ( - edge_df.sort_values(["source_id", "target_id", "distance"]) - .drop_duplicates(subset=["source_id", "target_id"], keep="first") - .reset_index(drop=True) - ) - return edge_df - - -def _resolve_instance_root_node_id(graph: object) -> int: - root_n_id = getattr(graph, "root_n_id", None) - if not isinstance(root_n_id, torch.Tensor) or root_n_id.numel() == 0: - raise ValueError( - "Each instance graph must include non-empty root_n_id tensor for spatial export." - ) - root_local = int(root_n_id.detach().cpu().view(-1)[0].item()) - root_is_global = getattr(graph, "root_n_id_is_global", None) - if isinstance(root_is_global, torch.Tensor) and root_is_global.numel() > 0: - if bool(root_is_global.detach().cpu().view(-1)[0].item()): - return root_local - elif isinstance(root_is_global, bool) and root_is_global: - return root_local - n_id = getattr(graph, "n_id", None) - if isinstance(n_id, torch.Tensor): - n_id = n_id.detach().cpu().long().view(-1) - if 0 <= root_local < int(n_id.numel()): - return int(n_id[root_local].item()) - return root_local - - -def _resolve_graph_edge_distances(graph: object, n_edges: int) -> Optional[torch.Tensor]: - edge_attr = getattr(graph, "edge_attr", None) - if not isinstance(edge_attr, torch.Tensor): - return None - edge_attr = edge_attr.detach().cpu() - if int(edge_attr.shape[0]) != int(n_edges): - return None - if edge_attr.ndim == 1: - return edge_attr.float().view(-1) - if edge_attr.ndim != 2 or int(edge_attr.shape[1]) == 0: - return None - dist_col = -1 - names = getattr(graph, "edge_attr_names", None) - if isinstance(names, (list, tuple)): - for idx, name in enumerate(names): - if str(name) == "distance": - dist_col = int(idx) - break - if dist_col == -1: - warnings.warn("No 'distance' column found in edge_attr, using column 1.") - dist_col = 1 - if dist_col >= int(edge_attr.shape[1]): - return None - return edge_attr[:, dist_col].float().view(-1) - - -def _edge_distance_column(graph: Any, n_edges: int) -> Optional[torch.Tensor]: - """Per-edge distances, preferring the column named ``distance``.""" - return _resolve_graph_edge_distances(graph, n_edges) - - -def build_cell_tables( - payload: BatchPredictionPayload, - instance_ids: Sequence[str], - *, - composition_column_names: Optional[Sequence[str]] = None, - composition_prefix: str = "comp_", - composition_label: str = "cell_type", - cell_id_column: str = "cell_id", - instance_id_column: str = "instance_id", -) -> Optional[tuple]: - """Cell-level node and edge tables pooled over every sampled ego-graph. - - The instance tables describe neighbourhoods: one row per ego-graph, with the - cell content reduced to composition fractions. Several analyses are only - meaningful over the cells themselves -- cell-type neighbourhood enrichment - within a niche, per-cell-type spatial autocorrelation, filtration over edge - distances -- and cannot be recovered from that summary. - - These tables restore it. Every node of every sampled subgraph becomes a row, - carrying its own cell type and coordinates, and every intra-subgraph edge - becomes a row with its distance. Joining ``instance_id`` against niche labels - gives each cell the niche of the ego-graph it came from, which is what makes - per-niche cell-level statistics possible. - - Cells are pooled *with duplicates*: overlapping ego-graphs revisit the same - cell, and each visit is a row. ``node_id`` carries the id in the source graph - so callers can deduplicate when that is the right thing to do. - - Returns ``(cell_table, cell_edge_table)``, or None when the payload carries - no subgraphs. - """ - graphs = payload.instance_graphs - if graphs is None: - return None - if len(graphs) != len(instance_ids): - raise ValueError( - "Mismatch between instance_graphs and instance ids: " - f"{len(graphs)} vs {len(instance_ids)}." - ) - - labels: Optional[List[str]] = None - if composition_column_names is not None: - prefix_len = len(composition_prefix) - labels = [ - name[prefix_len:] if name.startswith(composition_prefix) else name - for name in composition_column_names - ] - - cell_rows: List[Dict[str, Any]] = [] - edge_rows: List[Dict[str, Any]] = [] - - for instance_id, graph in zip(instance_ids, graphs): - codes = getattr(graph, "categorical_codes", None) - if not isinstance(codes, torch.Tensor): - return None - codes_2d = codes if codes.ndim == 2 else codes.view(-1, 1) - column = 0 - slices = getattr(graph, "categorical_slices", None) - if isinstance(slices, dict) and composition_label in slices: - column = int(slices[composition_label]) - if column >= int(codes_2d.shape[1]): - return None - - n_nodes = int(codes_2d.shape[0]) - pos = getattr(graph, "pos", None) - pos = pos.detach().cpu() if isinstance(pos, torch.Tensor) else None - n_id = getattr(graph, "n_id", None) - n_id = n_id.detach().cpu().long().view(-1) if isinstance(n_id, torch.Tensor) else None - - try: - root_global = _resolve_instance_root_node_id(graph) - except ValueError: - root_global = None - - cell_ids = [f"{instance_id}::n{idx}" for idx in range(n_nodes)] - for idx in range(n_nodes): - code = int(codes_2d[idx, column]) - row: Dict[str, Any] = { - cell_id_column: cell_ids[idx], - instance_id_column: instance_id, - composition_label: ( - labels[code] if labels is not None and 0 <= code < len(labels) else str(code) - ), - } - node_id = int(n_id[idx]) if n_id is not None and idx < int(n_id.numel()) else None - row["node_id"] = node_id - # node_id is patch-local, so (patch_id, node_id) identifies a real - # cell. Overlapping ego-graphs revisit cells; dedupe on that pair. - patch = getattr(graph, "patch_id", None) - if isinstance(patch, (list, tuple)) and patch: - patch = patch[0] - row["patch_id"] = None if patch is None else str(patch) - row["is_root"] = bool(root_global is not None and node_id == int(root_global)) - if pos is not None and idx < int(pos.shape[0]) and int(pos.shape[1]) >= 2: - row["x"] = float(pos[idx, 0]) - row["y"] = float(pos[idx, 1]) - cell_rows.append(row) - - edge_index = getattr(graph, "edge_index", None) - if isinstance(edge_index, torch.Tensor) and edge_index.numel(): - edge_index = edge_index.detach().cpu().long() - n_edges = int(edge_index.shape[1]) - distances = _edge_distance_column(graph, n_edges) - for e in range(n_edges): - src = int(edge_index[0, e]) - dst = int(edge_index[1, e]) - if src >= n_nodes or dst >= n_nodes: - continue - edge_rows.append( - { - "source_id": cell_ids[src], - "target_id": cell_ids[dst], - "distance": ( - float(distances[e]) if distances is not None else float("nan") - ), - instance_id_column: instance_id, - } - ) - - cell_table = pd.DataFrame(cell_rows) - edge_table = pd.DataFrame( - edge_rows, columns=["source_id", "target_id", "distance", instance_id_column] - ) - return cell_table, edge_table diff --git a/src/grass_mil/inference/io.py b/src/grass_mil/inference/io.py deleted file mode 100644 index dd9983e..0000000 --- a/src/grass_mil/inference/io.py +++ /dev/null @@ -1,107 +0,0 @@ -from __future__ import annotations - -import json -import math -from pathlib import Path -from typing import Any, Dict, Optional - -import pandas as pd -import torch - -from grass_mil.inference.schemas import BatchPredictionPayload, EmbeddingPayload - - -def ensure_output_dir(base_output_dir: Path, output_subdir: str) -> Path: - output_dir = base_output_dir / output_subdir - output_dir.mkdir(parents=True, exist_ok=True) - return output_dir - - -def _tensor_columns(prefix: str, values: torch.Tensor) -> Dict[str, Any]: - values = values.detach().cpu() - if values.is_floating_point(): - values = values.float() # bf16/fp16 autocast outputs have no numpy dtype - if values.ndim == 1: - values = values.unsqueeze(-1) - data: Dict[str, Any] = {} - for col in range(values.shape[1]): - data[f"{prefix}_{col}"] = values[:, col].tolist() - return data - - -def predictions_to_dataframe( - payload: BatchPredictionPayload, *, include_targets: bool, include_attention: bool -) -> pd.DataFrame: - data: Dict[str, Any] = {"bag_id": payload.bag_ids} - data.update(_tensor_columns("logit", payload.bag_logits)) - - if include_targets and payload.bag_targets is not None: - data.update(_tensor_columns("target", payload.bag_targets)) - - if include_attention and payload.bag_attention is not None: - if len(payload.bag_attention) != len(payload.bag_ids): - raise ValueError( - "Mismatch between bag_ids and attention rows while building dataframe: " - f"{len(payload.bag_ids)} ids vs {len(payload.bag_attention)} attention rows." - ) - data["attention"] = [ - json.dumps(attention.detach().cpu().tolist() if attention is not None else []) - for attention in payload.bag_attention - ] - - frame = pd.DataFrame(data) - return frame.sort_values(by=["bag_id"]).reset_index(drop=True) - - -def embeddings_to_dataframe(payload: EmbeddingPayload) -> pd.DataFrame: - data: Dict[str, Any] = {"bag_id": payload.bag_ids} - data.update(_tensor_columns("graph_emb", payload.graph_embeddings)) - frame = pd.DataFrame(data) - return frame.sort_values(by=["bag_id"]).reset_index(drop=True) - - -def node_embeddings_to_dataframe(payload: EmbeddingPayload) -> pd.DataFrame: - if payload.node_embeddings is None or payload.node_bag_ids is None: - return pd.DataFrame(columns=["bag_id"]) - data: Dict[str, Any] = {"bag_id": payload.node_bag_ids} - data.update(_tensor_columns("node_emb", payload.node_embeddings)) - frame = pd.DataFrame(data) - return frame.sort_values(by=["bag_id"]).reset_index(drop=True) - - -def write_dataframe(frame: pd.DataFrame, path: Path) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - frame.to_csv(path, index=False) - - -def _json_sanitize(value: Any) -> Any: - if isinstance(value, float) and not math.isfinite(value): - return None - if isinstance(value, dict): - return {key: _json_sanitize(item) for key, item in value.items()} - if isinstance(value, list): - return [_json_sanitize(item) for item in value] - if isinstance(value, tuple): - return [_json_sanitize(item) for item in value] - return value - - -def write_json(payload: Dict[str, Any], path: Path) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - sanitized_payload = _json_sanitize(payload) - path.write_text(json.dumps(sanitized_payload, indent=2, allow_nan=False)) - - -def build_summary_payload( - *, - prediction_path: Optional[Path], - metrics_path: Optional[Path], - embeddings_path: Optional[Path], - row_count: int, -) -> Dict[str, Any]: - return { - "prediction_path": str(prediction_path) if prediction_path else None, - "metrics_path": str(metrics_path) if metrics_path else None, - "embeddings_path": str(embeddings_path) if embeddings_path else None, - "row_count": int(row_count), - } diff --git a/src/grass_mil/inference/metrics.py b/src/grass_mil/inference/metrics.py deleted file mode 100644 index eaee6c2..0000000 --- a/src/grass_mil/inference/metrics.py +++ /dev/null @@ -1,273 +0,0 @@ -from __future__ import annotations - -import math -from typing import Any, Dict, List - -import torch -from omegaconf import DictConfig -from sklearn.metrics import balanced_accuracy_score, roc_auc_score - -from grass_mil.inference.schemas import BatchPredictionPayload - - -def _safe_div(num: float, den: float) -> float: - if den == 0.0: - return float("nan") - return float(num / den) - - -def _binary_roc_auc(scores: torch.Tensor, labels: torch.Tensor) -> float: - y_true = labels.long().view(-1).detach().cpu().numpy() - y_score = scores.float().view(-1).detach().cpu().numpy() - try: - return float(roc_auc_score(y_true, y_score)) - except ValueError: - # Preserve previous behavior when only one class is present. - return float("nan") - - -def _classification_binary( - logits: torch.Tensor, targets: torch.Tensor, threshold: float -) -> Dict[str, Any]: - logits = logits.float() - targets = targets.float() - if logits.ndim == 1: - logits = logits.unsqueeze(-1) - if targets.ndim == 1: - targets = targets.unsqueeze(-1) - if logits.shape[0] != targets.shape[0]: - raise ValueError( - "Binary metrics require logits and targets to have the same number of rows. " - f"Got logits {tuple(logits.shape)} and targets {tuple(targets.shape)}." - ) - if logits.shape[1] != targets.shape[1]: - raise ValueError( - "Binary metrics require matching task dimensions for logits and targets. " - f"Got logits {tuple(logits.shape)} and targets {tuple(targets.shape)}." - ) - - per_task: Dict[str, Dict[str, float]] = {} - for task_idx in range(int(logits.shape[1])): - scores = torch.sigmoid(logits[:, task_idx].reshape(-1)) - y_true = targets[:, task_idx].reshape(-1) - y_pred = (scores >= threshold).float() - - tp = float(((y_pred == 1) & (y_true == 1)).sum().item()) - tn = float(((y_pred == 0) & (y_true == 0)).sum().item()) - fp = float(((y_pred == 1) & (y_true == 0)).sum().item()) - fn = float(((y_pred == 0) & (y_true == 1)).sum().item()) - - precision = _safe_div(tp, tp + fp) - recall = _safe_div(tp, tp + fn) - f1 = _safe_div(2.0 * precision * recall, precision + recall) - accuracy = _safe_div(tp + tn, tp + tn + fp + fn) - roc_auc = _binary_roc_auc(scores=scores, labels=y_true.long()) - try: - balanced_accuracy = float( - balanced_accuracy_score( - y_true.detach().cpu().numpy().astype(int), - y_pred.detach().cpu().numpy().astype(int), - ) - ) - except ValueError: - balanced_accuracy = float("nan") - per_task[f"task_{task_idx}"] = { - "accuracy": accuracy, - "balanced_accuracy": balanced_accuracy, - "precision": precision, - "recall": recall, - "f1": f1, - "roc_auc": roc_auc, - } - - def _metric_mean(metric_name: str) -> float: - values = [ - float(metrics[metric_name]) - for metrics in per_task.values() - if not math.isnan(float(metrics[metric_name])) - ] - if not values: - return float("nan") - return float(sum(values) / len(values)) - - result: Dict[str, Any] = { - "accuracy": _metric_mean("accuracy"), - "balanced_accuracy": _metric_mean("balanced_accuracy"), - "precision": _metric_mean("precision"), - "recall": _metric_mean("recall"), - "f1": _metric_mean("f1"), - "roc_auc": _metric_mean("roc_auc"), - } - if len(per_task) > 1: - result["per_task"] = per_task - return result - - -def _classification_multiclass(logits: torch.Tensor, targets: torch.Tensor) -> Dict[str, float]: - pred_labels = torch.argmax(logits, dim=1) - true_labels = targets.long().view(-1) - num_classes = int(logits.shape[1]) - - acc = float((pred_labels == true_labels).float().mean().item()) - try: - bal_acc = float( - balanced_accuracy_score( - true_labels.detach().cpu().numpy(), pred_labels.detach().cpu().numpy() - ) - ) - except ValueError: - bal_acc = float("nan") - - per_class_precision: List[float] = [] - per_class_recall: List[float] = [] - per_class_f1: List[float] = [] - for cls in range(num_classes): - tp = float(((pred_labels == cls) & (true_labels == cls)).sum().item()) - fp = float(((pred_labels == cls) & (true_labels != cls)).sum().item()) - fn = float(((pred_labels != cls) & (true_labels == cls)).sum().item()) - precision = _safe_div(tp, tp + fp) - recall = _safe_div(tp, tp + fn) - f1 = _safe_div(2.0 * precision * recall, precision + recall) - per_class_precision.append(precision) - per_class_recall.append(recall) - per_class_f1.append(f1) - - def _nanmean(values: List[float]) -> float: - valid = [v for v in values if not math.isnan(v)] - if not valid: - return float("nan") - return float(sum(valid) / len(valid)) - - return { - "accuracy": acc, - "balanced_accuracy": bal_acc, - "precision": _nanmean(per_class_precision), - "recall": _nanmean(per_class_recall), - "f1": _nanmean(per_class_f1), - "roc_auc": float("nan"), - } - - -def _regression_metrics(logits: torch.Tensor, targets: torch.Tensor) -> Dict[str, float]: - preds = logits.float().view(-1) - y_true = targets.float().view(-1) - err = preds - y_true - - mae = float(err.abs().mean().item()) - rmse = float(torch.sqrt((err**2).mean()).item()) - ss_res = float((err**2).sum().item()) - centered = y_true - y_true.mean() - ss_tot = float((centered**2).sum().item()) - r2 = float("nan") if ss_tot == 0.0 else float(1.0 - ss_res / ss_tot) - return {"mae": mae, "rmse": rmse, "r2": r2} - - -def _survival_concordance_index( - risk_scores: torch.Tensor, - times: torch.Tensor, - events: torch.Tensor, - *, - chunk_size: int = 4096, -) -> float: - """Harrell's concordance index. - - A pair is comparable when the earlier subject had an observed event; it is - concordant when that subject also carries the higher risk, and tied - contributes a half. - - Evaluated in row chunks of the pair matrix, which bounds peak memory at - ``chunk_size * n`` booleans. The previous implementation was a Python - double loop with two ``.item()`` calls per pair, i.e. a device - synchronisation per pair. - """ - r = risk_scores.float().view(-1) - t = times.float().view(-1) - e = events.float().view(-1) - - n = int(r.shape[0]) - if n == 0: - return float("nan") - - concordant = 0.0 - ties = 0.0 - comparable = 0.0 - for start in range(0, n, chunk_size): - stop = min(start + chunk_size, n) - t_i = t[start:stop].unsqueeze(1) - r_i = r[start:stop].unsqueeze(1) - e_i = e[start:stop].unsqueeze(1) - - # Strict inequality on time also excludes the diagonal. - is_comparable = (e_i > 0) & (t_i < t.unsqueeze(0)) - if not bool(is_comparable.any()): - continue - comparable += float(is_comparable.sum().item()) - concordant += float((is_comparable & (r_i > r.unsqueeze(0))).sum().item()) - ties += float((is_comparable & (r_i == r.unsqueeze(0))).sum().item()) - - if comparable == 0.0: - return float("nan") - return float((concordant + 0.5 * ties) / comparable) - - -def _survival_metrics(logits: torch.Tensor, targets: torch.Tensor) -> Dict[str, float]: - if targets.ndim != 2 or targets.shape[1] < 2: - raise ValueError("Survival metrics require target columns [time, event].") - c_index = _survival_concordance_index( - risk_scores=logits.view(-1), - times=targets[:, 0], - events=targets[:, 1], - ) - return {"c_index": c_index} - - -def compute_task_metrics_from_tensors( - *, target_type: str, logits: torch.Tensor, targets: torch.Tensor, threshold: float -) -> Dict[str, Any]: - if target_type == "binary": - return _classification_binary(logits=logits, targets=targets, threshold=threshold) - if target_type == "regression": - return _regression_metrics(logits=logits, targets=targets) - if target_type == "survival": - return _survival_metrics(logits=logits, targets=targets) - if target_type == "categorical": - return _classification_multiclass(logits=logits, targets=targets) - raise ValueError(f"Unsupported target_type '{target_type}'.") - - -def compute_task_metrics( - *, - payload: BatchPredictionPayload, - target_type: str, - metrics_cfg: DictConfig, -) -> Dict[str, Any]: - """Compute global and per-group metrics for a prediction payload.""" - if payload.bag_targets is None: - return {"global": {}, "per_group": {}} - - threshold = float(metrics_cfg.categorical.threshold) - global_metrics = compute_task_metrics_from_tensors( - target_type=target_type, - logits=payload.bag_logits, - targets=payload.bag_targets, - threshold=threshold, - ) - - per_group: Dict[str, Dict[str, float]] = {} - if bool(metrics_cfg.per_group): - grouped_idx: Dict[str, List[int]] = {} - for idx, bag_id in enumerate(payload.bag_ids): - grouped_idx.setdefault(str(bag_id), []).append(idx) - - for bag_id, indices in grouped_idx.items(): - idx_tensor = torch.tensor(indices, dtype=torch.long) - group_logits = payload.bag_logits.index_select(0, idx_tensor) - group_targets = payload.bag_targets.index_select(0, idx_tensor) - per_group[bag_id] = compute_task_metrics_from_tensors( - target_type=target_type, - logits=group_logits, - targets=group_targets, - threshold=threshold, - ) - - return {"global": global_metrics, "per_group": per_group} diff --git a/src/grass_mil/inference/model_loading.py b/src/grass_mil/inference/model_loading.py deleted file mode 100644 index 046745b..0000000 --- a/src/grass_mil/inference/model_loading.py +++ /dev/null @@ -1,69 +0,0 @@ -"""Build a supervised module from config and load trained weights outside ``Trainer``.""" - -from __future__ import annotations - -from pathlib import Path -from typing import Any, Dict, Optional - -import hydra -import torch -from omegaconf import DictConfig - -from grass_mil.models.training.probe import attach_datamodule - -__all__ = ["load_trained_module", "read_checkpoint_state"] - - -def read_checkpoint_state(ckpt_path: str | Path) -> Dict[str, torch.Tensor]: - raw = torch.load(str(ckpt_path), map_location="cpu", weights_only=False) - state = raw["state_dict"] if isinstance(raw, dict) and "state_dict" in raw else raw - if not isinstance(state, dict): - raise ValueError(f"{ckpt_path} does not hold a state dict.") - return state - - -def load_trained_module( - cfg: DictConfig, - *, - ckpt_path: str | Path, - datamodule: Any, - device: Optional[torch.device | str] = None, - strict: bool = True, -): - """Instantiate ``cfg.model``, bind ``datamodule`` and load ``ckpt_path``. - - The datamodule must already be set up (``setup("fit")``) so the module can - infer its input widths. Returns the module in eval mode on ``device``. - """ - import copy - - from omegaconf import open_dict - - cfg = copy.deepcopy(cfg) - with open_dict(cfg): - # The checkpoint carries the trained head; a probe pre-fit request - # (linear_* regimes: task.head_init_path=auto) is a training-time step. - if "task" in cfg and cfg.task.get("head_init_path") is not None: - cfg.task.head_init_path = None - if "init_from_ckpt" in cfg.model: - cfg.model.init_from_ckpt = None - model = hydra.utils.instantiate(cfg.model) - trainer = hydra.utils.instantiate(cfg.trainer, logger=False, enable_checkpointing=False) - attach_datamodule(model, datamodule, trainer) - model.setup("fit") - state = read_checkpoint_state(ckpt_path) - missing, unexpected = model.load_state_dict(state, strict=False) - if strict: - missing = [k for k in missing if not k.startswith("loss_fn.")] - if missing or unexpected: - raise ValueError( - f"Checkpoint {ckpt_path} does not match the model: missing {missing[:5]}, " - f"unexpected {unexpected[:5]}. Compose the same model/regime as the run." - ) - if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - model.to(torch.device(device)) - model.eval() - for p in model.parameters(): - p.requires_grad_(False) - return model diff --git a/src/grass_mil/inference/predict.py b/src/grass_mil/inference/predict.py deleted file mode 100644 index a2aa769..0000000 --- a/src/grass_mil/inference/predict.py +++ /dev/null @@ -1,529 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple - -import hydra -from lightning import LightningDataModule, LightningModule, Trainer -from lightning.pytorch.loggers import Logger -from omegaconf import DictConfig - -from grass_mil._paths import setup_project_root - -setup_project_root() - -from grass_mil.inference.aggregation import aggregate_group_logits # noqa: E402 -from grass_mil.inference.collectors import collect_inference_payload # noqa: E402 -from grass_mil.inference.interpretability_export import ( # noqa: E402 - build_instance_table, - build_cell_tables, - build_spatial_table, - resolve_composition_column_names, -) -from grass_mil.inference.io import ( # noqa: E402 - build_summary_payload, - embeddings_to_dataframe, - ensure_output_dir, - node_embeddings_to_dataframe, - predictions_to_dataframe, - write_dataframe, - write_json, -) -from grass_mil.inference.metrics import compute_task_metrics # noqa: E402 -from grass_mil.inference.schemas import ( # noqa: E402 - BatchPredictionPayload, - PreforwardSubsamplingDecision, -) -from grass_mil.utils import ( # noqa: E402 - RankedLogger, - extras, - instantiate_loggers, - log_hyperparameters, - task_wrapper, -) - -log = RankedLogger(__name__, rank_zero_only=True) - - -def _supports_preforward_subsampling(mode: str) -> bool: - return mode in {"mean", "max", "attention_weighted"} - - -def _needs_instance_payload(cfg: DictConfig) -> bool: - if not bool(cfg.aggregation.enabled): - return False - mode = str(cfg.aggregation.mode) - return mode in {"mean", "max", "attention_weighted"} - - -def _plan_preforward_subsampling( - datamodule: LightningDataModule, cfg: DictConfig -) -> PreforwardSubsamplingDecision: - aggregation_enabled = bool(cfg.aggregation.enabled) - mode = str(cfg.aggregation.mode) - subsample_fraction = float(cfg.aggregation.get("subsample_fraction", 1.0)) - subsample_seed_raw = cfg.aggregation.get("subsample_seed") - subsample_seed = None if subsample_seed_raw is None else int(subsample_seed_raw) - requested = aggregation_enabled and subsample_fraction < 1.0 - - if not aggregation_enabled: - return PreforwardSubsamplingDecision( - requested=False, - effective=False, - fraction=subsample_fraction, - seed=subsample_seed, - reason="aggregation_disabled", - ) - if not requested: - return PreforwardSubsamplingDecision( - requested=False, - effective=False, - fraction=subsample_fraction, - seed=subsample_seed, - reason="not_requested", - ) - if not _supports_preforward_subsampling(mode): - return PreforwardSubsamplingDecision( - requested=True, - effective=False, - fraction=subsample_fraction, - seed=subsample_seed, - reason="unsupported_mode", - ) - - sampler_cfg = getattr(datamodule, "sampler_config", None) - if sampler_cfg is None: - return PreforwardSubsamplingDecision( - requested=True, - effective=False, - fraction=subsample_fraction, - seed=subsample_seed, - reason="missing_sampler_config", - ) - strategy_name_raw = getattr(sampler_cfg, "name", None) - strategy_name = None if strategy_name_raw is None else str(strategy_name_raw).strip().lower() - if not strategy_name: - return PreforwardSubsamplingDecision( - requested=True, - effective=False, - fraction=subsample_fraction, - seed=subsample_seed, - reason="missing_sampler_config", - ) - if strategy_name == "identity": - return PreforwardSubsamplingDecision( - requested=True, - effective=False, - fraction=subsample_fraction, - seed=subsample_seed, - reason="identity_sampler", - strategy_name=strategy_name, - runtime_enabled=None, - ) - - runtime = getattr(sampler_cfg, "runtime", None) - if runtime is None: - runtime = {} - runtime = dict(runtime) - runtime_enabled = bool(runtime.get("enabled", False)) - if not runtime_enabled: - return PreforwardSubsamplingDecision( - requested=True, - effective=False, - fraction=subsample_fraction, - seed=subsample_seed, - reason="runtime_disabled", - strategy_name=strategy_name, - runtime_enabled=False, - ) - - runtime["subsample_fraction"] = subsample_fraction - runtime["subsample_seed"] = subsample_seed - sampler_cfg.runtime = runtime - - # Force a strategy rebuild in case setup() was invoked previously. - if hasattr(datamodule, "sampler_strategy"): - setattr(datamodule, "sampler_strategy", None) - if hasattr(datamodule, "val_sampler_strategy"): - setattr(datamodule, "val_sampler_strategy", None) - return PreforwardSubsamplingDecision( - requested=True, - effective=True, - fraction=subsample_fraction, - seed=subsample_seed, - reason="applied", - strategy_name=strategy_name, - runtime_enabled=True, - ) - - -def _as_batch_payload(payload: Any) -> BatchPredictionPayload: - return BatchPredictionPayload( - bag_ids=list(payload.bag_ids), - bag_logits=payload.bag_logits, - bag_targets=payload.bag_targets, - bag_attention=getattr(payload, "bag_attention", None), - row_region_ids=getattr(payload, "row_region_ids", None), - row_sample_ids=getattr(payload, "row_sample_ids", None), - instance_logits=getattr(payload, "instance_logits", None), - instance_attention_logits=getattr(payload, "instance_attention_logits", None), - instance_patch_ids=getattr(payload, "instance_patch_ids", None), - instance_bag_ids=getattr(payload, "instance_bag_ids", None), - instance_region_ids=getattr(payload, "instance_region_ids", None), - instance_sample_ids=getattr(payload, "instance_sample_ids", None), - instance_embeddings=getattr(payload, "instance_embeddings", None), - instance_composition=getattr(payload, "instance_composition", None), - instance_centroids=getattr(payload, "instance_centroids", None), - instance_graphs=getattr(payload, "instance_graphs", None), - ) - - -def _preforward_decision_to_dict( - decision: PreforwardSubsamplingDecision, -) -> Dict[str, Any]: - return { - "requested": bool(decision.requested), - "effective": bool(decision.effective), - "fraction": float(decision.fraction), - "seed": decision.seed, - "reason": str(decision.reason), - "strategy_name": decision.strategy_name, - "runtime_enabled": decision.runtime_enabled, - } - - -def _load_patient_map(cfg_map) -> Optional[Dict[str, str]]: - """Region -> patient mapping from ``aggregation.patient_map`` (label CSV columns).""" - if not cfg_map or not cfg_map.get("label_file"): - return None - import pandas as pd - - df = pd.read_csv(str(cfg_map.label_file)) - df.columns = [str(c).strip() for c in df.columns] - id_col = str(cfg_map.get("id_column", "REGION_ID")) - pat_col = str(cfg_map.get("patient_column", "PATIENT")) - for col in (id_col, pat_col): - if col not in df.columns: - raise ValueError( - f"aggregation.patient_map: column '{col}' not in {cfg_map.label_file}" - ) - return {str(r): str(p) for r, p in zip(df[id_col], df[pat_col])} - - -@task_wrapper -def predict(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: - ckpt_path = cfg.get("ckpt_path") - if ckpt_path is None or not str(ckpt_path).strip(): - raise ValueError("Missing required `ckpt_path` for inference.") - - log.info(f"Instantiating datamodule <{cfg.data._target_}>") - datamodule: LightningDataModule = hydra.utils.instantiate(cfg.data) - - log.info(f"Instantiating model <{cfg.model._target_}>") - model: LightningModule = hydra.utils.instantiate(cfg.model) - - log.info("Instantiating loggers...") - logger: List[Logger] = instantiate_loggers(cfg.get("logger")) - - log.info(f"Instantiating trainer <{cfg.trainer._target_}>") - trainer: Trainer = hydra.utils.instantiate(cfg.trainer, logger=logger) - - object_dict = { - "cfg": cfg, - "datamodule": datamodule, - "model": model, - "logger": logger, - "trainer": trainer, - } - if logger: - log.info("Logging hyperparameters!") - log_hyperparameters(object_dict) - - base_output_dir = Path(str(cfg.paths.output_dir)) - output_dir = ensure_output_dir( - base_output_dir=base_output_dir, - output_subdir=str(cfg.predict.output_subdir), - ) - preforward_subsampling = _plan_preforward_subsampling(datamodule, cfg) - if preforward_subsampling.requested and not preforward_subsampling.effective: - raise ValueError( - "Preforward subsampling was requested via aggregation.subsample_fraction < 1.0 " - "but cannot be applied effectively. " - "Use data.sampler.name in {'shadow_native','shadow_custom'} with " - "data.sampler.runtime.enabled=true, or set aggregation.subsample_fraction=1.0. " - f"reason={preforward_subsampling.reason!r}, " - f"strategy={preforward_subsampling.strategy_name!r}, " - f"runtime_enabled={preforward_subsampling.runtime_enabled!r}." - ) - aggregation_metadata: Dict[str, Any] | None = None - interpret_cfg = cfg.get("interpretability") - export_interpretability = bool(interpret_cfg and interpret_cfg.get("enabled", False)) - include_instance_payload = _needs_instance_payload(cfg) or export_interpretability - include_instance_embeddings = export_interpretability - include_embeddings_payload = bool(cfg.embeddings.enabled) and bool(cfg.embeddings.save) - include_node_embeddings = include_embeddings_payload and bool(cfg.embeddings.extract_node) - previous_emit_setting = getattr(model, "_predict_emit_instance_payload", None) - had_emit_setting = hasattr(model, "_predict_emit_instance_payload") - previous_emit_instance_embedding_setting = getattr( - model, "_predict_emit_instance_embeddings", None - ) - had_emit_instance_embedding_setting = hasattr(model, "_predict_emit_instance_embeddings") - previous_emit_instance_comp_setting = getattr( - model, "_predict_emit_instance_composition", None - ) - had_emit_instance_comp_setting = hasattr(model, "_predict_emit_instance_composition") - previous_emit_instance_centroid_setting = getattr( - model, "_predict_emit_instance_centroids", None - ) - had_emit_instance_centroid_setting = hasattr(model, "_predict_emit_instance_centroids") - previous_predict_comp_label = getattr(model, "_predict_composition_label", None) - had_predict_comp_label = hasattr(model, "_predict_composition_label") - previous_emit_embeddings_setting = getattr(model, "_predict_emit_embeddings_payload", None) - had_emit_embeddings_setting = hasattr(model, "_predict_emit_embeddings_payload") - previous_emit_node_embeddings_setting = getattr(model, "_predict_emit_node_embeddings", None) - had_emit_node_embeddings_setting = hasattr(model, "_predict_emit_node_embeddings") - setattr(model, "_predict_emit_instance_payload", include_instance_payload) - setattr(model, "_predict_emit_instance_embeddings", include_instance_embeddings) - setattr(model, "_predict_emit_instance_composition", include_instance_embeddings) - setattr(model, "_predict_emit_instance_centroids", include_instance_embeddings) - setattr( - model, - "_predict_composition_label", - str((interpret_cfg or {}).get("composition_label", "cell_type")), - ) - setattr(model, "_predict_emit_embeddings_payload", include_embeddings_payload) - setattr(model, "_predict_emit_node_embeddings", include_node_embeddings) - - try: - collected = collect_inference_payload( - trainer=trainer, - model=model, - datamodule=datamodule, - ckpt_path=cfg.ckpt_path, - include_instance_payload=include_instance_payload, - include_instance_embeddings=include_instance_embeddings, - include_embeddings=include_embeddings_payload, - include_node_embeddings=include_node_embeddings, - ) - finally: - if had_emit_setting: - setattr(model, "_predict_emit_instance_payload", previous_emit_setting) - else: - delattr(model, "_predict_emit_instance_payload") - if had_emit_instance_embedding_setting: - setattr( - model, - "_predict_emit_instance_embeddings", - previous_emit_instance_embedding_setting, - ) - else: - delattr(model, "_predict_emit_instance_embeddings") - if had_emit_instance_comp_setting: - setattr( - model, "_predict_emit_instance_composition", previous_emit_instance_comp_setting - ) - else: - delattr(model, "_predict_emit_instance_composition") - if had_emit_instance_centroid_setting: - setattr( - model, "_predict_emit_instance_centroids", previous_emit_instance_centroid_setting - ) - else: - delattr(model, "_predict_emit_instance_centroids") - if had_predict_comp_label: - setattr(model, "_predict_composition_label", previous_predict_comp_label) - else: - delattr(model, "_predict_composition_label") - if had_emit_embeddings_setting: - setattr( - model, - "_predict_emit_embeddings_payload", - previous_emit_embeddings_setting, - ) - else: - delattr(model, "_predict_emit_embeddings_payload") - if had_emit_node_embeddings_setting: - setattr( - model, - "_predict_emit_node_embeddings", - previous_emit_node_embeddings_setting, - ) - else: - delattr(model, "_predict_emit_node_embeddings") - - raw_pred_payload = collected.prediction_payload - pred_payload = raw_pred_payload - - if bool(cfg.aggregation.enabled): - aggregation_subsample_fraction = float(cfg.aggregation.get("subsample_fraction", 1.0)) - aggregation_subsample_seed = cfg.aggregation.get("subsample_seed") - if preforward_subsampling.effective: - aggregation_subsample_fraction = 1.0 - aggregation_subsample_seed = None - aggregated = aggregate_group_logits( - pred_payload, - mode=str(cfg.aggregation.mode), - bag_scope=str(cfg.aggregation.get("bag_scope", "patch")), - subsample_fraction=aggregation_subsample_fraction, - subsample_seed=aggregation_subsample_seed, - group_map=_load_patient_map(cfg.aggregation.get("patient_map")), - ) - aggregation_metadata = dict(getattr(aggregated, "metadata", {})) - aggregation_metadata["preforward_subsampling_applied"] = bool( - preforward_subsampling.effective - ) - aggregation_metadata["preforward_subsampling"] = _preforward_decision_to_dict( - preforward_subsampling - ) - pred_payload = _as_batch_payload(aggregated) - - prediction_path = None - if bool(cfg.predict.save_predictions): - pred_frame = predictions_to_dataframe( - pred_payload, - include_targets=bool(cfg.predict.include_targets), - include_attention=bool(cfg.predict.include_attention), - ) - prediction_path = output_dir / str(cfg.output.predictions_filename) - write_dataframe(pred_frame, prediction_path) - else: - pred_frame = predictions_to_dataframe( - pred_payload, - include_targets=False, - include_attention=False, - ) - - metrics_path = None - metrics_payload: Dict[str, Any] = {} - if bool(cfg.predict.save_metrics) and bool(cfg.metrics.enabled): - metrics_payload = compute_task_metrics( - payload=pred_payload, - target_type=str(cfg.task.target_type), - metrics_cfg=cfg.metrics, - ) - metrics_path = output_dir / str(cfg.output.metrics_filename) - write_json(metrics_payload, metrics_path) - - embeddings_path = None - if include_embeddings_payload: - emb_payload = collected.embedding_payload - if emb_payload is None: - raise ValueError("Embeddings were requested but no embedding payload was collected.") - emb_frame = embeddings_to_dataframe(emb_payload) - embeddings_path = output_dir / str(cfg.embeddings.filename) - write_dataframe(emb_frame, embeddings_path) - - if include_node_embeddings: - node_frame = node_embeddings_to_dataframe(emb_payload) - node_path = output_dir / str(cfg.embeddings.node_filename) - write_dataframe(node_frame, node_path) - - interpretability_payload: Dict[str, Any] | None = None - if export_interpretability: - composition_prefix = str((interpret_cfg or {}).get("composition_prefix", "comp_")) - composition_label = str((interpret_cfg or {}).get("composition_label", "cell_type")) - composition_names = None - if raw_pred_payload.instance_composition is not None: - composition_names = resolve_composition_column_names( - datamodule, - composition_label=composition_label, - composition_prefix=composition_prefix, - width=int(raw_pred_payload.instance_composition.shape[1]), - ) - instance_frame = build_instance_table( - raw_pred_payload, - composition_prefix=composition_prefix, - embedding_prefix=str((interpret_cfg or {}).get("embedding_prefix", "inst_emb")), - composition_column_names=composition_names, - require_composition=bool((interpret_cfg or {}).get("require_composition", True)), - id_column=str((interpret_cfg or {}).get("id_column", "instance_id")), - bag_id_column=str((interpret_cfg or {}).get("bag_id_column", "bag_id")), - score_column=str((interpret_cfg or {}).get("score_column", "score")), - score_mode=str((interpret_cfg or {}).get("score_mode", "sigmoid")), - score_logit_index=int((interpret_cfg or {}).get("score_logit_index", 0)), - attention_column=str((interpret_cfg or {}).get("attention_column", "attention")), - attention_class_index=(interpret_cfg or {}).get("attention_class_index"), - cell_type_column=(interpret_cfg or {}).get("cell_type_column", "cell_type"), - ) - instance_path = output_dir / str( - (interpret_cfg or {}).get("instance_filename", "instance_table.csv") - ) - write_dataframe(instance_frame, instance_path) - - spatial_path = None - spatial_cfg = (interpret_cfg or {}).get("spatial", {}) - if bool(spatial_cfg.get("enabled", False)): - spatial_frame = build_spatial_table( - instance_frame, - raw_pred_payload, - id_column=str((interpret_cfg or {}).get("id_column", "instance_id")), - undirected=bool(spatial_cfg.get("undirected", True)), - ) - spatial_path = output_dir / str( - (interpret_cfg or {}).get("spatial_filename", "spatial_table.csv") - ) - write_dataframe(spatial_frame, spatial_path) - - cell_path = None - cell_edge_path = None - cell_cfg = (interpret_cfg or {}).get("cells", {}) - if bool(cell_cfg.get("enabled", False)): - # Cell-level tables: one row per node of every sampled ego-graph. - # Needed for the analyses that operate on cells rather than on - # neighbourhood summaries (cell-type enrichment within a niche, - # per-cell-type autocorrelation, distance filtration). - tables = build_cell_tables( - raw_pred_payload, - list(instance_frame[str((interpret_cfg or {}).get("id_column", "instance_id"))]), - composition_column_names=composition_names, - composition_prefix=composition_prefix, - composition_label=composition_label, - ) - if tables is None: - raise ValueError( - "interpretability.cells.enabled=true requires per-instance subgraphs " - "in the predict payload. Use a shadow sampler " - "(data.sampler.name=shadow_custom, data.sampler.runtime.enabled=true)." - ) - cell_frame, cell_edge_frame = tables - cell_path = output_dir / str(cell_cfg.get("cell_filename", "cell_table.csv")) - cell_edge_path = output_dir / str(cell_cfg.get("cell_edge_filename", "cell_edges.csv")) - write_dataframe(cell_frame, cell_path) - write_dataframe(cell_edge_frame, cell_edge_path) - - interpretability_payload = { - "instance_table_path": str(instance_path), - "spatial_table_path": str(spatial_path) if spatial_path is not None else None, - "cell_table_path": str(cell_path) if cell_path is not None else None, - "cell_edge_table_path": str(cell_edge_path) if cell_edge_path is not None else None, - "rows": int(len(instance_frame)), - } - - summary_path = output_dir / str(cfg.output.summary_filename) - summary_payload = build_summary_payload( - prediction_path=prediction_path, - metrics_path=metrics_path, - embeddings_path=embeddings_path, - row_count=len(pred_frame), - ) - summary_payload["metrics"] = metrics_payload - if aggregation_metadata is not None: - summary_payload["aggregation"] = aggregation_metadata - if interpretability_payload is not None: - summary_payload["interpretability"] = interpretability_payload - write_json(summary_payload, summary_path) - return summary_payload, object_dict - - -@hydra.main( - version_base="1.3", - config_path="../configs", - config_name="inference/predict.yaml", -) -def main(cfg: DictConfig) -> None: - extras(cfg) - predict(cfg) - - -if __name__ == "__main__": - main() diff --git a/src/grass_mil/inference/predictions.py b/src/grass_mil/inference/predictions.py new file mode 100644 index 0000000..a91318b --- /dev/null +++ b/src/grass_mil/inference/predictions.py @@ -0,0 +1,92 @@ +"""Unit-level prediction tables, shared by training, cross-validation and prediction.""" + +from __future__ import annotations + +from typing import Sequence + +import numpy as np +import pandas as pd +import torch + +__all__ = ["prediction_frame", "patient_frame", "frame_metrics", "score_column"] + + +def score_column(target_type: str) -> str: + """The column holding the score that metrics such as AUROC are computed on.""" + return {"classification": "prob_1", "regression": "pred_0", "survival": "risk"}[target_type] + + +def prediction_frame( + unit_ids: Sequence[str], + logits: torch.Tensor, + *, + target_type: str, + targets: torch.Tensor | None = None, + meta: pd.DataFrame | None = None, +) -> pd.DataFrame: + """One row per unit: ids, class probabilities / predictions / risk, and targets. + + ``meta`` (indexed by unit id) adds columns such as ``region_id``, + ``sample_id`` and ``patient_id``. + """ + logits = logits.detach().float().cpu() + df = pd.DataFrame({"unit_id": [str(u) for u in unit_ids]}) + if target_type == "classification": + probs = torch.softmax(logits, dim=1).numpy() + for c in range(probs.shape[1]): + df[f"prob_{c}"] = probs[:, c] + df["pred"] = probs.argmax(1) + elif target_type == "regression": + for c in range(logits.shape[1]): + df[f"pred_{c}"] = logits[:, c].numpy() + else: + df["risk"] = logits[:, 0].numpy() + if targets is not None: + t = targets.detach().float().cpu().numpy() + if target_type == "classification": + df["label"] = np.floor(t[:, 0] + 0.5).astype(int) + elif target_type == "survival": + df["time"], df["event"] = t[:, 0], t[:, 1] + else: + for c in range(t.shape[1]): + df[f"target_{c}"] = t[:, c] + if meta is not None: + df = df.join(meta, on="unit_id") + return df + + +def patient_frame(units: pd.DataFrame, target_type: str) -> pd.DataFrame: + """Patient scores as the mean of the patient's unit scores.""" + if "patient_id" not in units.columns: + raise ValueError("Patient-level scores need a patient_id column.") + value_cols = [c for c in units.columns if c.startswith(("prob_", "pred_")) or c == "risk"] + agg: dict[str, tuple[str, str]] = {c: (c, "mean") for c in value_cols} + agg["n_units"] = ("unit_id", "nunique") + for c in ("label", "time", "event", *[c for c in units.columns if c.startswith("target_")]): + if c in units.columns: + agg[c] = (c, "first") + for c in ("outer",): + if c in units.columns: + agg[c] = (c, "first") + out = units.groupby("patient_id").agg(**agg).reset_index() + if target_type == "classification": + probs = out[[c for c in out.columns if c.startswith("prob_")]].to_numpy() + out["pred"] = probs.argmax(1) + return out + + +def frame_metrics(frame: pd.DataFrame, target_type: str) -> dict[str, float]: + """Region-level metrics of a prediction table (unit or patient rows).""" + from grass_mil.models.objectives import region_metrics + + if target_type == "classification": + probs = frame[[c for c in frame.columns if c.startswith("prob_")]].to_numpy() + logits = torch.log(torch.as_tensor(probs, dtype=torch.float32).clamp_min(1e-12)) + targets = torch.as_tensor(frame[["label"]].to_numpy(), dtype=torch.float32) + elif target_type == "regression": + logits = torch.as_tensor(frame[[c for c in frame.columns if c.startswith("pred_")]].to_numpy(), dtype=torch.float32) + targets = torch.as_tensor(frame[[c for c in frame.columns if c.startswith("target_")]].to_numpy(), dtype=torch.float32) + else: + logits = torch.as_tensor(frame[["risk"]].to_numpy(), dtype=torch.float32) + targets = torch.as_tensor(frame[["time", "event"]].to_numpy(), dtype=torch.float32) + return region_metrics(target_type, logits, targets) diff --git a/src/grass_mil/inference/schemas.py b/src/grass_mil/inference/schemas.py deleted file mode 100644 index feed766..0000000 --- a/src/grass_mil/inference/schemas.py +++ /dev/null @@ -1,72 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Dict, List, Optional - -import torch - - -@dataclass(frozen=True) -class BatchPredictionPayload: - """Normalized prediction payload emitted by collector utilities.""" - - bag_ids: List[str] - bag_logits: torch.Tensor - bag_targets: Optional[torch.Tensor] - bag_attention: Optional[List[Optional[torch.Tensor]]] - row_region_ids: Optional[List[Optional[str]]] = None - row_sample_ids: Optional[List[Optional[str]]] = None - instance_logits: Optional[torch.Tensor] = None - instance_attention_logits: Optional[torch.Tensor] = None - instance_patch_ids: Optional[List[str]] = None - instance_bag_ids: Optional[List[str]] = None - instance_region_ids: Optional[List[Optional[str]]] = None - instance_sample_ids: Optional[List[Optional[str]]] = None - instance_embeddings: Optional[torch.Tensor] = None - instance_composition: Optional[torch.Tensor] = None - instance_centroids: Optional[torch.Tensor] = None - instance_graphs: Optional[List[Any]] = None - - -@dataclass(frozen=True) -class AggregatedPredictionPayload: - """Prediction payload after optional group-level aggregation.""" - - bag_ids: List[str] - bag_logits: torch.Tensor - bag_targets: Optional[torch.Tensor] - bag_attention: Optional[List[Optional[torch.Tensor]]] - metadata: Dict[str, Any] - row_region_ids: Optional[List[Optional[str]]] = None - row_sample_ids: Optional[List[Optional[str]]] = None - - -@dataclass(frozen=True) -class EmbeddingPayload: - """Embeddings aligned to bag-level ids.""" - - bag_ids: List[str] - graph_embeddings: torch.Tensor - node_embeddings: Optional[torch.Tensor] - node_bag_ids: Optional[List[str]] - - -@dataclass(frozen=True) -class PreforwardSubsamplingDecision: - """Decision record for preforward subsampling orchestration.""" - - requested: bool - effective: bool - fraction: float - seed: Optional[int] - reason: str - strategy_name: Optional[str] = None - runtime_enabled: Optional[bool] = None - - -@dataclass(frozen=True) -class CollectedInferencePayload: - """Normalized payloads gathered from a single trainer.predict pass.""" - - prediction_payload: BatchPredictionPayload - embedding_payload: Optional[EmbeddingPayload] diff --git a/src/grass_mil/loocv.py b/src/grass_mil/loocv.py deleted file mode 100644 index e2fbac4..0000000 --- a/src/grass_mil/loocv.py +++ /dev/null @@ -1,285 +0,0 @@ -from __future__ import annotations - -import copy -import hashlib -import json -import re -import warnings -from pathlib import Path -from statistics import mean, pstdev -from typing import Any, Dict, List, Optional - -import hydra -import lightning as L -import pandas as pd -from hydra.utils import instantiate -from omegaconf import DictConfig, open_dict - -from grass_mil._paths import setup_project_root - -setup_project_root() - -from grass_mil.data.components.precompute import discover_viable_loocv_fold_ids # noqa: E402 -from grass_mil.train import train # noqa: E402 -from grass_mil.utils import RankedLogger, extras # noqa: E402 - -log = RankedLogger(__name__, rank_zero_only=True) - - -def _validate_loocv_isolation_settings(cfg: DictConfig) -> None: - uses_per_fold_dirs = bool(cfg.loocv.per_fold_processed_dir) - uses_per_fold_recompute = bool(cfg.loocv.force_precompute_per_fold) - if uses_per_fold_dirs or uses_per_fold_recompute: - return - raise ValueError( - "Unsafe LOOCV configuration: both loocv.per_fold_processed_dir=false " - "and loocv.force_precompute_per_fold=false. This can silently reuse " - "a previous fold's processed_index/splits. Enable at least one of " - "these flags to guarantee fold isolation." - ) - - -def _canonical_fold_id(fold_unit: str, sample_id: str, region_id: Optional[str]) -> str: - if fold_unit == "sample": - return str(sample_id) - region = "__NONE__" if region_id is None else str(region_id) - return f"{sample_id}::{region}" - - -def _discover_folds(cfg: DictConfig) -> List[Dict[str, Optional[str]]]: - manifest_path = Path(str(cfg.data.raw_manifest_path)) - if not manifest_path.exists(): - raise FileNotFoundError(f"Manifest file not found: {manifest_path}") - df = pd.read_csv(manifest_path) - df.columns = [str(c).strip() for c in df.columns] - - sample_col = str(cfg.data.manifest.sample_id) - if sample_col not in df.columns: - raise ValueError(f"Manifest missing sample id column '{sample_col}'.") - - fold_unit = str(cfg.data.split.loocv.fold_unit) - folds: List[Dict[str, Optional[str]]] = [] - if fold_unit == "sample": - for sample_id in sorted(df[sample_col].astype(str).unique().tolist()): - folds.append( - { - "fold_id": str(sample_id), - "sample_id": str(sample_id), - "region_id": None, - } - ) - return folds - - if fold_unit != "region": - raise ValueError("data.split.loocv.fold_unit must be one of ['region', 'sample'].") - - region_col = str(cfg.data.manifest.region_id) - if region_col not in df.columns: - raise ValueError( - "Region-based LOOCV requires a region column in manifest. " f"Missing '{region_col}'." - ) - - pairs = ( - df[[sample_col, region_col]] - .drop_duplicates() - .sort_values(by=[sample_col, region_col], na_position="first") - ) - for _, row in pairs.iterrows(): - sample_id = str(row[sample_col]) - region_id = None if pd.isna(row[region_col]) else str(row[region_col]) - folds.append( - { - "fold_id": _canonical_fold_id("region", sample_id, region_id), - "sample_id": sample_id, - "region_id": region_id, - } - ) - return folds - - -def _resolve_selected_folds( - cfg: DictConfig, discovered: List[Dict[str, Optional[str]]] -) -> List[str]: - mode = str(cfg.loocv.mode) - if mode not in {"single", "all"}: - raise ValueError("loocv.mode must be one of ['single', 'all'].") - - canonical_ids = [str(item["fold_id"]) for item in discovered] - if mode == "all": - return canonical_ids - - requested = cfg.data.split.loocv.holdout_id - fold_index = cfg.loocv.fold_index - if requested is None and fold_index is None: - raise ValueError( - "Single LOOCV mode requires either data.split.loocv.holdout_id or loocv.fold_index." - ) - if fold_index is not None: - if requested is not None: - warnings.warn( - "Both data.split.loocv.holdout_id and loocv.fold_index are set in " - "single mode. loocv.fold_index takes priority.", - UserWarning, - stacklevel=2, - ) - idx = int(fold_index) - if idx < 0 or idx >= len(canonical_ids): - raise ValueError( - f"loocv.fold_index out of range: {idx}. Valid range: [0, {len(canonical_ids) - 1}]" - ) - return [canonical_ids[idx]] - - requested = str(requested) - if requested in canonical_ids: - return [requested] - - if str(cfg.data.split.loocv.fold_unit) == "region": - # Convenience support for plain region_id when unique. - matches = [item for item in discovered if item["region_id"] == requested] - if len(matches) == 1: - return [str(matches[0]["fold_id"])] - if len(matches) > 1: - raise ValueError( - f"Ambiguous region alias '{requested}'. Use canonical fold id " - f"'sample_id::region_id'. Matches: {[m['fold_id'] for m in matches]}" - ) - - raise ValueError( - f"Requested holdout id '{requested}' not found. Available fold ids: {canonical_ids}" - ) - - -def _discover_viable_fold_ids(cfg: DictConfig) -> List[str]: - datamodule = instantiate(cfg.data) - preprocessor = datamodule._build_preprocessor() - return discover_viable_loocv_fold_ids( - preprocessor=preprocessor, fold_unit=str(cfg.data.split.loocv.fold_unit) - ) - - -def _fold_slug(fold_id: str) -> str: - base = re.sub(r"[^a-zA-Z0-9._-]+", "_", fold_id).strip("_") - if not base: - base = "fold" - digest = hashlib.sha1(fold_id.encode("utf-8")).hexdigest()[:8] - return f"{base}_{digest}" - - -def _to_float_metrics(metrics: Dict[str, Any]) -> Dict[str, float]: - converted: Dict[str, float] = {} - for key, value in metrics.items(): - if hasattr(value, "item"): - try: - converted[key] = float(value.item()) - continue - except Exception: - pass - try: - converted[key] = float(value) - except Exception: - warnings.warn( - f"Skipping non-numeric LOOCV metric '{key}' of type " - f"'{type(value).__name__}' because it cannot be " - "converted to float.", - UserWarning, - stacklevel=2, - ) - continue - return converted - - -def _aggregate_fold_metrics( - rows: List[Dict[str, float]], -) -> Dict[str, Dict[str, float]]: - if not rows: - return {} - metric_names = sorted({name for row in rows for name in row}) - aggregated: Dict[str, Dict[str, float]] = {} - for name in metric_names: - vals = [row[name] for row in rows if name in row] - if not vals: - continue - aggregated[name] = { - "mean": float(mean(vals)), - "std": float(pstdev(vals)) if len(vals) > 1 else 0.0, - "n": float(len(vals)), - } - return aggregated - - -def run_loocv(cfg: DictConfig) -> Dict[str, Any]: - if cfg.get("seed"): - L.seed_everything(cfg.seed, workers=True) - _validate_loocv_isolation_settings(cfg) - - discovered = _discover_folds(cfg) - fold_ids = _resolve_selected_folds(cfg, discovered) - viable_fold_ids = _discover_viable_fold_ids(cfg) - viable_fold_set = set(viable_fold_ids) - missing_folds = [fold_id for fold_id in fold_ids if fold_id not in viable_fold_set] - if missing_folds: - raise ValueError( - "LOOCV preflight failed: selected folds are not viable after preprocessing " - "filters (e.g., min_cells). " - f"Missing selected folds: {missing_folds}. " - f"Available viable folds: {viable_fold_ids}" - ) - log.info(f"LOOCV fold count: {len(fold_ids)}") - - base_output_dir = Path(str(cfg.paths.output_dir)) - base_processed_dir = Path(str(cfg.data.processed_dir)) - base_output_dir.mkdir(parents=True, exist_ok=True) - - fold_results: List[Dict[str, Any]] = [] - fold_numeric_metrics: List[Dict[str, float]] = [] - - for fold_id in fold_ids: - fold_cfg = copy.deepcopy(cfg) - fold_slug = _fold_slug(fold_id) - with open_dict(fold_cfg): - fold_cfg.data.split.loocv.enabled = True - fold_cfg.data.split.loocv.holdout_id = fold_id - if bool(cfg.loocv.force_precompute_per_fold): - fold_cfg.data.force_precompute = True - if bool(cfg.loocv.per_fold_processed_dir): - fold_cfg.data.processed_dir = str(base_processed_dir / f"fold_{fold_slug}") - fold_cfg.paths.output_dir = str(base_output_dir / f"fold_{fold_slug}") - fold_cfg.task_name = f"{cfg.task_name}_fold_{fold_slug}" - fold_cfg.tags = ( - list(cfg.tags) + [f"loocv:{fold_id}"] if cfg.get("tags") else [f"loocv:{fold_id}"] - ) - - log.info(f"Starting fold: {fold_id}") - metric_dict, _ = train(fold_cfg) - numeric = _to_float_metrics(metric_dict) - fold_numeric_metrics.append(numeric) - fold_results.append( - { - "fold_id": fold_id, - "output_dir": str(fold_cfg.paths.output_dir), - "processed_dir": str(fold_cfg.data.processed_dir), - "metrics": numeric, - } - ) - - summary = { - "mode": str(cfg.loocv.mode), - "fold_unit": str(cfg.data.split.loocv.fold_unit), - "validation_strategy": str(cfg.data.split.loocv.validation_strategy), - "folds": fold_results, - "aggregate_metrics": _aggregate_fold_metrics(fold_numeric_metrics), - } - summary_path = base_output_dir / str(cfg.loocv.summary_filename) - summary_path.write_text(json.dumps(summary, indent=2)) - log.info(f"Wrote LOOCV summary to {summary_path}") - return summary - - -@hydra.main(version_base="1.3", config_path="configs", config_name="loocv.yaml") -def main(cfg: DictConfig) -> None: - extras(cfg) - run_loocv(cfg) - - -if __name__ == "__main__": - main() diff --git a/src/grass_mil/models/__init__.py b/src/grass_mil/models/__init__.py index f3fab0b..03024ac 100644 --- a/src/grass_mil/models/__init__.py +++ b/src/grass_mil/models/__init__.py @@ -1,47 +1,17 @@ -from .components import ( - BGRL, - AttnNetGated, - AttnNetGatedProjected, - CoxSGDLoss, - EncoderConfig, - GNNEncoder, - GraphPooling, - GraphPredictionHead, - MLPPredictor, - WeightedBCEWithLogitsLoss, - WeightedCrossEntropyLoss, - WeightedHuberLoss, - WeightedMSELoss, - build_attention, - build_encoder, - build_graph_head, - build_loss, - build_pooling, - build_ssl, -) -from .bgrl_module import BGRLModule -from .supervised_module import SupervisedModule +"""Models: the ego-graph encoder, attention-based MIL module and BGRL pretraining.""" + +from .encoder import EncoderConfig, GINEEncoder +from .mil import GatedAttention, InstanceHead, attention_pool, pool_by_bag +from .pretraining import BGRLModule +from .supervised import SupervisedModule __all__ = [ "BGRLModule", - "SupervisedModule", - "AttnNetGated", - "AttnNetGatedProjected", - "BGRL", - "CoxSGDLoss", "EncoderConfig", - "GNNEncoder", - "GraphPooling", - "GraphPredictionHead", - "MLPPredictor", - "WeightedBCEWithLogitsLoss", - "WeightedCrossEntropyLoss", - "WeightedHuberLoss", - "WeightedMSELoss", - "build_attention", - "build_encoder", - "build_graph_head", - "build_loss", - "build_pooling", - "build_ssl", + "GINEEncoder", + "GatedAttention", + "InstanceHead", + "SupervisedModule", + "attention_pool", + "pool_by_bag", ] diff --git a/src/grass_mil/models/augmentations.py b/src/grass_mil/models/augmentations.py new file mode 100644 index 0000000..8a3a10d --- /dev/null +++ b/src/grass_mil/models/augmentations.py @@ -0,0 +1,183 @@ +"""Stochastic graph augmentations (pretraining views and fine-tuning regularisation). + +All functions take an explicit ``torch.Generator`` (so seeded runs are +reproducible) and return a new batch instead of modifying their input. Edges +are always dropped in both directions together, so the graph stays undirected. +A masked node keeps its place in the graph, its features are zeroed and its +cell type is set to the encoder's "unassigned" code. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import Any, Callable, Mapping, Sequence + +import torch + +__all__ = [ + "PretrainViews", + "FinetuneAugmentation", + "degree_importance", + "drop_importance", + "drop_edges", + "mask_nodes", + "feature_noise", +] + + +def _rand(shape, generator, device) -> torch.Tensor: + # A CPU generator cannot drive a CUDA draw: draw on CPU and move. + return torch.rand(shape, generator=generator).to(device) + + +def _undirected(edge_index: torch.Tensor, keep: torch.Tensor) -> torch.Tensor: + """Restrict ``keep`` to edges whose reverse edge is kept as well.""" + if edge_index.numel() == 0: + return keep + n = int(edge_index.max()) + 1 + src, dst = edge_index + kept = src[keep] * n + dst[keep] + return keep & torch.isin(dst * n + src, kept) + + +def _keep_edges(batch, keep: torch.Tensor): + batch.edge_index = batch.edge_index[:, keep] + if getattr(batch, "edge_attr", None) is not None: + batch.edge_attr = batch.edge_attr[keep] + return batch + + +def mask_nodes(batch, masked: torch.Tensor, unassigned_index: int | None): + """Zero the features of ``masked`` nodes and mark their cell type unassigned.""" + if batch.x.numel(): + batch.x = batch.x.masked_fill(masked.view(-1, 1), 0.0) + if unassigned_index is not None and getattr(batch, "cell_type", None) is not None: + batch.cell_type = batch.cell_type.masked_fill(masked, int(unassigned_index)) + return batch + + +def degree_importance(edge_index: torch.Tensor, num_nodes: int) -> torch.Tensor: + """Per-node importance in ``[0, 1]``: log degree normalised as ``(v - mean) / (max - mean)``, + clamped at zero (every below-average node gets 0).""" + from torch_geometric.utils import degree + + v = torch.log1p(degree(edge_index[0], num_nodes)) + return torch.clamp((v - v.mean()) / (v.max() - v.mean() + 1e-8), min=0.0) + + +def drop_importance(batch, *, mu: float, p_lambda: float, unassigned_index: int | None, generator): + """Mask low-degree nodes and drop low-importance edges. + + The drop probability of a node (or edge) with importance ``s`` is + ``min((1 - s) * mu, p_lambda)``; an edge's importance is the mean of its + end points', renormalised the same way. + """ + if not (0.0 < mu <= 1.0 and 0.0 < p_lambda <= 1.0): + raise ValueError("mu and p_lambda must be in (0, 1].") + out = batch.clone() + imp = degree_importance(out.edge_index, int(out.num_nodes)) + p_node = torch.clamp((1.0 - imp) * mu, max=p_lambda) + out = mask_nodes(out, _rand(p_node.shape, generator, p_node.device) <= p_node, unassigned_index) + if out.edge_index.numel(): + e = (imp[out.edge_index[0]] + imp[out.edge_index[1]]) / 2.0 + e = torch.clamp((e - e.mean()) / (e.max() - e.mean() + 1e-8), min=0.0) + p_edge = torch.clamp((1.0 - e) * mu, max=p_lambda) + keep = _rand(p_edge.shape, generator, p_edge.device) > p_edge + out = _keep_edges(out, _undirected(out.edge_index, keep)) + return out + + +def drop_edges(batch, *, p: float, generator): + """Drop each undirected edge with probability ``p``.""" + if p <= 0.0 or batch.edge_index.numel() == 0: + return batch + out = batch.clone() + keep = _rand((out.edge_index.size(1),), generator, out.edge_index.device) >= p + return _keep_edges(out, _undirected(out.edge_index, keep)) + + +def feature_noise(batch, *, std: float, columns: Sequence[int], generator): + """Add Gaussian noise of ``std`` to the given feature columns.""" + if std <= 0.0 or not columns or batch.x.numel() == 0: + return batch + out = batch.clone() + x = out.x.clone() + for c in columns: + x[:, c] += torch.randn(x.size(0), generator=generator).to(x.device) * std + out.x = x + return out + + +def _columns(columns: Sequence[Any] | None, feature_names: Sequence[str]) -> list[int]: + out = [] + for c in columns or []: + if isinstance(c, int): + out.append(c) + elif str(c) in feature_names: + out.append(list(feature_names).index(str(c))) + else: + raise ValueError(f"Feature '{c}' not among the node features {list(feature_names)}.") + return out + + +@dataclass(frozen=True) +class PretrainViews: + """Two stochastic views per batch: importance-weighted node/edge drop plus + Gaussian noise on selected continuous features.""" + + mu: float = 0.25 + p_lambda: float = 0.45 + noise_std: float = 0.2 + noise_features: tuple = () + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any] | None) -> PretrainViews: + cfg = dict(cfg or {}) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown pretrain.augmentation keys: {sorted(unknown)}") + cfg["noise_features"] = tuple(cfg.get("noise_features") or ()) + return cls(**cfg) + + def build(self, *, feature_names: Sequence[str], unassigned_index: int | None) -> Callable: + columns = _columns(self.noise_features, feature_names) + + def view(batch, generator): + out = drop_importance(batch, mu=self.mu, p_lambda=self.p_lambda, unassigned_index=unassigned_index, generator=generator) + return feature_noise(out, std=self.noise_std, columns=columns, generator=generator) + + return view + + +@dataclass(frozen=True) +class FinetuneAugmentation: + """Uniform edge drop and node masking during fine-tuning (both off by default).""" + + drop_edge_p: float = 0.0 + mask_node_p: float = 0.0 + seed: int = 0 + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any] | None) -> FinetuneAugmentation: + cfg = dict(cfg or {}) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown finetune.augmentation keys: {sorted(unknown)}") + return cls(**cfg) + + @property + def enabled(self) -> bool: + return self.drop_edge_p > 0 or self.mask_node_p > 0 + + def build(self, *, unassigned_index: int | None) -> Callable: + generator = torch.Generator().manual_seed(int(self.seed)) + + def augment(batch): + out = drop_edges(batch, p=self.drop_edge_p, generator=generator) + if self.mask_node_p > 0: + out = out.clone() if out is batch else out + masked = _rand((int(out.num_nodes),), generator, out.x.device) < self.mask_node_p + out = mask_nodes(out, masked, unassigned_index) + return out + + return augment diff --git a/src/grass_mil/models/bgrl_module.py b/src/grass_mil/models/bgrl_module.py deleted file mode 100644 index 1ab2931..0000000 --- a/src/grass_mil/models/bgrl_module.py +++ /dev/null @@ -1,266 +0,0 @@ -from __future__ import annotations - -from typing import Any, Dict, Optional - -import lightning as L -import torch -import torch.nn.functional as F - -from .components import EncoderConfig, build_encoder, build_ssl -from .training import ( - CosineWarmup, - build_augmentation, - infer_encoder_input_dim, - resolve_encoder_cfg, - instantiate_optimizer, - instantiate_scheduler_with_warmup, - load_state_dict_with_optional_mapping, -) - - -def _molecular_feature_names(module) -> Optional[list]: - """Feature names from the datamodule, for name-based column resolution.""" - try: - trainer = module.trainer - except RuntimeError: - return None - datamodule = getattr(trainer, "datamodule", None) if trainer is not None else None - dataset = getattr(datamodule, "dataset_train", None) - if dataset is None or len(dataset) == 0: - return None - names = getattr(dataset[0], "molecular_feature_names", None) - if names: - return list(names) - # Older precomputes record the names only in metadata.json next to the index. - index_path = getattr(getattr(dataset, "dataset", dataset), "index_path", None) - if index_path is not None: - import json - from pathlib import Path - - meta = Path(index_path).parent / "metadata.json" - if meta.exists(): - try: - names = json.loads(meta.read_text()).get("molecular_feature_names") - except (OSError, ValueError): - names = None - if names: - return list(map(str, names)) - return None - - -class BGRLModule(L.LightningModule): - def __init__( - self, - encoder: Dict[str, Any], - ssl: Dict[str, Any], - optim: Optional[Dict[str, Any]] = None, - scheduler: Optional[Dict[str, Any]] = None, - task: Optional[Dict[str, Any]] = None, - init_from_ckpt: Optional[str] = None, - init_strict: bool = False, - encoder_init_map: str = "auto_bgrl_or_identity", - ) -> None: - super().__init__() - self.save_hyperparameters(logger=False) - - self._built = False - self.encoder = None - self.ssl_model = None - self.optim_cfg = optim - self.scheduler_cfg = scheduler - self.task_cfg = task or {} - self.init_from_ckpt = init_from_ckpt - self.init_strict = init_strict - self.encoder_init_map = encoder_init_map - - # Per-view uniform drop rates, used when no augmentation config is set. - self._view_params = ( - { - "mode": "uniform", - "drop_edge_p": float(self.task_cfg.get("drop_edge_p1", 0.0)), - "drop_feat_p": float(self.task_cfg.get("drop_feat_p1", 0.0)), - }, - { - "mode": "uniform", - "drop_edge_p": float(self.task_cfg.get("drop_edge_p2", 0.0)), - "drop_feat_p": float(self.task_cfg.get("drop_feat_p2", 0.0)), - }, - ) - self._augmentation_seed = int(self.task_cfg.get("augmentation_seed", 0)) - self._view_generator: Optional[torch.Generator] = None - # EMA momentum of the target encoder: cosine from `momentum` to - # `momentum_final` (preprint: 0.99 -> 1.0, no warmup). `momentum_min` is - # the legacy name of the final value. - self._momentum_base = float(self.task_cfg.get("momentum", 0.99)) - final = self.task_cfg.get("momentum_final") - if final is None: - final = self.task_cfg.get("momentum_min", self._momentum_base) - self._momentum_min = float(final) - self._warmup_steps = int(self.task_cfg.get("warmup_steps", 0)) - momentum_warmup = self.task_cfg.get("momentum_warmup_steps") - self._momentum_warmup_steps = ( - self._warmup_steps if momentum_warmup is None else int(momentum_warmup) - ) - self._total_steps = int(self.task_cfg.get("total_steps", 0)) - - self._momentum_scheduler: Optional[CosineWarmup] = None - self._last_momentum_step = 0 - # When set, replaces the uniform drop_edge/drop_feat views. - self._augmentation_cfg = self.task_cfg.get("augmentation") - self._augmentation = None - self._fallback_augmentations = None - - def setup(self, stage: Optional[str] = None) -> None: - if self._built: - return - encoder_cfg = resolve_encoder_cfg( - dict(self.hparams.encoder), infer_encoder_input_dim(self), module=self - ) - self.encoder = build_encoder(EncoderConfig(**encoder_cfg)) - input_proj = getattr(self.encoder, "input_proj", None) - categorical = getattr(input_proj, "categorical", None) - build_kwargs = { - "feature_names": _molecular_feature_names(self), - "unassigned_index": getattr(input_proj, "unassigned_index", None), - "categorical_column": (None if categorical is None else categorical.column_index or 0), - } - if self._augmentation_cfg: - self._augmentation = build_augmentation(self._augmentation_cfg, **build_kwargs) - else: - self._fallback_augmentations = tuple( - build_augmentation(params, **build_kwargs) for params in self._view_params - ) - self.ssl_model = build_ssl(True, self.encoder, dict(self.hparams.ssl)) - if self.ssl_model is None: - raise ValueError("BGRLModule requires ssl.method=bgrl config.") - result = load_state_dict_with_optional_mapping( - self, - init_from_ckpt=self.init_from_ckpt, - init_strict=self.init_strict, - encoder_init_map=self.encoder_init_map, - ) - if result["missing_keys"] or result["unexpected_keys"]: - self.log( - "init/state_dict_warnings", - float(len(result["missing_keys"]) + len(result["unexpected_keys"])), - prog_bar=False, - ) - total_steps = self._total_steps - if total_steps <= 0 and self.trainer is not None: - total_steps = int(getattr(self.trainer, "estimated_stepping_batches", 0)) - # Give a warning if total_steps is not set - self.log("init/total_steps_warning", float(total_steps), prog_bar=False) - self._momentum_scheduler = CosineWarmup( - base_value=self._momentum_base, - warmup_steps=max(0, self._momentum_warmup_steps), - total_steps=max(1, total_steps), - min_value=self._momentum_min, - ) - self._built = True - - def _encode_tuple(self, batch): - return ( - batch.x, - batch.edge_index, - getattr(batch, "edge_attr", None), - getattr(batch, "batch", None), - getattr(batch, "categorical_codes", None), - ) - - def _view_rng(self) -> torch.Generator: - """Generator for the augmentation draws. - - Seeded once from ``task.augmentation_seed`` and then advanced, so the - two views of a batch differ while the whole pretraining run stays - reproducible. Without an explicit generator the augmentations would - consume the global RNG and silently depend on unrelated call order. - """ - if self._view_generator is None: - self._view_generator = torch.Generator() - self._view_generator.manual_seed(self._augmentation_seed) - return self._view_generator - - def _make_views(self, batch): - generator = self._view_rng() - if self._augmentation is not None: - return self._augmentation(batch, generator), self._augmentation(batch, generator) - first, second = self._fallback_augmentations - return first(batch, generator), second(batch, generator) - - def _step(self, batch) -> torch.Tensor: - view1, view2 = self._make_views(batch) - q1, y2 = self.ssl_model(self._encode_tuple(view1), self._encode_tuple(view2)) - q2, y1 = self.ssl_model(self._encode_tuple(view2), self._encode_tuple(view1)) - loss = ( - 2.0 - - F.cosine_similarity(q1, y2.detach(), dim=-1).mean() - - F.cosine_similarity(q2, y1.detach(), dim=-1).mean() - ) - return loss - - def training_step(self, batch, batch_idx): - loss = self._step(batch) - self.log("train/loss", loss, on_step=True, on_epoch=True, prog_bar=True) - return loss - - def validation_step(self, batch, batch_idx): - loss = self._step(batch) - self.log("val/loss", loss, on_step=False, on_epoch=True, prog_bar=True) - return loss - - def test_step(self, batch, batch_idx): - loss = self._step(batch) - self.log("test/loss", loss, on_step=False, on_epoch=True, prog_bar=True) - return loss - - def on_train_batch_end(self, outputs, batch, batch_idx) -> None: - """EMA update of the target encoder after every optimizer step. - - Runs once per completed optimizer step (``global_step`` has just been - incremented), with the momentum of that zero-based step, exactly like - the reference loop (``update_target_network(mm_now)`` after ``step()``). - """ - if self.ssl_model is None or self._momentum_scheduler is None: - return - completed = int(self.global_step) - if completed <= self._last_momentum_step: - return # gradient-accumulation sub-batch: no optimizer step happened - self._last_momentum_step = completed - momentum = self._momentum_scheduler.value(completed - 1) - self.ssl_model.update_target_network(momentum) - self.log("train/momentum", momentum, on_step=True, on_epoch=False, prog_bar=False) - - def configure_optimizers(self): - if not self._built: - self.setup("fit") - if self.optim_cfg is None: - raise ValueError( - "Missing optimizer config for training. Compose an optim group " - "(for example '+optim=adamw') or set model.optim explicitly." - ) - optimizer = instantiate_optimizer(self.optim_cfg, self.ssl_model.trainable_parameters()) - # `scheduler.T_max: auto` = cosine over the steps after the LR warmup, so - # the learning rate reaches eta_min exactly at `task.total_steps`. - warm_cfg = self.task_cfg.get("lr_warmup") or {} - lr_warmup_steps = ( - int(warm_cfg.get("warmup_steps", 0)) if bool(warm_cfg.get("enabled", False)) else 0 - ) - t_max_auto = None - if self._total_steps > 0: - t_max_auto = max(1, self._total_steps - max(0, lr_warmup_steps)) - scheduler = instantiate_scheduler_with_warmup( - self.scheduler_cfg, - optimizer, - warmup_cfg=self.task_cfg.get("lr_warmup"), - t_max_auto=t_max_auto, - ) - if scheduler is None: - return optimizer - return { - "optimizer": optimizer, - "lr_scheduler": { - "scheduler": scheduler, - "interval": "step", - "frequency": 1, - }, - } diff --git a/src/grass_mil/models/checkpoints.py b/src/grass_mil/models/checkpoints.py new file mode 100644 index 0000000..19561c9 --- /dev/null +++ b/src/grass_mil/models/checkpoints.py @@ -0,0 +1,26 @@ +"""Loading encoder weights from a pretraining (or fine-tuning) checkpoint.""" + +from __future__ import annotations + +from pathlib import Path + +import torch + +__all__ = ["load_encoder_weights"] + + +def load_encoder_weights(encoder: torch.nn.Module, path: str | Path) -> None: + """Copy the ``encoder.*`` weights of a checkpoint into ``encoder``. + + Both module types store their (online) encoder under ``encoder``, so a BGRL + checkpoint and a fine-tuned checkpoint initialise an encoder the same way. + """ + path = Path(path) + if not path.exists(): + raise FileNotFoundError(f"Encoder checkpoint not found: {path}") + state = torch.load(path, map_location="cpu", weights_only=False) + state = state.get("state_dict", state) + weights = {k[len("encoder.") :]: v for k, v in state.items() if k.startswith("encoder.")} + if not weights: + raise ValueError(f"{path} holds no encoder weights (no 'encoder.*' keys).") + encoder.load_state_dict(weights, strict=True) diff --git a/src/grass_mil/models/components/__init__.py b/src/grass_mil/models/components/__init__.py deleted file mode 100644 index 2d89b69..0000000 --- a/src/grass_mil/models/components/__init__.py +++ /dev/null @@ -1,45 +0,0 @@ -from .attention import AttnNetGated, AttnNetGatedProjected -from .backbones import EncoderConfig, GNNEncoder -from .embeddings import CategoricalEmbeddingConfig, NodeInputEmbedding -from .factory import ( - build_attention, - build_encoder, - build_graph_head, - build_loss, - build_pooling, - build_ssl, -) -from .heads import GraphPredictionHead -from .losses import ( - CoxSGDLoss, - WeightedBCEWithLogitsLoss, - WeightedCrossEntropyLoss, - WeightedHuberLoss, - WeightedMSELoss, -) -from .pooling import GraphPooling -from .ssl import BGRL, MLPPredictor - -__all__ = [ - "AttnNetGated", - "AttnNetGatedProjected", - "BGRL", - "CoxSGDLoss", - "CategoricalEmbeddingConfig", - "EncoderConfig", - "NodeInputEmbedding", - "GNNEncoder", - "GraphPooling", - "GraphPredictionHead", - "MLPPredictor", - "WeightedBCEWithLogitsLoss", - "WeightedCrossEntropyLoss", - "WeightedHuberLoss", - "WeightedMSELoss", - "build_attention", - "build_encoder", - "build_graph_head", - "build_loss", - "build_pooling", - "build_ssl", -] diff --git a/src/grass_mil/models/components/attention.py b/src/grass_mil/models/components/attention.py deleted file mode 100644 index c064271..0000000 --- a/src/grass_mil/models/components/attention.py +++ /dev/null @@ -1,134 +0,0 @@ -from __future__ import annotations - -from typing import Literal - -import torch -from torch import nn - -InitType = Literal["default", "kaiming"] - - -def _init_linears(module: nn.Module, init: InitType) -> None: - """Apply the requested initialisation to every ``nn.Linear`` in ``module``. - - ``kaiming`` is the reference GRASS-MIL initialisation (``kaiming_uniform_`` on - the weights, biases left at the torch default); ``default`` keeps torch's own. - """ - if init == "default": - return - if init != "kaiming": - raise ValueError(f"init must be 'default' or 'kaiming', got {init!r}.") - for layer in module.modules(): - if isinstance(layer, nn.Linear): - nn.init.kaiming_uniform_(layer.weight) - - -class AttnNetGated(nn.Module): - """Gated MIL attention network (Ilse et al. 2018), one attention channel per class. - - ``forward`` returns the raw (pre-softmax) attention logits of shape - ``[num_instances, n_classes]``; the softmax over instances is applied by the - bag aggregation so that it can be taken within each bag. - """ - - def __init__( - self, - input_dim: int, - hidden_dim: int = 64, - dropout: bool = False, - n_classes: int = 1, - init: InitType = "kaiming", - ) -> None: - super().__init__() - a_layers = [nn.Linear(input_dim, hidden_dim), nn.Tanh()] - b_layers = [nn.Linear(input_dim, hidden_dim), nn.Sigmoid()] - if dropout: - a_layers.append(nn.Dropout(0.25)) - b_layers.append(nn.Dropout(0.25)) - self.attention_a = nn.Sequential(*a_layers) - self.attention_b = nn.Sequential(*b_layers) - self.attention_c = nn.Linear(hidden_dim, n_classes) - self.init = init - _init_linears(self, init) - - def forward(self, x): - a = self.attention_a(x) - b = self.attention_b(x) - logits = self.attention_c(a * b) - return logits, x - - def attention_space(self, x: torch.Tensor) -> torch.Tensor: - """The ``hidden_dim``-d pre-tanh activation, the attention head's own feature space.""" - return self.attention_a[0](x) - - def head_spaces(self, x): - """Intermediate activations of the attention head. - - ``a_pre_tanh`` is the argument of the tanh gate, before saturation. It - is a low-dimensional space directly upstream of the attention weights; - clustering the saturated output collapses the extremes together. - """ - pre_tanh = self.attention_space(x) - a = self.attention_a(x) - b = self.attention_b(x) - return {"a_pre_tanh": pre_tanh, "gate_product": a * b} - - -class AttnNetGatedProjected(nn.Module): - """Projected gated MIL attention: ``Linear + LeakyReLU`` projection, then the gate. - - GRASS-MIL default: ``projection_dim=32``, ``hidden_dim=8``, dropout 0.25 in - training, one channel per class, kaiming-initialised linears. - """ - - def __init__( - self, - input_dim: int, - projection_dim: int = 32, - hidden_dim: int = 8, - dropout: bool = True, - n_classes: int = 1, - init: InitType = "kaiming", - ) -> None: - super().__init__() - projection = [nn.Linear(input_dim, projection_dim), nn.LeakyReLU()] - attn_a = [nn.Linear(projection_dim, hidden_dim), nn.Tanh()] - attn_b = [nn.Linear(projection_dim, hidden_dim), nn.Sigmoid()] - if dropout: - projection.append(nn.Dropout(0.25)) - attn_a.append(nn.Dropout(0.25)) - attn_b.append(nn.Dropout(0.25)) - self.projection = nn.Sequential(*projection) - self.attention_a = nn.Sequential(*attn_a) - self.attention_b = nn.Sequential(*attn_b) - self.attention_c = nn.Linear(hidden_dim, n_classes) - self.init = init - _init_linears(self, init) - - def forward(self, x): - h = self.projection(x) - logits = self.attention_c(self.attention_a(h) * self.attention_b(h)) - return logits, x - - def attention_space(self, x: torch.Tensor) -> torch.Tensor: - """The ``hidden_dim``-d pre-tanh activation ``V · LeakyReLU(W_proj x)``. - - Computed without dropout (``projection[0]`` and ``attention_a[0]`` only), so - it is deterministic in eval and train mode alike. - """ - h = self.projection[1](self.projection[0](x)) - return self.attention_a[0](h) - - def head_spaces(self, x): - """Intermediate activations of the attention head. - - ``a_pre_tanh`` is ``V * LeakyReLU(W_proj x)``, the argument of the tanh - gate: low-dimensional, directly upstream of the attention weights, and - not yet collapsed by tanh saturation. - """ - h = self.projection(x) - pre_tanh = self.attention_a[0](h) - return { - "a_pre_tanh": pre_tanh, - "gate_product": self.attention_a(h) * self.attention_b(h), - } diff --git a/src/grass_mil/models/components/backbones.py b/src/grass_mil/models/components/backbones.py deleted file mode 100644 index a49e9f3..0000000 --- a/src/grass_mil/models/components/backbones.py +++ /dev/null @@ -1,267 +0,0 @@ -from __future__ import annotations - -import warnings -from dataclasses import dataclass -from typing import Any, Dict, Literal, Optional - -import torch -from torch import nn - -from .embeddings import CategoricalEmbeddingConfig, NodeInputEmbedding -from .pooling import GraphPooling - -try: - from torch_geometric.nn import GATConv, GCNConv, GINConv, GINEConv, SAGEConv -except Exception as exc: # pragma: no cover - raise ImportError( - "torch_geometric is required for grass_mil.models.components.backbones" - ) from exc - - -NormType = Literal["batchnorm", "layernorm", "none"] -JKType = Literal["last", "concat", "max", "sum"] -ConvType = Literal["gin", "gcn", "gat", "graphsage", "gine"] - - -@dataclass -class EncoderConfig: - input_dim: int - hidden_dim: int = 128 - out_dim: Optional[int] = None - num_layers: int = 3 - dropout: float = 0.0 - conv_type: ConvType = "gin" - norm: NormType = "batchnorm" - jk: JKType = "last" - act: Literal["relu", "gelu", "leaky_relu"] = "relu" - pooling: Optional[str] = None - gat_heads: int = 4 - use_edge_attr: bool = False - edge_weight_index: int = 0 - edge_attr_dim: Optional[int] = None - # Selects a single scalar column of `edge_attr` (the edge length) for - # edge-aware message passing. When set, the conv sees only that column and - # `edge_attr_dim` is 1. Leave unset to pass the whole edge_attr vector. - edge_feature_index: Optional[int] = None - categorical_embedding: Optional[Dict[str, Any]] = None - # GIN/GINE epsilon: learnable (the GRASS-MIL preprint trains it) or fixed at 0. - train_eps: bool = True - # Initialisation of the input projection (embedding table + linear): "default" - # keeps the torch defaults, "xavier" applies xavier_uniform_ to both (preprint). - input_init: Literal["default", "xavier"] = "xavier" - - -def _get_activation(name: str) -> nn.Module: - if name == "relu": - return nn.ReLU() - if name == "gelu": - return nn.GELU() - if name == "leaky_relu": - return nn.LeakyReLU() - raise ValueError(f"Unsupported activation: {name}") - - -def _get_norm(name: NormType, dim: int) -> nn.Module: - if name == "batchnorm": - return nn.BatchNorm1d(dim) - if name == "layernorm": - return nn.LayerNorm(dim) - if name == "none": - return nn.Identity() - raise ValueError(f"Unsupported norm type: {name}") - - -class GNNEncoder(nn.Module): - """PyG-first encoder with thin adapters for consistent edge handling.""" - - def __init__(self, cfg: EncoderConfig): - super().__init__() - if cfg.num_layers < 1: - raise ValueError("num_layers must be >= 1") - self.cfg = cfg - self._validate_edge_attr_config() - self.act = _get_activation(cfg.act) - self.dropout = nn.Dropout(cfg.dropout) - - self.input_proj = NodeInputEmbedding( - input_dim=cfg.input_dim, - hidden_dim=cfg.hidden_dim, - categorical=CategoricalEmbeddingConfig.from_dict(cfg.categorical_embedding), - init=cfg.input_init, - ) - - self.layers = nn.ModuleList([self._build_conv_layer() for _ in range(cfg.num_layers)]) - self.norms = nn.ModuleList( - [_get_norm(cfg.norm, cfg.hidden_dim) for _ in range(cfg.num_layers)] - ) - - combined_dim = self._combined_node_dim() - final_dim = cfg.out_dim if cfg.out_dim is not None else combined_dim - self.output_proj = ( - nn.Identity() if final_dim == combined_dim else nn.Linear(combined_dim, final_dim) - ) - self.output_dim = final_dim - self.graph_pool = ( - GraphPooling( - name=cfg.pooling, - input_dim=final_dim, - ) - if cfg.pooling is not None - else None - ) - - def _build_conv_layer(self) -> nn.Module: - ctype = self.cfg.conv_type - hdim = self.cfg.hidden_dim - if ctype == "gin": - mlp = nn.Sequential( - nn.Linear(hdim, hdim * 2), - nn.ReLU(), - nn.Linear(hdim * 2, hdim), - ) - return GINConv(mlp, train_eps=self.cfg.train_eps) - if ctype == "gcn": - return GCNConv(hdim, hdim, add_self_loops=True, normalize=True) - if ctype == "gat": - if hdim % self.cfg.gat_heads != 0: - raise ValueError("hidden_dim must be divisible by gat_heads") - out_per_head = hdim // self.cfg.gat_heads - return GATConv( - in_channels=hdim, - out_channels=out_per_head, - heads=self.cfg.gat_heads, - concat=True, - edge_dim=None, - add_self_loops=True, - ) - if ctype == "graphsage": - return SAGEConv(hdim, hdim, aggr="mean") - if ctype == "gine": - if self.cfg.edge_attr_dim is None: - raise ValueError("conv_type='gine' requires edge_attr_dim in EncoderConfig.") - mlp = nn.Sequential( - nn.Linear(hdim, hdim * 2), - nn.ReLU(), - nn.Linear(hdim * 2, hdim), - ) - return GINEConv(mlp, edge_dim=self.cfg.edge_attr_dim, train_eps=self.cfg.train_eps) - raise ValueError(f"Unsupported conv_type: {ctype}") - - def _supports_edge_attr(self) -> bool: - return self.cfg.conv_type in {"gcn", "gine"} - - def _validate_edge_attr_config(self) -> None: - if self.cfg.conv_type == "gine" and self.cfg.edge_attr_dim is None: - raise ValueError("conv_type='gine' requires edge_attr_dim to be set in EncoderConfig.") - if self.cfg.edge_feature_index is not None: - if self.cfg.edge_feature_index < 0: - raise ValueError( - f"edge_feature_index must be >= 0, got {self.cfg.edge_feature_index}." - ) - if self.cfg.conv_type != "gine": - raise ValueError( - "edge_feature_index selects a scalar edge feature for GINE message " - f"passing, but conv_type='{self.cfg.conv_type}'. Use conv_type='gine', " - "or 'gcn' with edge_weight_index for scalar edge weights." - ) - if self.cfg.edge_attr_dim != 1: - raise ValueError( - "edge_feature_index selects exactly one edge column, so edge_attr_dim " - f"must be 1; got {self.cfg.edge_attr_dim}." - ) - if self.cfg.use_edge_attr and not self._supports_edge_attr(): - warnings.warn( - f"use_edge_attr=True is ignored for conv_type='{self.cfg.conv_type}'. " - "Choose 'gcn' (scalar edge_weight) or 'gine' (edge_attr) for edge-aware message passing.", - stacklevel=2, - ) - - def _combined_node_dim(self) -> int: - if self.cfg.jk == "concat": - return self.cfg.hidden_dim * (self.cfg.num_layers + 1) - return self.cfg.hidden_dim - - def _apply_conv( - self, - layer: nn.Module, - x: torch.Tensor, - edge_index: torch.Tensor, - edge_attr: Optional[torch.Tensor], - ) -> torch.Tensor: - if self.cfg.conv_type == "gcn" and self.cfg.use_edge_attr: - if edge_attr is None: - raise ValueError( - "use_edge_attr=True with conv_type='gcn' requires edge_attr in forward inputs." - ) - if edge_attr.dim() != 2: - raise ValueError("edge_attr must have shape [num_edges, num_edge_features].") - if self.cfg.edge_weight_index >= edge_attr.size(1): - raise ValueError("edge_weight_index is out of range for provided edge_attr.") - edge_weight = edge_attr[:, self.cfg.edge_weight_index] - return layer(x, edge_index, edge_weight=edge_weight) - if self.cfg.conv_type == "gine": - if edge_attr is None: - raise ValueError("conv_type='gine' requires edge_attr in forward inputs.") - return layer(x, edge_index, self._select_edge_features(edge_attr)) - return layer(x, edge_index) - - def _select_edge_features(self, edge_attr: torch.Tensor) -> torch.Tensor: - """Narrow ``edge_attr`` to the configured scalar column. - - The production encoder conditions message passing on the edge *length* - alone, so only that column reaches the convolution. Without - ``edge_feature_index`` the whole edge-attribute vector is passed - through unchanged. - """ - index = self.cfg.edge_feature_index - if index is None: - return edge_attr - if edge_attr.dim() != 2: - raise ValueError("edge_attr must have shape [num_edges, num_edge_features].") - if index < 0 or index >= edge_attr.size(1): - raise ValueError( - f"edge_feature_index={index} is out of range for edge_attr with " - f"{edge_attr.size(1)} column(s)." - ) - return edge_attr[:, index : index + 1].float() - - def _apply_jk(self, states: list[torch.Tensor]) -> torch.Tensor: - mode = self.cfg.jk - if mode == "last": - return states[-1] - if mode == "concat": - return torch.cat(states, dim=-1) - if mode == "max": - return torch.stack(states, dim=0).max(dim=0).values - if mode == "sum": - return torch.stack(states, dim=0).sum(dim=0) - raise ValueError(f"Unsupported JK mode: {mode}") - - def forward( - self, - x: torch.Tensor, - edge_index: torch.Tensor, - edge_attr: Optional[torch.Tensor] = None, - batch: Optional[torch.Tensor] = None, - categorical_codes: Optional[torch.Tensor] = None, - return_graph_embedding: bool = False, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - h = self.input_proj(x, categorical_codes) - states = [h] - for idx, layer in enumerate(self.layers): - h = self._apply_conv(layer, h, edge_index, edge_attr) - h = self.norms[idx](h) - if idx != len(self.layers) - 1: - h = self.act(h) - h = self.dropout(h) - states.append(h) - - node_emb = self.output_proj(self._apply_jk(states)) - if not return_graph_embedding: - return node_emb - if self.graph_pool is None: - raise ValueError("return_graph_embedding=True requires cfg.pooling to be set") - if batch is None: - batch = torch.zeros(node_emb.size(0), dtype=torch.long, device=node_emb.device) - graph_emb = self.graph_pool(node_emb, batch) - return node_emb, graph_emb diff --git a/src/grass_mil/models/components/embeddings.py b/src/grass_mil/models/components/embeddings.py deleted file mode 100644 index 8408435..0000000 --- a/src/grass_mil/models/components/embeddings.py +++ /dev/null @@ -1,280 +0,0 @@ -"""Input embeddings for node features. - -The published encoder consumes cell identity through a learned embedding table -rather than as a one-hot column of ``x``: - -.. math:: h_v^{(0)} = E[t_v] + W_x x_v - -with :math:`t_v` the categorical cell-type code and :math:`x_v` the continuous -node features (cell size, where available). The two terms are **summed**, not -concatenated, so the embedding dimension is the hidden dimension. - -This matters for the cohorts whose only node input is cell type: there ``x`` has -zero columns and the embedding is the entire input representation. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Dict, Mapping, Optional - -import torch -from torch import nn - -__all__ = ["CategoricalEmbeddingConfig", "NodeInputEmbedding"] - - -@dataclass(frozen=True) -class CategoricalEmbeddingConfig: - """Selects one categorical node label to embed. - - Attributes: - label: Key into ``data.categorical_slices`` naming the column of - ``data.categorical_codes`` to embed (typically ``cell_type``). - num_embeddings: Vocabulary size. ``None`` defers to the cardinality - recorded by precompute. - column_index: Column of ``categorical_codes``. ``None`` defers to - ``categorical_slices[label]``. - reserve_unassigned: Reserve one extra row so augmentations can mark a - dropped node as "unassigned" without resizing the table. The - reserved code is the last row. - unassigned_label: Alternatively, name an EXISTING category (for example - ``"other"``) as the unassigned token; augmentations then write that - category's code instead of a reserved row. Resolved to - ``unassigned_index`` from the datamodule label map. Mutually - exclusive with ``reserve_unassigned``. - unassigned_index: Explicit code of the unassigned token (set directly, - or resolved from ``unassigned_label``). - """ - - label: str = "cell_type" - num_embeddings: Optional[int] = None - column_index: Optional[int] = None - reserve_unassigned: bool = True - unassigned_label: Optional[str] = None - unassigned_index: Optional[int] = None - - @classmethod - def from_dict(cls, cfg: Optional[Mapping[str, Any]]) -> Optional["CategoricalEmbeddingConfig"]: - if cfg is None: - return None - allowed = { - "label", - "num_embeddings", - "column_index", - "reserve_unassigned", - "unassigned_label", - "unassigned_index", - } - unknown = set(cfg) - allowed - if unknown: - raise ValueError( - f"Unknown categorical embedding keys: {sorted(unknown)}. " - f"Expected a subset of {sorted(allowed)}." - ) - return cls(**dict(cfg)) - - -class NodeInputEmbedding(nn.Module): - """Project node inputs to the hidden dimension. - - With no categorical config this reproduces the plain input projection: an - ``Identity`` when the dimensions already agree, otherwise a ``Linear``. - """ - - def __init__( - self, - *, - input_dim: int, - hidden_dim: int, - categorical: Optional[CategoricalEmbeddingConfig] = None, - init: str = "default", - ) -> None: - super().__init__() - if input_dim < 0: - raise ValueError(f"input_dim must be >= 0, got {input_dim}.") - if init not in {"default", "xavier"}: - raise ValueError(f"init must be 'default' or 'xavier', got {init!r}.") - self.input_dim = int(input_dim) - self.hidden_dim = int(hidden_dim) - self.categorical = categorical - self.init = init - - if self.input_dim == 0: - # Cell-type-only cohorts: precompute emits an empty feature matrix - # and the embedding carries the whole representation. - if categorical is None: - raise ValueError( - "input_dim=0 requires a categorical embedding; the encoder would " - "otherwise receive no input features at all." - ) - self.linear: nn.Module = nn.Identity() - self._uses_linear = False - else: - self.linear = ( - nn.Identity() - if self.input_dim == self.hidden_dim - else nn.Linear(self.input_dim, self.hidden_dim) - ) - self._uses_linear = True - - self.embedding: Optional[nn.Embedding] = None - self._column_index: Optional[int] = None - self._unassigned_index: Optional[int] = None - if categorical is not None: - if categorical.num_embeddings is None: - raise ValueError( - "categorical.num_embeddings is unset. Resolve it from the datamodule " - "metadata (categorical_cardinalities) or set it explicitly in config." - ) - if categorical.num_embeddings < 1: - raise ValueError( - f"categorical.num_embeddings must be >= 1, got {categorical.num_embeddings}." - ) - if categorical.reserve_unassigned and ( - categorical.unassigned_label is not None - or categorical.unassigned_index is not None - ): - raise ValueError( - "categorical.reserve_unassigned=True cannot be combined with " - "unassigned_label/unassigned_index: either reserve an extra row or name an " - "existing category as the unassigned token." - ) - num_rows = int(categorical.num_embeddings) + int(categorical.reserve_unassigned) - self.embedding = nn.Embedding(num_rows, self.hidden_dim) - self._column_index = ( - 0 if categorical.column_index is None else int(categorical.column_index) - ) - if self._column_index < 0: - raise ValueError( - f"categorical.column_index must be >= 0, got {self._column_index}." - ) - if categorical.reserve_unassigned: - self._unassigned_index = num_rows - 1 - elif categorical.unassigned_index is not None: - idx = int(categorical.unassigned_index) - if idx < 0 or idx >= num_rows: - raise ValueError( - f"categorical.unassigned_index={idx} is out of range for an embedding " - f"with {num_rows} rows." - ) - self._unassigned_index = idx - elif categorical.unassigned_label is not None: - raise ValueError( - f"categorical.unassigned_label={categorical.unassigned_label!r} was not " - "resolved to a code; use resolve_categorical_embedding_config with the " - "datamodule label map, or set unassigned_index." - ) - self.reset_parameters() - - def reset_parameters(self) -> None: - if self.init == "xavier": - if self.embedding is not None: - nn.init.xavier_uniform_(self.embedding.weight) - if isinstance(self.linear, nn.Linear): - # Weight only, as in the reference implementation; the bias keeps - # the torch default. - nn.init.xavier_uniform_(self.linear.weight) - - @property - def unassigned_index(self) -> Optional[int]: - """Reserved code marking a dropped node, or ``None`` when not reserved.""" - return self._unassigned_index - - def extra_repr(self) -> str: - label = self.categorical.label if self.categorical is not None else None - return f"input_dim={self.input_dim}, hidden_dim={self.hidden_dim}, categorical={label}" - - def forward( - self, - x: Optional[torch.Tensor], - categorical_codes: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - if self.embedding is None: - if x is None: - raise ValueError("NodeInputEmbedding requires `x` when no categorical embedding.") - return self.linear(x) - - if categorical_codes is None: - raise ValueError( - f"A categorical embedding on '{self.categorical.label}' is configured but the " - "batch carries no `categorical_codes`. Enable the categorical feature in the " - "datamodule, or unset model.encoder.categorical_embedding." - ) - if categorical_codes.dim() != 2: - raise ValueError( - "categorical_codes must have shape [num_nodes, num_labels]; got " - f"{tuple(categorical_codes.shape)}." - ) - if self._column_index >= categorical_codes.size(1): - raise ValueError( - f"categorical column_index={self._column_index} is out of range for " - f"categorical_codes with {categorical_codes.size(1)} column(s)." - ) - - codes = categorical_codes[:, self._column_index].long() - num_rows = int(self.embedding.num_embeddings) - if codes.numel() and (int(codes.min()) < 0 or int(codes.max()) >= num_rows): - # Fail loudly rather than clamping: an out-of-range code means the - # vocabulary and the data disagree, and silently folding it into a - # valid row would train on a wrong cell type. - raise ValueError( - f"Categorical code out of range for '{self.categorical.label}': got " - f"[{int(codes.min())}, {int(codes.max())}] for an embedding with " - f"{num_rows} rows." - ) - embedded = self.embedding(codes) - - if not self._uses_linear or x is None or x.size(1) == 0: - return embedded - return self.linear(x) + embedded - - -def resolve_categorical_embedding_config( - cfg: Optional[Mapping[str, Any]], - *, - cardinalities: Optional[Dict[str, int]] = None, - slices: Optional[Mapping[str, int]] = None, - label_maps: Optional[Mapping[str, Mapping[str, int]]] = None, -) -> Optional[CategoricalEmbeddingConfig]: - """Fill in vocabulary size, column index and unassigned code from datamodule metadata. - - Explicit config values always win; ``None`` triggers inference, mirroring - the ``encoder.input_dim`` contract. ``label_maps`` (``{label: {category: code}}``) - is used to turn ``unassigned_label`` into ``unassigned_index``. - """ - resolved = CategoricalEmbeddingConfig.from_dict(cfg) - if resolved is None: - return None - - num_embeddings = resolved.num_embeddings - if num_embeddings is None and cardinalities: - num_embeddings = cardinalities.get(resolved.label) - column_index = resolved.column_index - if column_index is None and slices: - column_index = slices.get(resolved.label) - - if num_embeddings is None: - raise ValueError( - f"Could not resolve the vocabulary size for categorical label " - f"'{resolved.label}'. Set model.encoder.categorical_embedding.num_embeddings " - "explicitly, or include the label in data.categorical_features.include_labels." - ) - unassigned_index = resolved.unassigned_index - if unassigned_index is None and resolved.unassigned_label is not None: - mapping = (label_maps or {}).get(resolved.label) - if not mapping or resolved.unassigned_label not in mapping: - raise ValueError( - f"categorical.unassigned_label={resolved.unassigned_label!r} is not a category " - f"of '{resolved.label}'. Known categories: " - f"{sorted(mapping) if mapping else 'none (no label map available)'}." - ) - unassigned_index = int(mapping[resolved.unassigned_label]) - return CategoricalEmbeddingConfig( - label=resolved.label, - num_embeddings=int(num_embeddings), - column_index=None if column_index is None else int(column_index), - reserve_unassigned=resolved.reserve_unassigned, - unassigned_label=resolved.unassigned_label, - unassigned_index=unassigned_index, - ) diff --git a/src/grass_mil/models/components/factory.py b/src/grass_mil/models/components/factory.py deleted file mode 100644 index 557aab7..0000000 --- a/src/grass_mil/models/components/factory.py +++ /dev/null @@ -1,105 +0,0 @@ -from __future__ import annotations - -from typing import Any, Dict, Optional - -import torch - -from .attention import AttnNetGated, AttnNetGatedProjected -from .backbones import EncoderConfig, GNNEncoder -from .heads import GraphPredictionHead -from .losses import ( - CoxSGDLoss, - WeightedBCEWithLogitsLoss, - WeightedCrossEntropyLoss, - WeightedHuberLoss, - WeightedMSELoss, -) -from .pooling import GraphPooling -from .ssl import BGRL, MLPPredictor - - -def build_encoder(config: Dict[str, Any] | EncoderConfig) -> GNNEncoder: - if isinstance(config, EncoderConfig): - return GNNEncoder(config) - return GNNEncoder(EncoderConfig(**config)) - - -def build_pooling(config: Dict[str, Any]) -> GraphPooling: - return GraphPooling(**config) - - -def build_attention( - use_attention: bool, attention_type: str = "gated", **kwargs: Any -) -> Optional[object]: - if not use_attention: - return None - if attention_type == "gated": - return AttnNetGated(**kwargs) - if attention_type == "gated_projected": - return AttnNetGatedProjected(**kwargs) - raise ValueError(f"Unsupported attention type: {attention_type}") - - -def build_graph_head(config: Dict[str, Any]) -> GraphPredictionHead: - return GraphPredictionHead(**config) - - -def build_ssl(use_ssl: bool, encoder, ssl_config: Optional[Dict[str, Any]] = None): - if not use_ssl: - return None - cfg = ssl_config or {} - method = cfg.get("method", "bgrl") - if method != "bgrl": - raise ValueError(f"Unsupported SSL method: {method}") - predictor_cfg = cfg.get("predictor") - if predictor_cfg is None: - raise ValueError( - "SSL predictor configuration is required when use_ssl=True " - "(expected key: ssl.predictor)." - ) - if "hidden_size" not in predictor_cfg: - raise ValueError("ssl.predictor.hidden_size is required when use_ssl=True.") - predictor_cfg = dict(predictor_cfg) - predictor_cfg.pop("_target_", None) - predictor_cfg.pop("_recursive_", None) - predictor_cfg.pop("_convert_", None) - predictor_cfg.setdefault("input_size", encoder.output_dim) - predictor_cfg.setdefault("output_size", encoder.output_dim) - predictor = MLPPredictor(**predictor_cfg) - return BGRL(encoder=encoder, predictor=predictor) - - -def _as_weight_tensor(value): - """Accept tensors, lists/tuples and OmegaConf sequences for loss weights.""" - if value is None or isinstance(value, torch.Tensor): - return value - try: - from omegaconf import ListConfig, OmegaConf - - if isinstance(value, ListConfig): - value = OmegaConf.to_container(value, resolve=True) - except ImportError: # pragma: no cover - pass - if isinstance(value, (list, tuple)): - return torch.as_tensor([float(v) for v in value], dtype=torch.float32) - if isinstance(value, (int, float)): - return torch.as_tensor([float(value)], dtype=torch.float32) - raise TypeError(f"Loss weight must be a tensor or a sequence of floats, got {type(value)!r}.") - - -def build_loss(loss_type: str, **kwargs): - name = loss_type.lower() - for key in ("class_weight", "pos_weight"): - if key in kwargs: - kwargs[key] = _as_weight_tensor(kwargs[key]) - if name == "categorical_ce": - return WeightedCrossEntropyLoss(**kwargs) - if name == "categorical_bce": - return WeightedBCEWithLogitsLoss(**kwargs) - if name == "regression_mse": - return WeightedMSELoss(**kwargs) - if name == "regression_huber": - return WeightedHuberLoss(**kwargs) - if name == "survival_coxsgd": - return CoxSGDLoss(**kwargs) - raise ValueError(f"Unsupported loss type: {loss_type}") diff --git a/src/grass_mil/models/components/heads.py b/src/grass_mil/models/components/heads.py deleted file mode 100644 index 8c5aaec..0000000 --- a/src/grass_mil/models/components/heads.py +++ /dev/null @@ -1,34 +0,0 @@ -from __future__ import annotations - -from torch import nn - - -class _MLPHead(nn.Module): - def __init__( - self, - input_dim: int, - output_dim: int, - hidden_dim: int | None = None, - num_layers: int = 2, - dropout: float = 0.0, - ): - super().__init__() - if num_layers < 1: - raise ValueError("num_layers must be >= 1") - hidden = hidden_dim or input_dim - layers = [] - curr_dim = input_dim - for _ in range(num_layers - 1): - layers.extend([nn.Linear(curr_dim, hidden), nn.LeakyReLU()]) - if dropout > 0: - layers.append(nn.Dropout(dropout)) - curr_dim = hidden - layers.append(nn.Linear(curr_dim, output_dim)) - self.net = nn.Sequential(*layers) - - def forward(self, x): - return self.net(x) - - -class GraphPredictionHead(_MLPHead): - pass diff --git a/src/grass_mil/models/components/losses.py b/src/grass_mil/models/components/losses.py deleted file mode 100644 index 4b6813b..0000000 --- a/src/grass_mil/models/components/losses.py +++ /dev/null @@ -1,145 +0,0 @@ -from __future__ import annotations - -from typing import Optional - -import torch -from torch import nn - - -def _reshape_sample_weight( - sample_weight: Optional[torch.Tensor], - target_shape: torch.Size, - device: torch.device, -) -> Optional[torch.Tensor]: - if sample_weight is None: - return None - w = sample_weight.to(device) - if w.ndim == 1 and len(target_shape) > 1: - w = w.unsqueeze(-1) - return w - - -class WeightedCrossEntropyLoss(nn.Module): - def __init__(self, class_weight: Optional[torch.Tensor] = None): - super().__init__() - self.class_weight = class_weight - - def forward( - self, - logits: torch.Tensor, - target: torch.Tensor, - sample_weight: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - if logits.ndim != 2: - raise ValueError("WeightedCrossEntropyLoss expects logits with shape [N, C]") - target = target.long().view(-1) - class_weight = ( - self.class_weight.to(logits.device) if self.class_weight is not None else None - ) - loss = nn.functional.cross_entropy(logits, target, weight=class_weight, reduction="none") - if sample_weight is not None: - sw = sample_weight.view(-1).to(loss.device) - loss = loss * sw - return loss.mean() - - -class WeightedBCEWithLogitsLoss(nn.Module): - def __init__(self, pos_weight: Optional[torch.Tensor] = None): - super().__init__() - self.pos_weight = pos_weight - - def forward( - self, - logits: torch.Tensor, - target: torch.Tensor, - sample_weight: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - if logits.shape != target.shape: - raise ValueError( - f"BCE logits and target must have identical shape, got {logits.shape} and {target.shape}" - ) - pos_weight = self.pos_weight.to(logits.device) if self.pos_weight is not None else None - loss = nn.functional.binary_cross_entropy_with_logits( - logits, target.float(), pos_weight=pos_weight, reduction="none" - ) - sw = _reshape_sample_weight(sample_weight, target.shape, logits.device) - if sw is not None: - loss = loss * sw - return loss.mean() - - -class WeightedMSELoss(nn.Module): - def forward( - self, - pred: torch.Tensor, - target: torch.Tensor, - sample_weight: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - loss = (pred - target) ** 2 - sw = _reshape_sample_weight(sample_weight, target.shape, pred.device) - if sw is not None: - loss = loss * sw - return loss.mean() - - -class WeightedHuberLoss(nn.Module): - def __init__(self, delta: float = 1.0): - super().__init__() - self.delta = delta - - def forward( - self, - pred: torch.Tensor, - target: torch.Tensor, - sample_weight: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - loss = nn.functional.huber_loss(pred, target, delta=self.delta, reduction="none") - sw = _reshape_sample_weight(sample_weight, target.shape, pred.device) - if sw is not None: - loss = loss * sw - return loss.mean() - - -class CoxSGDLoss(nn.Module): - """Pairwise Cox-style objective with optional random top-n pair pruning.""" - - def __init__(self, top_n: int = 2, regularizer_weight: float = 0.05): - super().__init__() - self.top_n = top_n - self.regularizer_weight = regularizer_weight - - def forward( - self, y_pred: torch.Tensor, length: torch.Tensor, event: torch.Tensor - ) -> torch.Tensor: - y_pred = y_pred.view(-1) - length = length.view(-1) - event = event.view(-1) - if not (y_pred.shape[0] == length.shape[0] == event.shape[0]): - raise ValueError("y_pred, length, and event must have same first dimension") - - n_samples = y_pred.shape[0] - pair_mat = (length.view(1, -1) - length.view(-1, 1) > 0) * event.view(-1, 1) - - if self.top_n > 0 and n_samples > 1: - top_k = min(int(self.top_n), n_samples - 1) - if top_k <= 0: - top_k = 1 - noise = 1 + torch.rand_like(pair_mat.float()) - p_with_rand = pair_mat.float() * noise - rand_thr_ind = torch.argsort(p_with_rand, dim=1)[:, -(top_k + 1)] - rand_thr = p_with_rand[torch.arange(n_samples), rand_thr_ind].view(-1, 1) - pair_mat = pair_mat * (p_with_rand > rand_thr) - - valid_sample_is = torch.nonzero(pair_mat.sum(1)).flatten() - if valid_sample_is.numel() == 0: - return torch.zeros((), device=y_pred.device, dtype=y_pred.dtype) - - pair_mat = pair_mat.clone() - pair_mat[(valid_sample_is, valid_sample_is)] = 1 - - score_diff = y_pred.view(1, -1) - y_pred.view(-1, 1) - row_max = torch.max(score_diff, dim=1, keepdim=True).values - exp_terms = torch.exp(score_diff - row_max) * pair_mat - loss = row_max[:, 0][valid_sample_is] + torch.log(exp_terms.sum(1)[valid_sample_is]) - regularizer = torch.abs(pair_mat.sum(0) * y_pred).sum() - return loss.sum() + self.regularizer_weight * regularizer diff --git a/src/grass_mil/models/components/pooling.py b/src/grass_mil/models/components/pooling.py deleted file mode 100644 index 46b7e9e..0000000 --- a/src/grass_mil/models/components/pooling.py +++ /dev/null @@ -1,44 +0,0 @@ -from __future__ import annotations - -from typing import Optional - -from torch import nn - -try: - from torch_geometric.nn import global_add_pool, global_max_pool, global_mean_pool -except Exception as exc: # pragma: no cover - raise ImportError( - "torch_geometric is required for grass_mil.models.components.pooling" - ) from exc - -_SUPPORTED = ("sum", "mean", "max") - - -class GraphPooling(nn.Module): - """Pool node embeddings into one vector per graph. - - Only order-invariant reductions are offered. Learned pooling - (``GlobalAttention``, ``Set2Set``) was removed: instance-level selection is - the job of the MIL attention head, which pools *across* ego-graphs, and - duplicating it inside each ego-graph adds parameters without a - corresponding claim. ``GlobalAttention`` is also deprecated upstream. - """ - - def __init__(self, name: Optional[str], input_dim: int, **_: object): - super().__init__() - self.name = (name or "mean").lower() - self.input_dim = input_dim - self.output_dim = input_dim - if self.name == "sum": - self.pool = global_add_pool - elif self.name == "mean": - self.pool = global_mean_pool - elif self.name == "max": - self.pool = global_max_pool - else: - raise ValueError( - f"Unsupported graph pooling '{name}'. Expected one of {list(_SUPPORTED)}." - ) - - def forward(self, x, batch): - return self.pool(x, batch) diff --git a/src/grass_mil/models/components/ssl.py b/src/grass_mil/models/components/ssl.py deleted file mode 100644 index d8d5b6e..0000000 --- a/src/grass_mil/models/components/ssl.py +++ /dev/null @@ -1,67 +0,0 @@ -from __future__ import annotations - -import copy - -import torch -from torch import nn - - -class MLPPredictor(nn.Module): - def __init__(self, input_size: int, output_size: int, hidden_size: int = 512): - super().__init__() - self.net = nn.Sequential( - nn.Linear(input_size, hidden_size, bias=True), - nn.PReLU(1), - nn.Linear(hidden_size, output_size, bias=True), - ) - self.reset_parameters() - - def forward(self, x): - return self.net(x) - - def reset_parameters(self): - for module in self.modules(): - if isinstance(module, nn.Linear): - module.reset_parameters() - - -class BGRL(nn.Module): - """Bootstrap Graph Representation Learning wrapper.""" - - def __init__(self, encoder: nn.Module, predictor: nn.Module): - super().__init__() - self.online_encoder = encoder - self.predictor = predictor - self.target_encoder = copy.deepcopy(encoder) - - if hasattr(self.target_encoder, "reset_parameters"): - self.target_encoder.reset_parameters() - for param in self.target_encoder.parameters(): - param.requires_grad = False - - def trainable_parameters(self): - return list(self.online_encoder.parameters()) + list(self.predictor.parameters()) - - @torch.no_grad() - def update_target_network(self, momentum: float): - if momentum < 0.0 or momentum > 1.0: - raise ValueError(f"Momentum must be in [0,1], got {momentum}") - for online_p, target_p in zip( - self.online_encoder.parameters(), self.target_encoder.parameters() - ): - target_p.data.mul_(momentum).add_(online_p.data, alpha=1.0 - momentum) - - def forward(self, online_x, target_x): - online_y = ( - self.online_encoder(*online_x) - if isinstance(online_x, tuple) - else self.online_encoder(online_x) - ) - online_q = self.predictor(online_y) - with torch.no_grad(): - target_y = ( - self.target_encoder(*target_x) - if isinstance(target_x, tuple) - else self.target_encoder(target_x) - ).detach() - return online_q, target_y diff --git a/src/grass_mil/models/encoder.py b/src/grass_mil/models/encoder.py new file mode 100644 index 0000000..f8c1355 --- /dev/null +++ b/src/grass_mil/models/encoder.py @@ -0,0 +1,177 @@ +"""The ego-graph encoder: GINE message passing conditioned on edge length. + +Node inputs are a learned cell-type embedding summed with a linear projection +of the continuous node features, + +.. math:: h_v^{(0)} = E[t_v] + W x_v, + +followed by ``num_layers`` GINE layers (norm, activation and dropout between +layers) and a pooling of the node states into one vector per ego-graph. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import Any, Mapping + +import torch +from torch import nn +from torch_geometric.nn import GINEConv, global_add_pool, global_max_pool, global_mean_pool + +__all__ = ["EncoderConfig", "GINEEncoder"] + +_POOL = {"max": global_max_pool, "mean": global_mean_pool, "sum": global_add_pool} +_ACT = {"relu": nn.ReLU, "gelu": nn.GELU, "leaky_relu": nn.LeakyReLU} +_NORM = {"layernorm": nn.LayerNorm, "batchnorm": nn.BatchNorm1d} + + +@dataclass(frozen=True) +class EncoderConfig: + """Encoder architecture (input widths come from the data, not the config). + + Attributes: + hidden_dim: Width of every GINE layer. + out_dim: Width of the ego-graph embedding (a final linear map when it + differs from the jumping-knowledge output). + num_layers: GINE layers (match ``data.sampler.depth``). + dropout: Dropout after every layer. + norm: ``layernorm`` | ``batchnorm`` | ``none``. + jk: Jumping knowledge over layer outputs: ``last`` | ``concat`` | + ``max`` | ``sum``. + act: ``relu`` | ``gelu`` | ``leaky_relu``. + pooling: Ego-graph readout: ``max`` | ``mean`` | ``sum``. + mlp_ratio: Hidden width of each GINE MLP, as a multiple of ``hidden_dim``. + train_eps: Learn the GIN epsilon. + edge_feature: Column of ``edge_attr`` the convolutions see (``None``: + every column). + input_init: ``xavier`` or ``default`` initialisation of the input + embedding and projection. + unassigned_cell_type: Cell type used to mark a masked node in + augmentations; ``None`` reserves an extra embedding row for it. + """ + + hidden_dim: int = 256 + out_dim: int = 256 + num_layers: int = 4 + dropout: float = 0.0 + norm: str = "layernorm" + jk: str = "last" + act: str = "relu" + pooling: str = "max" + mlp_ratio: int = 2 + train_eps: bool = True + edge_feature: int | None = 0 + input_init: str = "xavier" + unassigned_cell_type: str | None = None + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any] | None) -> EncoderConfig: + cfg = dict(cfg or {}) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown model.encoder keys: {sorted(unknown)}") + out = cls(**cfg) + checks = { + "norm": (out.norm, {"layernorm", "batchnorm", "none"}), + "jk": (out.jk, {"last", "concat", "max", "sum"}), + "act": (out.act, set(_ACT)), + "pooling": (out.pooling, set(_POOL)), + "input_init": (out.input_init, {"xavier", "default"}), + } + for key, (value, allowed) in checks.items(): + if value not in allowed: + raise ValueError(f"model.encoder.{key} must be one of {sorted(allowed)}") + return out + + +class NodeInput(nn.Module): + """``E[cell_type] + W x``; either term is absent when the data lacks it.""" + + def __init__(self, feature_dim: int, cell_types: tuple[str, ...], cfg: EncoderConfig) -> None: + super().__init__() + if feature_dim == 0 and not cell_types: + raise ValueError("The encoder needs cell types or continuous node features.") + self.linear = nn.Linear(feature_dim, cfg.hidden_dim) if feature_dim else None + self.embedding = None + self.unassigned_index = None + if cell_types: + if cfg.unassigned_cell_type is None: + rows, self.unassigned_index = len(cell_types) + 1, len(cell_types) + elif cfg.unassigned_cell_type in cell_types: + rows, self.unassigned_index = len(cell_types), cell_types.index(cfg.unassigned_cell_type) + else: + raise ValueError( + f"model.encoder.unassigned_cell_type '{cfg.unassigned_cell_type}' is not a " + f"cell type of the data ({list(cell_types)})." + ) + self.embedding = nn.Embedding(rows, cfg.hidden_dim) + if cfg.input_init == "xavier": + if self.embedding is not None: + nn.init.xavier_uniform_(self.embedding.weight) + if self.linear is not None: + nn.init.xavier_uniform_(self.linear.weight) + + def forward(self, x: torch.Tensor, cell_type: torch.Tensor | None) -> torch.Tensor: + h = self.linear(x) if self.linear is not None else 0.0 + if self.embedding is not None: + if cell_type is None: + raise ValueError("The encoder embeds cell types but the batch has none.") + h = h + self.embedding(cell_type) + return h + + +class GINEEncoder(nn.Module): + def __init__(self, cfg: EncoderConfig, *, feature_dim: int, cell_types: tuple[str, ...], edge_dim: int) -> None: + super().__init__() + self.cfg = cfg + if cfg.edge_feature is not None and not 0 <= cfg.edge_feature < edge_dim: + raise ValueError(f"model.encoder.edge_feature={cfg.edge_feature} but edges have {edge_dim} feature(s).") + conv_edge_dim = 1 if cfg.edge_feature is not None else edge_dim + h, k = cfg.hidden_dim, cfg.mlp_ratio * cfg.hidden_dim + self.input_proj = NodeInput(feature_dim, cell_types, cfg) + self.layers = nn.ModuleList( + GINEConv(nn.Sequential(nn.Linear(h, k), nn.ReLU(), nn.Linear(k, h)), edge_dim=conv_edge_dim, train_eps=cfg.train_eps) + for _ in range(cfg.num_layers) + ) + self.norms = nn.ModuleList( + _NORM[cfg.norm](h) if cfg.norm != "none" else nn.Identity() for _ in range(cfg.num_layers) + ) + self.act = _ACT[cfg.act]() + self.dropout = nn.Dropout(cfg.dropout) + jk_dim = h * (cfg.num_layers + 1) if cfg.jk == "concat" else h + self.output_proj = nn.Identity() if jk_dim == cfg.out_dim else nn.Linear(jk_dim, cfg.out_dim) + self.output_dim = cfg.out_dim + + @property + def unassigned_index(self) -> int | None: + return self.input_proj.unassigned_index + + def node_embeddings(self, batch) -> torch.Tensor: + """One embedding per node of ``batch``.""" + edge_attr = batch.edge_attr + if self.cfg.edge_feature is not None: + edge_attr = edge_attr[:, self.cfg.edge_feature : self.cfg.edge_feature + 1] + h = self.input_proj(batch.x, getattr(batch, "cell_type", None)) + states = [h] + for i, (conv, norm) in enumerate(zip(self.layers, self.norms)): + h = norm(conv(h, batch.edge_index, edge_attr.float())) + if i < len(self.layers) - 1: + h = self.act(h) + h = self.dropout(h) + states.append(h) + if self.cfg.jk == "last": + nodes = states[-1] + elif self.cfg.jk == "concat": + nodes = torch.cat(states, dim=-1) + else: + stacked = torch.stack(states) + nodes = stacked.max(0).values if self.cfg.jk == "max" else stacked.sum(0) + return self.output_proj(nodes) + + def forward(self, batch) -> torch.Tensor: + """One embedding per ego-graph of ``batch`` (the pooled node embeddings).""" + nodes = self.node_embeddings(batch) + index = getattr(batch, "batch", None) + if index is None: + index = torch.zeros(nodes.size(0), dtype=torch.long, device=nodes.device) + return _POOL[self.cfg.pooling](nodes, index) diff --git a/src/grass_mil/models/mil.py b/src/grass_mil/models/mil.py new file mode 100644 index 0000000..9998650 --- /dev/null +++ b/src/grass_mil/models/mil.py @@ -0,0 +1,143 @@ +"""Instance head, gated attention and logit-space attention pooling. + +For a bag (unit) with instances ``i`` and classes ``c``, the instance head +gives logits ``l[i, c]``, the attention head gives scores ``a[i, c]``, and + +.. math:: A_{i,c} = \\mathrm{softmax}_i(a_{i,c}), \\qquad L_c = \\sum_i A_{i,c}\\, l_{i,c}. + +Pooling in logit space with one attention channel per class makes every bag +logit an exact sum of per-instance contributions ``A[i, c] l[i, c]``. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import Any, Mapping, Sequence + +import torch +from torch import nn + +__all__ = ["HeadConfig", "AttentionConfig", "InstanceHead", "GatedAttention", "attention_pool", "pool_by_bag"] + + +@dataclass(frozen=True) +class HeadConfig: + """Instance head: ``num_layers=1`` is linear, more layers make an MLP.""" + + num_layers: int = 1 + hidden_dim: int | None = None + dropout: float = 0.0 + act: str = "leaky_relu" + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any] | None) -> HeadConfig: + cfg = dict(cfg or {}) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown model.head keys: {sorted(unknown)}") + return cls(**cfg) + + +@dataclass(frozen=True) +class AttentionConfig: + """Projected gated attention (``Linear+LeakyReLU`` projection, tanh/sigmoid gate).""" + + projection_dim: int = 32 + hidden_dim: int = 8 + dropout: float = 0.25 + init: str = "kaiming" + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any] | None) -> AttentionConfig: + cfg = dict(cfg or {}) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown model.attention keys: {sorted(unknown)}") + out = cls(**cfg) + if out.init not in {"kaiming", "default"}: + raise ValueError("model.attention.init must be 'kaiming' or 'default'.") + return out + + +_ACT = {"leaky_relu": nn.LeakyReLU, "relu": nn.ReLU, "gelu": nn.GELU} + + +class InstanceHead(nn.Module): + def __init__(self, input_dim: int, num_classes: int, cfg: HeadConfig) -> None: + super().__init__() + if cfg.num_layers < 1: + raise ValueError("model.head.num_layers must be >= 1") + hidden = cfg.hidden_dim or input_dim + layers: list[nn.Module] = [] + width = input_dim + for _ in range(cfg.num_layers - 1): + layers += [nn.Linear(width, hidden), _ACT[cfg.act]()] + if cfg.dropout > 0: + layers.append(nn.Dropout(cfg.dropout)) + width = hidden + layers.append(nn.Linear(width, num_classes)) + self.net = nn.Sequential(*layers) + + @property + def is_linear(self) -> bool: + return len(self.net) == 1 + + @property + def bias(self) -> torch.Tensor: + return self.net[-1].bias + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +class GatedAttention(nn.Module): + """Raw (pre-softmax) attention scores ``[n_instances, num_classes]``.""" + + def __init__(self, input_dim: int, num_classes: int, cfg: AttentionConfig) -> None: + super().__init__() + p, h = cfg.projection_dim, cfg.hidden_dim + self.projection = nn.Sequential(nn.Linear(input_dim, p), nn.LeakyReLU(), nn.Dropout(cfg.dropout)) + self.attention_a = nn.Sequential(nn.Linear(p, h), nn.Tanh(), nn.Dropout(cfg.dropout)) + self.attention_b = nn.Sequential(nn.Linear(p, h), nn.Sigmoid(), nn.Dropout(cfg.dropout)) + self.attention_c = nn.Linear(h, num_classes) + if cfg.init == "kaiming": + for layer in self.modules(): + if isinstance(layer, nn.Linear): + nn.init.kaiming_uniform_(layer.weight) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.projection(x) + return self.attention_c(self.attention_a(h) * self.attention_b(h)) + + def space(self, x: torch.Tensor) -> torch.Tensor: + """The low-dimensional pre-tanh activation, the attention head's own feature space. + + Computed without dropout, so it is the same in train and eval mode. + """ + return self.attention_a[0](self.projection[1](self.projection[0](x))) + + +def attention_pool(logits: torch.Tensor, scores: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Pool one bag: ``(bag_logits [C], attention weights [n, C])``.""" + if scores.shape != logits.shape: + raise ValueError( + f"Attention scores {tuple(scores.shape)} and instance logits {tuple(logits.shape)} " + "must match: the attention needs one channel per class." + ) + weights = torch.softmax(scores.float(), dim=0) + return (weights * logits.float()).sum(dim=0), weights + + +def pool_by_bag( + logits: torch.Tensor, scores: torch.Tensor, bag_ids: Sequence[str] +) -> tuple[list[str], torch.Tensor]: + """Pool instances grouped by ``bag_ids``: sorted bag ids and ``[n_bags, C]`` logits.""" + groups: dict[str, list[int]] = {} + for i, bag in enumerate(bag_ids): + groups.setdefault(str(bag), []).append(i) + names = sorted(groups) + pooled = [] + for name in names: + idx = torch.tensor(groups[name], device=logits.device) + pooled.append(attention_pool(logits[idx], scores[idx])[0]) + return names, torch.stack(pooled) if pooled else logits.new_zeros((0, logits.shape[1])) diff --git a/src/grass_mil/models/objectives.py b/src/grass_mil/models/objectives.py new file mode 100644 index 0000000..66f20c9 --- /dev/null +++ b/src/grass_mil/models/objectives.py @@ -0,0 +1,162 @@ +"""Bag-level objectives and metrics for the three target types. + +* ``classification``: weighted cross-entropy over ``num_classes`` logits; AUROC + (macro one-vs-rest beyond two classes), balanced accuracy and accuracy. +* ``regression``: mean squared error or Huber loss per target column; RMSE, + MAE, R² and Pearson correlation (averaged over columns). +* ``survival``: a Cox partial-likelihood ranking loss on one risk logit, + computed over all bags of an optimisation step; Harrell's concordance index. + Targets are ``[time, event]``. + +Every bag may carry a sample weight; class weights apply to classification. +""" + +from __future__ import annotations + +import math +from typing import Any, Mapping + +import numpy as np +import torch +import torch.nn.functional as F + +__all__ = ["TARGET_TYPES", "output_dim", "Objective", "region_metrics", "concordance_index", "class_weights_from_labels"] + +TARGET_TYPES = ("classification", "regression", "survival") + + +def output_dim(task: Mapping[str, Any]) -> int: + """Logits per instance for a task config.""" + kind = task["target_type"] + if kind == "classification": + n = int(task.get("num_classes", 2)) + if n < 2: + raise ValueError("task.num_classes must be >= 2 (one logit per class).") + return n + if kind == "regression": + return len(task.get("target_columns") or [None]) + if kind == "survival": + return 1 + raise ValueError(f"task.target_type must be one of {TARGET_TYPES}, got '{kind}'.") + + +class Objective: + """The loss of a stack of bag logits against their targets.""" + + def __init__(self, task: Mapping[str, Any], class_weight: torch.Tensor | None = None) -> None: + self.kind = task["target_type"] + if self.kind not in TARGET_TYPES: + raise ValueError(f"task.target_type must be one of {TARGET_TYPES}") + self.class_weight = class_weight + self.loss_name = str(task.get("loss", "mse")) + self.huber_delta = float(task.get("huber_delta", 1.0)) + self.cox_regularization = float(task.get("cox_regularization", 0.0)) + if self.kind == "regression" and self.loss_name not in {"mse", "huber"}: + raise ValueError("task.loss must be 'mse' or 'huber' for regression.") + if self.kind == "survival" and len(task.get("target_columns") or []) != 2: + raise ValueError("Survival tasks need task.target_columns: [time, event].") + + def __call__(self, logits: torch.Tensor, targets: torch.Tensor, weights: torch.Tensor | None = None) -> torch.Tensor: + logits = logits.float() + if self.kind == "classification": + weight = self.class_weight.to(logits.device) if self.class_weight is not None else None + per_bag = F.cross_entropy(logits, targets[:, 0].long(), weight=weight, reduction="none") + elif self.kind == "regression": + if self.loss_name == "huber": + per_bag = F.huber_loss(logits, targets, delta=self.huber_delta, reduction="none").mean(1) + else: + per_bag = F.mse_loss(logits, targets, reduction="none").mean(1) + else: + return cox_loss(logits[:, 0], targets[:, 0], targets[:, 1], weights, self.cox_regularization) + if weights is not None: + per_bag = per_bag * weights.view(-1).to(per_bag.device) + return per_bag.mean() + + +def cox_loss( + risk: torch.Tensor, + time: torch.Tensor, + event: torch.Tensor, + weights: torch.Tensor | None = None, + regularization: float = 0.0, +) -> torch.Tensor: + """Negative Cox partial log-likelihood (Breslow ties), averaged over events. + + ``regularization`` adds an L2 penalty on the risk scores, which keeps them + centred (the partial likelihood is invariant to a shared shift). + """ + order = torch.argsort(time, descending=True) + risk, event = risk[order], event[order].float() + w = torch.ones_like(risk) if weights is None else weights.view(-1)[order].float().to(risk.device) + # Cumulative log-sum-exp over the risk set (everyone with time >= t_i). + log_risk_set = torch.logcumsumexp(risk + torch.log(w), dim=0) + n_events = (event * w).sum() + if n_events <= 0: + # No observed event in this step: nothing to rank. Keep the graph so the + # optimiser step is a well-defined no-op. + return (risk * 0.0).sum() + loss = -((risk - log_risk_set) * event * w).sum() / n_events + return loss + regularization * (risk**2).mean() + + +def class_weights_from_labels(labels: np.ndarray, num_classes: int) -> torch.Tensor: + """Inverse class frequency ``n / (k * count_c)``, normalised to mean 1 (absent classes get 1).""" + counts = np.bincount(np.asarray(labels, dtype=int), minlength=num_classes).astype(float) + weights = np.ones(num_classes) + present = counts > 0 + weights[present] = counts.sum() / (present.sum() * counts[present]) + return torch.tensor(weights / weights.mean(), dtype=torch.float32) + + +def concordance_index(risk: np.ndarray, time: np.ndarray, event: np.ndarray) -> float: + """Harrell's C: the fraction of comparable pairs ordered by risk (ties count half).""" + risk, time, event = (np.asarray(a, dtype=float).ravel() for a in (risk, time, event)) + comparable = (event[:, None] > 0) & (time[:, None] < time[None, :]) + n = comparable.sum() + if n == 0: + return float("nan") + higher = risk[:, None] > risk[None, :] + tied = risk[:, None] == risk[None, :] + return float(((comparable & higher).sum() + 0.5 * (comparable & tied).sum()) / n) + + +def region_metrics(kind: str, logits: torch.Tensor, targets: torch.Tensor) -> dict[str, float]: + """Metrics of bag (region) logits; ``loss`` is the unweighted objective.""" + from sklearn.metrics import balanced_accuracy_score, roc_auc_score + + logits = logits.detach().float().cpu() + targets = targets.detach().float().cpu() + out: dict[str, float] = {} + if kind == "classification": + y = targets[:, 0].long().numpy() + out["loss"] = float(F.cross_entropy(logits, torch.from_numpy(y))) + prob = torch.softmax(logits, dim=1).numpy() + pred = prob.argmax(1) + out["acc"] = float((pred == y).mean()) + # Undefined on a single-class set: logged as NaN, which checkpoint + # selection treats as "no improvement". + out["bal_acc"] = out["auroc"] = float("nan") + if len(np.unique(y)) > 1: + out["bal_acc"] = float(balanced_accuracy_score(y, pred)) + try: + out["auroc"] = float( + roc_auc_score(y, prob[:, 1]) + if prob.shape[1] == 2 + else roc_auc_score(y, prob, multi_class="ovr", average="macro", labels=list(range(prob.shape[1]))) + ) + except ValueError: + pass + elif kind == "regression": + err = logits - targets + out["loss"] = float((err**2).mean()) + out["rmse"] = math.sqrt(out["loss"]) + out["mae"] = float(err.abs().mean()) + ss_tot = ((targets - targets.mean(0)) ** 2).sum(0) + out["r2"] = out["pearson"] = float("nan") + if bool((ss_tot > 0).all()): + out["r2"] = float((1 - (err**2).sum(0) / ss_tot).mean()) + out["pearson"] = float(np.mean([np.corrcoef(logits[:, j], targets[:, j])[0, 1] for j in range(targets.shape[1])])) + else: + out["loss"] = float(cox_loss(logits[:, 0], targets[:, 0], targets[:, 1])) + out["c_index"] = concordance_index(logits[:, 0].numpy(), targets[:, 0].numpy(), targets[:, 1].numpy()) + return out diff --git a/src/grass_mil/models/optimization.py b/src/grass_mil/models/optimization.py new file mode 100644 index 0000000..3cd31c3 --- /dev/null +++ b/src/grass_mil/models/optimization.py @@ -0,0 +1,54 @@ +"""Optimizer and learning-rate schedule construction.""" + +from __future__ import annotations + +from typing import Any, Mapping + +import hydra +import torch +from omegaconf import OmegaConf + +__all__ = ["build_optimizer", "build_scheduler"] + + +def _plain(cfg: Any) -> dict: + return dict(OmegaConf.to_container(cfg, resolve=True)) if OmegaConf.is_config(cfg) else dict(cfg) + + +def build_optimizer(cfg: Mapping[str, Any], groups: list[dict]) -> torch.optim.Optimizer: + """Instantiate ``cfg`` (``_target_`` + kwargs) over parameter groups. + + Groups may carry their own ``lr`` / ``weight_decay``; the rest inherit the + optimizer's. + """ + kwargs = _plain(cfg) + cls = hydra.utils.get_class(kwargs.pop("_target_")) + return cls(groups, **kwargs) + + +def build_scheduler( + cfg: Mapping[str, Any] | None, + optimizer: torch.optim.Optimizer, + *, + total_steps: int | None, + warmup_steps: int = 0, +): + """A linear warmup over ``warmup_steps`` followed by ``cfg``. + + ``T_max: auto`` in ``cfg`` becomes the number of optimizer steps after the + warmup, so a cosine schedule reaches its minimum exactly at the last step. + """ + if cfg is None: + return None + kwargs = _plain(cfg) + if str(kwargs.get("T_max", "")).lower() == "auto": + if not total_steps: + raise ValueError("scheduler.T_max=auto needs the total number of optimizer steps.") + kwargs["T_max"] = max(1, int(total_steps) - int(warmup_steps)) + main = hydra.utils.instantiate(kwargs, optimizer=optimizer) + if warmup_steps <= 0: + return main + warmup = torch.optim.lr_scheduler.LinearLR( + optimizer, start_factor=1.0 / (warmup_steps + 1), end_factor=1.0, total_iters=warmup_steps + ) + return torch.optim.lr_scheduler.SequentialLR(optimizer, [warmup, main], milestones=[warmup_steps]) diff --git a/src/grass_mil/models/pretraining.py b/src/grass_mil/models/pretraining.py new file mode 100644 index 0000000..d12093b --- /dev/null +++ b/src/grass_mil/models/pretraining.py @@ -0,0 +1,172 @@ +"""Label-free BGRL pretraining of the ego-graph encoder. + +Two augmented views of every ego-graph batch are encoded node by node by an +online encoder (plus a predictor) and by a target encoder that is an +exponential moving average of the online one; the loss is the symmetric +negative cosine similarity between the online predictions and the target node +embeddings. The EMA momentum follows a +cosine schedule from ``momentum.start`` to ``momentum.end`` over the run. +""" + +from __future__ import annotations + +import copy +import math +from dataclasses import dataclass, fields +from typing import Any, Mapping + +import lightning as L +import torch +import torch.nn.functional as F +from lightning.pytorch.utilities import rank_zero_warn +from omegaconf import OmegaConf +from torch import nn + +from grass_mil.data.datamodule import InputSpec + +from .augmentations import PretrainViews +from .checkpoints import load_encoder_weights +from .encoder import EncoderConfig, GINEEncoder +from .optimization import build_optimizer, build_scheduler + +__all__ = ["PretrainConfig", "BGRLModule"] + + +@dataclass(frozen=True) +class PretrainConfig: + """Attributes: + total_steps: Optimizer steps of the run (also ``trainer.max_steps``). + warmup_steps: Linear learning-rate warmup steps. + momentum_start / momentum_end: EMA momentum of the target encoder. + predictor_hidden_dim: Hidden width of the predictor MLP. + augmentation: The two views. + seed: Seed of the augmentation draws. + """ + + total_steps: int = 200_000 + warmup_steps: int = 1000 + momentum_start: float = 0.99 + momentum_end: float = 1.0 + predictor_hidden_dim: int = 512 + augmentation: PretrainViews = PretrainViews() + seed: int = 0 + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any] | None) -> PretrainConfig: + cfg = dict(cfg or {}) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown pretrain keys: {sorted(unknown)}") + cfg["augmentation"] = PretrainViews.from_dict(cfg.get("augmentation")) + return cls(**cfg) + + def momentum(self, step: int) -> float: + progress = min(max(step, 0), self.total_steps) / max(1, self.total_steps) + return self.momentum_end - (self.momentum_end - self.momentum_start) * 0.5 * (1 + math.cos(math.pi * progress)) + + +class Predictor(nn.Module): + def __init__(self, dim: int, hidden: int) -> None: + super().__init__() + self.net = nn.Sequential(nn.Linear(dim, hidden), nn.PReLU(1), nn.Linear(hidden, dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +def _plain(value: Any) -> Any: + return OmegaConf.to_container(value, resolve=True) if OmegaConf.is_config(value) else value + + +class BGRLModule(L.LightningModule): + def __init__( + self, + encoder: Mapping[str, Any], + pretrain: Mapping[str, Any], + optim: Mapping[str, Any] | None = None, + scheduler: Mapping[str, Any] | None = None, + init_from_ckpt: str | None = None, + ) -> None: + super().__init__() + hparams = {k: _plain(v) for k, v in locals().items() if k not in ("self", "__class__")} + self.save_hyperparameters(hparams, logger=False) + self.encoder_cfg = EncoderConfig.from_dict(hparams["encoder"]) + self.pretrain = PretrainConfig.from_dict(hparams["pretrain"]) + self.input_spec: InputSpec | None = None + self.encoder: GINEEncoder | None = None + self.target_encoder: GINEEncoder | None = None + self.predictor: Predictor | None = None + self._view = None + self._generator = torch.Generator().manual_seed(int(self.pretrain.seed)) + self._last_ema_step = 0 + + def build(self, spec: InputSpec, feature_names: list[str] | None = None) -> None: + self.input_spec = spec + self.encoder = GINEEncoder( + self.encoder_cfg, feature_dim=spec.feature_dim, cell_types=spec.cell_types, edge_dim=spec.edge_dim + ) + self.predictor = Predictor(self.encoder.output_dim, self.pretrain.predictor_hidden_dim) + self.target_encoder = copy.deepcopy(self.encoder) + self.target_encoder.requires_grad_(False) + self._view = self.pretrain.augmentation.build( + feature_names=feature_names or [], unassigned_index=self.encoder.unassigned_index + ) + + def setup(self, stage: str | None = None) -> None: + if self.encoder is not None: + return + datamodule = self.trainer.datamodule + self.build(datamodule.input_spec, feature_names=datamodule.index.feature_names) + if self.hparams.init_from_ckpt: + load_encoder_weights(self.encoder, self.hparams.init_from_ckpt) + self.target_encoder.load_state_dict(self.encoder.state_dict()) + if self.trainer.max_steps not in (-1, None) and self.trainer.max_steps != self.pretrain.total_steps: + rank_zero_warn( + f"trainer.max_steps={self.trainer.max_steps} differs from pretrain.total_steps=" + f"{self.pretrain.total_steps}; the schedules follow pretrain.total_steps." + ) + + def on_save_checkpoint(self, checkpoint: dict[str, Any]) -> None: + checkpoint["input_spec"] = self.input_spec.as_dict() + + def _loss(self, batch) -> torch.Tensor: + view1, view2 = self._view(batch, self._generator), self._view(batch, self._generator) + q1 = self.predictor(self.encoder.node_embeddings(view1)) + q2 = self.predictor(self.encoder.node_embeddings(view2)) + with torch.no_grad(): + y1 = self.target_encoder.node_embeddings(view1) + y2 = self.target_encoder.node_embeddings(view2) + return 2.0 - F.cosine_similarity(q1, y2, dim=-1).mean() - F.cosine_similarity(q2, y1, dim=-1).mean() + + def training_step(self, batch, batch_idx: int) -> torch.Tensor: + loss = self._loss(batch) + self.log("train/loss", loss, on_step=True, on_epoch=True, prog_bar=True, batch_size=int(batch.num_graphs)) + return loss + + def validation_step(self, batch, batch_idx: int) -> None: + self.log("val/loss", self._loss(batch), prog_bar=True, batch_size=int(batch.num_graphs)) + + def on_train_batch_end(self, outputs, batch, batch_idx: int) -> None: + """EMA update of the target encoder after every optimizer step.""" + step = int(self.global_step) + if step <= self._last_ema_step: + return + self._last_ema_step = step + momentum = self.pretrain.momentum(step - 1) + with torch.no_grad(): + for online, target in zip(self.encoder.parameters(), self.target_encoder.parameters()): + target.mul_(momentum).add_(online.detach(), alpha=1.0 - momentum) + self.log("train/momentum", momentum, on_step=True, on_epoch=False) + + def configure_optimizers(self): + params = list(self.encoder.parameters()) + list(self.predictor.parameters()) + optimizer = build_optimizer(self.hparams.optim, [{"params": params}]) + scheduler = build_scheduler( + self.hparams.scheduler, + optimizer, + total_steps=self.pretrain.total_steps, + warmup_steps=self.pretrain.warmup_steps, + ) + if scheduler is None: + return optimizer + return {"optimizer": optimizer, "lr_scheduler": {"scheduler": scheduler, "interval": "step"}} diff --git a/src/grass_mil/models/probe.py b/src/grass_mil/models/probe.py new file mode 100644 index 0000000..6ccd87a --- /dev/null +++ b/src/grass_mil/models/probe.py @@ -0,0 +1,91 @@ +"""Initialising a linear instance head from a probe on frozen ego-graph embeddings. + +Every ego-graph inherits the label of its unit. A standardised logistic +regression (classification) or ridge regression (regression) is fitted on the +embeddings of the training units and folded, scaler included, into the weights +of a ``Linear(embedding_dim, num_logits)`` head. For two classes the class-0 +row is zero and the class-1 row carries the decision function, so the softmax +of the two logits equals the logistic probability. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import Any, Mapping + +import numpy as np +import torch + +__all__ = ["ProbeConfig", "collect_embeddings", "fit_linear_head"] + + +@dataclass(frozen=True) +class ProbeConfig: + roots_per_unit: int | None = 96 + seed: int = 0 + C: float = 1.0 + class_weight: str | None = "balanced" + max_iter: int = 3000 + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any] | None) -> ProbeConfig: + cfg = dict(cfg or {}) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown finetune.probe keys: {sorted(unknown)}") + return cls(**cfg) + + +@torch.no_grad() +def collect_embeddings(encoder: torch.nn.Module, loader, device: torch.device) -> tuple[np.ndarray, np.ndarray]: + """Embeddings and inherited targets of every ego-graph of ``loader``.""" + was_training = encoder.training + encoder.eval() + embeddings, targets = [], [] + for batch in loader: + batch = batch.to(device) + embeddings.append(encoder(batch).float().cpu().numpy()) + targets.append(batch.graph_y.float().cpu().numpy()) + encoder.train(was_training) + if not embeddings: + raise ValueError("The probe found no ego-graphs in the training split.") + return np.concatenate(embeddings), np.concatenate(targets) + + +def fit_linear_head( + embeddings: np.ndarray, + targets: np.ndarray, + *, + kind: str, + num_logits: int, + cfg: ProbeConfig, +) -> tuple[torch.Tensor, torch.Tensor]: + """``(weight [num_logits, dim], bias [num_logits])`` of the folded probe.""" + from sklearn.linear_model import LogisticRegression, Ridge + from sklearn.preprocessing import StandardScaler + + scaler = StandardScaler().fit(embeddings) + z = scaler.transform(embeddings) + if kind == "classification": + y = np.floor(targets[:, 0] + 0.5).astype(int) + if len(np.unique(y)) < 2: + raise ValueError("The probe's training split holds a single class.") + model = LogisticRegression(C=cfg.C, class_weight=cfg.class_weight, max_iter=cfg.max_iter) + model.fit(z, y) + coef, intercept = model.coef_, model.intercept_ + elif kind == "regression": + model = Ridge(alpha=1.0 / cfg.C).fit(z, targets) + coef, intercept = np.atleast_2d(model.coef_), np.atleast_1d(model.intercept_) + else: + raise ValueError(f"No probe for target type '{kind}'.") + w = coef / scaler.scale_[None, :] + b = intercept - (coef * scaler.mean_[None, :] / scaler.scale_[None, :]).sum(axis=1) + if coef.shape[0] not in (1, num_logits): + raise ValueError(f"The probe saw {coef.shape[0]} classes but the head has {num_logits} logits.") + weight = torch.zeros(num_logits, embeddings.shape[1]) + bias = torch.zeros(num_logits) + if kind == "classification" and coef.shape[0] == 1: + weight[1], bias[1] = torch.from_numpy(w[0]).float(), float(b[0]) + else: + weight[:], bias[:] = torch.from_numpy(w).float(), torch.from_numpy(b).float() + return weight, bias diff --git a/src/grass_mil/models/supervised.py b/src/grass_mil/models/supervised.py new file mode 100644 index 0000000..2cde1b9 --- /dev/null +++ b/src/grass_mil/models/supervised.py @@ -0,0 +1,334 @@ +"""GRASS-MIL supervised module: ego-graph instances pooled by gated attention. + +Each training batch is one bag (all sampled ego-graphs of one unit). Bag +logits of ``finetune.hyperbatch_size`` consecutive bags are stacked and the +loss is computed over the stack before one optimizer step, so a step sees +several units (which the Cox objective needs) while a bag's attention softmax +always covers the whole unit. Validation and test units may be scored in +several chunks; their instances are pooled per unit at the end of the epoch. +""" + +from __future__ import annotations + +import logging +import math +from dataclasses import dataclass, fields +from typing import Any, Mapping + +import lightning as L +import torch +import torch.distributed as dist +from omegaconf import OmegaConf + +from grass_mil.data.datamodule import InputSpec + +from .augmentations import FinetuneAugmentation +from .checkpoints import load_encoder_weights +from .encoder import EncoderConfig, GINEEncoder +from .mil import AttentionConfig, GatedAttention, HeadConfig, InstanceHead, attention_pool, pool_by_bag +from .objectives import Objective, class_weights_from_labels, output_dim, region_metrics +from .optimization import build_optimizer, build_scheduler +from .probe import ProbeConfig, collect_embeddings, fit_linear_head + +__all__ = ["FinetuneConfig", "SupervisedModule"] + +log = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class FinetuneConfig: + """The fine-tuning procedure (set by the ``regime`` presets). + + Attributes: + lr_scales: Learning-rate multipliers of ``encoder``, ``head`` and + ``attention`` on ``optim.lr``; ``0`` freezes the component (a + frozen encoder also stays in eval mode). + head_weight_decay: Weight decay of the head (``None``: the optimizer's). + head_init: ``probe`` (fit a linear probe on the frozen embeddings of + the training units first) or ``random``. + freeze_head: Keep the initial head fixed. + hyperbatch_size: Bags per optimizer step. + grad_clip_norm: Gradient-norm clip per step (``None``: off). + warmup_steps: Linear learning-rate warmup steps. + augmentation: Edge drop / node masking during training. + probe: Probe settings for ``head_init: probe``. + """ + + lr_scales: Mapping[str, float] + head_weight_decay: float | None = None + head_init: str = "probe" + freeze_head: bool = False + hyperbatch_size: int = 12 + grad_clip_norm: float | None = 1.0 + warmup_steps: int = 0 + augmentation: FinetuneAugmentation = FinetuneAugmentation() + probe: ProbeConfig = ProbeConfig() + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any]) -> FinetuneConfig: + cfg = dict(cfg) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown finetune keys: {sorted(unknown)}") + scales = dict(cfg.get("lr_scales") or {}) + if set(scales) != {"encoder", "head", "attention"} or any(v is None for v in scales.values()): + raise ValueError("finetune.lr_scales needs numeric encoder, head and attention scales.") + cfg["lr_scales"] = {k: float(v) for k, v in scales.items()} + cfg["augmentation"] = FinetuneAugmentation.from_dict(cfg.get("augmentation")) + cfg["probe"] = ProbeConfig.from_dict(cfg.get("probe")) + out = cls(**cfg) + if out.head_init not in {"probe", "random"}: + raise ValueError("finetune.head_init must be 'probe' or 'random'.") + if out.hyperbatch_size < 1: + raise ValueError("finetune.hyperbatch_size must be >= 1.") + return out + + +def _plain(value: Any) -> Any: + return OmegaConf.to_container(value, resolve=True) if OmegaConf.is_config(value) else value + + +class SupervisedModule(L.LightningModule): + def __init__( + self, + encoder: Mapping[str, Any], + head: Mapping[str, Any], + attention: Mapping[str, Any], + task: Mapping[str, Any], + finetune: Mapping[str, Any], + optim: Mapping[str, Any] | None = None, + scheduler: Mapping[str, Any] | None = None, + init_from_ckpt: str | None = None, + ) -> None: + super().__init__() + hparams = {k: _plain(v) for k, v in locals().items() if k not in ("self", "__class__")} + self.save_hyperparameters(hparams, logger=False) + self.automatic_optimization = False + self.encoder_cfg = EncoderConfig.from_dict(hparams["encoder"]) + self.head_cfg = HeadConfig.from_dict(hparams["head"]) + self.attention_cfg = AttentionConfig.from_dict(hparams["attention"]) + self.task = dict(hparams["task"]) + self.finetune = FinetuneConfig.from_dict(hparams["finetune"]) + self.num_logits = output_dim(self.task) + self.input_spec: InputSpec | None = None + self.encoder: GINEEncoder | None = None + self.head: InstanceHead | None = None + self.attention: GatedAttention | None = None + self.objective = Objective(self.task) + self._augment = None + self._bag_buffer: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]] = [] + self._eval_buffer: dict[str, list] = {} + self.test_predictions: dict[str, Any] | None = None + + # ------------------------------------------------------------- building + @property + def built(self) -> bool: + return self.encoder is not None + + def build(self, spec: InputSpec) -> None: + """Create the layers for the given node/edge inputs.""" + self.input_spec = spec + self.encoder = GINEEncoder( + self.encoder_cfg, feature_dim=spec.feature_dim, cell_types=spec.cell_types, edge_dim=spec.edge_dim + ) + width = self.encoder.output_dim + self.head = InstanceHead(width, self.num_logits, self.head_cfg) + self.attention = GatedAttention(width, self.num_logits, self.attention_cfg) + if self.finetune.augmentation.enabled: + self._augment = self.finetune.augmentation.build(unassigned_index=self.encoder.unassigned_index) + + def setup(self, stage: str | None = None) -> None: + if self.built: + return + datamodule = self.trainer.datamodule + self.build(datamodule.input_spec) + if self.hparams.init_from_ckpt: + load_encoder_weights(self.encoder, self.hparams.init_from_ckpt) + if stage == "fit": + self._configure_class_weights(datamodule.dataset_train) + if self.finetune.head_init == "probe": + self._fit_probe_head(datamodule) + self._apply_freezing() + + def _configure_class_weights(self, dataset) -> None: + setting = self.task.get("class_weights") + if self.task["target_type"] != "classification" or setting is None: + return + if setting == "auto": + labels = [dataset.labels.classes(dataset.label_id(e)) for e in dataset.entries] + self.objective.class_weight = class_weights_from_labels(labels, self.num_logits) + else: + self.objective.class_weight = torch.tensor([float(w) for w in setting]) + + def _fit_probe_head(self, datamodule) -> None: + if self.task["target_type"] == "survival": + log.info("No probe for survival targets; the head starts from a random initialisation.") + return + if not self.head.is_linear: + raise ValueError("finetune.head_init=probe needs a linear head (model.head.num_layers=1).") + if not self.hparams.init_from_ckpt: + log.warning("Fitting the probe head on a randomly initialised encoder (no model.init_from_ckpt).") + cfg = self.finetune.probe + loader = datamodule.loader(datamodule.dataset_train, roots_per_unit=cfg.roots_per_unit, seed=cfg.seed) + device = self.trainer.strategy.root_device + self.encoder.to(device) + embeddings, targets = collect_embeddings(self.encoder, loader, device) + self.encoder.cpu() + weight, bias = fit_linear_head(embeddings, targets, kind=self.task["target_type"], num_logits=self.num_logits, cfg=cfg) + with torch.no_grad(): + self.head.net[0].weight.copy_(weight) + self.head.net[0].bias.copy_(bias) + log.info(f"Probe head fitted on {len(embeddings)} ego-graphs of {len(datamodule.dataset_train)} units.") + + def _apply_freezing(self) -> None: + scales = self.finetune.lr_scales + for module, frozen in ( + (self.encoder, scales["encoder"] == 0), + (self.head, scales["head"] == 0 or self.finetune.freeze_head), + (self.attention, scales["attention"] == 0), + ): + if frozen: + module.requires_grad_(False) + if scales["encoder"] == 0: + self.encoder.eval() + + def train(self, mode: bool = True): + super().train(mode) + if mode and self.built and self.finetune.lr_scales["encoder"] == 0: + self.encoder.eval() + return self + + def on_save_checkpoint(self, checkpoint: dict[str, Any]) -> None: + checkpoint["input_spec"] = self.input_spec.as_dict() + + @classmethod + def from_checkpoint(cls, path: str, map_location: str | torch.device = "cpu") -> SupervisedModule: + """Rebuild a trained module from its checkpoint alone (eval mode).""" + ckpt = torch.load(path, map_location=map_location, weights_only=False) + hparams = dict(ckpt["hyper_parameters"]) + hparams["init_from_ckpt"] = None + hparams["finetune"] = {**hparams["finetune"], "head_init": "random"} + module = cls(**hparams) + module.build(InputSpec.from_dict(ckpt["input_spec"])) + module.load_state_dict(ckpt["state_dict"]) + return module.eval() + + # -------------------------------------------------------------- forward + def forward_instances(self, batch) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Embeddings, instance logits and attention scores of every ego-graph.""" + if self.training and self._augment is not None: + batch = self._augment(batch) + embeddings = self.encoder(batch) + return embeddings, self.head(embeddings), self.attention(embeddings) + + def forward(self, batch) -> torch.Tensor: + """Bag logits of a batch holding the ego-graphs of one unit.""" + _, logits, scores = self.forward_instances(batch) + return attention_pool(logits, scores)[0] + + # ------------------------------------------------------------- training + def training_step(self, batch, batch_idx: int) -> None: + bag_logits = self(batch).unsqueeze(0) + weights = getattr(batch, "graph_w", None) + self._bag_buffer.append((bag_logits, batch.graph_y[:1].float(), None if weights is None else weights[:1])) + if len(self._bag_buffer) >= self.finetune.hyperbatch_size: + self._optimizer_step(on_step=True) + + def on_train_epoch_end(self) -> None: + if self._bag_buffer: + self._optimizer_step(on_step=False) + + def _optimizer_step(self, *, on_step: bool) -> None: + logits = torch.cat([b[0] for b in self._bag_buffer]) + targets = torch.cat([b[1] for b in self._bag_buffer]) + weights = None + if all(b[2] is not None for b in self._bag_buffer): + weights = torch.cat([b[2] for b in self._bag_buffer]) + self._bag_buffer.clear() + loss = self.objective(logits, targets, weights) + optimizer = self.optimizers() + optimizer.zero_grad(set_to_none=True) + self.manual_backward(loss) + scaler = getattr(self.trainer.precision_plugin, "scaler", None) + scale = scaler.get_scale() if scaler is not None else None + optimizer.step() + # Under fp16 a step with non-finite gradients is skipped (and the loss + # scale lowered); the schedule only advances on steps that happened. + if scaler is None or scaler.get_scale() >= scale: + scheduler = self.lr_schedulers() + if scheduler is not None: + scheduler.step() + self.log("train/loss", loss.detach(), on_step=on_step, on_epoch=True, prog_bar=True, batch_size=len(targets)) + + def on_before_optimizer_step(self, optimizer) -> None: + # Called after mixed-precision unscaling, so the norm is the true one. + if self.finetune.grad_clip_norm: + params = [p for p in self.parameters() if p.grad is not None] + torch.nn.utils.clip_grad_norm_(params, float(self.finetune.grad_clip_norm)) + + def configure_optimizers(self): + scales = self.finetune.lr_scales + lr = float(self.hparams.optim["lr"]) + groups = [] + for name, module in (("encoder", self.encoder), ("head", self.head), ("attention", self.attention)): + params = [p for p in module.parameters() if p.requires_grad] + if params: + group = {"params": params, "lr": lr * scales[name], "name": name} + if name == "head" and self.finetune.head_weight_decay is not None: + group["weight_decay"] = float(self.finetune.head_weight_decay) + groups.append(group) + if not groups: + raise ValueError("Nothing to train: every component is frozen.") + optimizer = build_optimizer(self.hparams.optim, groups) + n_units = len(self.trainer.datamodule.dataset_train) + steps = int(self.trainer.max_epochs) * math.ceil(n_units / self.finetune.hyperbatch_size) + scheduler = build_scheduler( + self.hparams.scheduler, optimizer, total_steps=steps, warmup_steps=self.finetune.warmup_steps + ) + if scheduler is None: + return optimizer + return {"optimizer": optimizer, "lr_scheduler": {"scheduler": scheduler, "interval": "step"}} + + # ----------------------------------------------------------- evaluation + def _eval_step(self, batch) -> None: + _, logits, scores = self.forward_instances(batch) + buf = self._eval_buffer + buf.setdefault("unit_ids", []).extend(batch.unit_id) + buf.setdefault("logits", []).append(logits.detach().float().cpu()) + buf.setdefault("scores", []).append(scores.detach().float().cpu()) + buf.setdefault("targets", {}).update( + {u: y for u, y in zip(batch.unit_id, batch.graph_y.detach().float().cpu())} + ) + + def validation_step(self, batch, batch_idx: int) -> None: + self._eval_step(batch) + + def test_step(self, batch, batch_idx: int) -> None: + self._eval_step(batch) + + def on_validation_epoch_end(self) -> None: + self._finish_eval("val") + + def on_test_epoch_end(self) -> None: + self._finish_eval("test") + + def _finish_eval(self, stage: str) -> None: + buf, self._eval_buffer = self._eval_buffer, {} + if not buf: + return + if dist.is_available() and dist.is_initialized(): + gathered: list[dict] = [None] * dist.get_world_size() # type: ignore[list-item] + dist.all_gather_object(gathered, buf) + buf = { + "unit_ids": [u for g in gathered for u in g["unit_ids"]], + "logits": [t for g in gathered for t in g["logits"]], + "scores": [t for g in gathered for t in g["scores"]], + "targets": {k: v for g in gathered for k, v in g["targets"].items()}, + } + units, bag_logits = pool_by_bag(torch.cat(buf["logits"]), torch.cat(buf["scores"]), buf["unit_ids"]) + targets = torch.stack([buf["targets"][u] for u in units]) + metrics = region_metrics(self.task["target_type"], bag_logits, targets) + for name, value in metrics.items(): + self.log(f"{stage}/{name}", value, prog_bar=name in ("loss", "auroc", "c_index", "rmse")) + if stage == "test": + self.test_predictions = {"unit_ids": units, "logits": bag_logits, "targets": targets} diff --git a/src/grass_mil/models/supervised_module.py b/src/grass_mil/models/supervised_module.py deleted file mode 100644 index d3b9eca..0000000 --- a/src/grass_mil/models/supervised_module.py +++ /dev/null @@ -1,1448 +0,0 @@ -from __future__ import annotations - -from typing import Any, Dict, List, Optional - -import lightning as L -import torch -import torch.distributed as dist -from omegaconf import DictConfig, OmegaConf - -from grass_mil.batching import extract_subgraphs -from grass_mil.inference.aggregation import aggregate_group_logits -from grass_mil.inference.metrics import compute_task_metrics_from_tensors -from grass_mil.inference.schemas import BatchPredictionPayload - -from .training import ( - aggregate_bag_logits_attention, - aggregate_bag_logits_mean, - build_finetune_augmentation, - build_param_groups, - build_supervised_components, - collect_unit_labels, - compute_binary_accuracy, - compute_categorical_accuracy, - compute_supervised_loss, - extract_bag_ids, - gather_bag_targets, - group_instance_indices_by_bag, - instantiate_optimizer, - instantiate_scheduler_with_warmup, - inverse_frequency_class_weights, - load_state_dict_with_optional_mapping, - resolve_t_max_auto, - validate_task_config, -) - - -class SupervisedModule(L.LightningModule): - """Unified supervised module supporting both mean-pooling and MIL-attention - aggregation. - - Behaviour is controlled by ``task.aggregation``: - - * ``"mean"`` -- simple mean-pooling over patch logits (no attention, no - auxiliary losses, no region accumulation). - * ``"mil_attention"`` -- attention-weighted MIL aggregation with optional - region accumulation, node auxiliary loss, and entropy regularisation. - """ - - def __init__( - self, - encoder: Dict[str, Any], - graph_head: Dict[str, Any], - loss: Dict[str, Any], - optim: Optional[Dict[str, Any]] = None, - attention: Optional[Dict[str, Any]] = None, - scheduler: Optional[Dict[str, Any]] = None, - task: Optional[Dict[str, Any]] = None, - flags: Optional[Dict[str, Any]] = None, - ssl: Optional[Dict[str, Any]] = None, - init_from_ckpt: Optional[str] = None, - init_strict: bool = False, - encoder_init_map: str = "auto_bgrl_or_identity", - freeze_encoder: bool = False, - ) -> None: - super().__init__() - self.save_hyperparameters(logger=False) - - self._built = False - self.encoder = None - self.graph_head = None - self.loss_fn = None - self.attention = None - self.ssl_model = None - - self.optim_cfg = optim - self.scheduler_cfg = scheduler - self.task_cfg = task or {} - validate_task_config(self.task_cfg) - self.flags_cfg = flags or {} - self.ssl_cfg = ssl or {} - self.init_from_ckpt = init_from_ckpt - self.init_strict = init_strict - self.encoder_init_map = encoder_init_map - self.freeze_encoder = freeze_encoder - - # Derive the aggregation mode from the task config. - aggregation = self.task_cfg.get("aggregation", "mean") - self.use_attention = aggregation == "mil_attention" - - # MIL-specific config -- only meaningful when use_attention is True. - if self.use_attention: - self.region_accum_cfg = dict(self.task_cfg.get("region_accumulation", {})) - self.region_accum_enabled = bool(self.region_accum_cfg.get("enabled", False)) - self.region_accum_hyperbatch_size = int( - self.region_accum_cfg.get("hyperbatch_size", 8) - ) - self.region_accum_flush_on_epoch_end = bool( - self.region_accum_cfg.get("flush_on_epoch_end", True) - ) - if self.region_accum_enabled and self.region_accum_hyperbatch_size < 1: - raise ValueError( - "task.region_accumulation.hyperbatch_size must be >= 1 " - "when region accumulation is enabled." - ) - else: - self.region_accum_cfg = {} - self.region_accum_enabled = False - self.region_accum_hyperbatch_size = 0 - self.region_accum_flush_on_epoch_end = False - - # Regime knobs (all optional; see configs/regime/*.yaml for the presets). - self.optimization_cfg = dict(self.task_cfg.get("optimization", {}) or {}) - self.class_weights_cfg = self.task_cfg.get("class_weights") - self.finetune_aug_cfg = self.task_cfg.get("finetune_augmentation") - self.head_init_path = self.task_cfg.get("head_init_path") - self.freeze_head = bool(self.task_cfg.get("freeze_head", False)) - grad_clip = self.task_cfg.get("grad_clip_norm") - self.grad_clip_norm = None if grad_clip in (None, "", "none") else float(grad_clip) - self.scheduler_interval = str(self.task_cfg.get("scheduler_interval", "epoch")) - if self.scheduler_interval not in {"epoch", "step"}: - raise ValueError("task.scheduler_interval must be 'epoch' or 'step'.") - self._finetune_augment = None - self._encoder_frozen = False - self._eval_stage = "val" - self.test_predictions: Optional[Dict[str, Any]] = None - self._skipped_optimizer_steps = 0 - - self._region_loss_buffer: list[torch.Tensor] = [] - self._manual_optimizer_steps = 0 - # In inference, callers can disable instance-level payload emission to reduce - # peak memory when aggregation only needs bag-level outputs. - self._predict_emit_instance_payload = True - self._predict_emit_instance_embeddings = False - self._predict_emit_instance_composition = False - self._predict_emit_instance_centroids = False - self._predict_composition_label = "cell_type" - self._predict_emit_embeddings_payload = False - self._predict_emit_node_embeddings = False - self._reset_val_epoch_buffers() - - # ------------------------------------------------------------------ - # Lightning hooks - # ------------------------------------------------------------------ - - @property - def automatic_optimization(self) -> bool: # noqa: D401 - return not self.region_accum_enabled - - def setup(self, stage: Optional[str] = None) -> None: - if self._built: - return - - loss_cfg = dict(self.hparams.loss) - loss_cfg.setdefault("loss_type", self.task_cfg.get("loss", "categorical_bce")) - - # Build attention config only when MIL is active. - if self.use_attention: - attention_cfg = dict(self.hparams.attention) - attention_cfg.pop("_target_", None) - else: - attention_cfg = None - - components = build_supervised_components( - module=self, - encoder_cfg=dict(self.hparams.encoder), - graph_head_cfg=dict(self.hparams.graph_head), - attention_cfg=attention_cfg, - use_attention=self.use_attention, - loss_cfg=loss_cfg, - use_ssl=bool(self.flags_cfg.get("use_ssl", False)), - ssl_cfg=dict(self.ssl_cfg), - ) - - self.encoder = components["encoder"] - self.graph_head = components["graph_head"] - self.loss_fn = components["loss_fn"] - self.attention = components["attention"] - self.ssl_model = components["ssl_model"] - - load_state_dict_with_optional_mapping( - self, - init_from_ckpt=self.init_from_ckpt, - init_strict=self.init_strict, - encoder_init_map=self.encoder_init_map, - ) - encoder_scale = float(self.optimization_cfg.get("encoder_lr_scale", 1.0) or 0.0) - if self.freeze_encoder or encoder_scale <= 0.0: - # A zero learning-rate scale is the preprint's "frozen encoder": no - # gradients, no optimizer group, and the encoder stays in eval mode - # (dropout/batchnorm off) throughout training. - for p in self.encoder.parameters(): - p.requires_grad = False - self._encoder_frozen = True - self.encoder.eval() - - if self.head_init_path: - if str(self.head_init_path).strip().lower() == "auto": - raise ValueError( - "task.head_init_path='auto' is only resolved by grass-mil-nested-cv, which " - "fits a probe head per split. For a single run, fit one with grass-mil-probe " - "and pass task.head_init_path=/path/to/probe_head.pt." - ) - self._load_head_init(str(self.head_init_path)) - if self.freeze_head: - for p in self.graph_head.parameters(): - p.requires_grad = False - - self._configure_class_weights(stage) - self._configure_finetune_augmentation() - self._built = True - - def _load_head_init(self, path: str) -> None: - """Initialise the graph head from a probe-head file (``graph_head.net.*`` keys).""" - state = torch.load(path, map_location="cpu", weights_only=False) - if isinstance(state, dict) and "state_dict" in state: - state = state["state_dict"] - head_state = { - k.split("graph_head.", 1)[1]: v - for k, v in state.items() - if isinstance(k, str) and k.startswith("graph_head.") and isinstance(v, torch.Tensor) - } - if not head_state: - raise ValueError(f"No 'graph_head.*' tensors found in head_init_path={path!r}.") - missing, unexpected = self.graph_head.load_state_dict(head_state, strict=False) - if unexpected: - raise ValueError(f"head_init_path carries keys the head does not have: {unexpected}") - if missing: - raise ValueError( - f"head_init_path does not cover the whole head (missing {missing}); a probe " - "pre-fit targets a linear head (model.graph_head.num_layers=1)." - ) - - def _configure_class_weights(self, stage: Optional[str]) -> None: - cfg = self.class_weights_cfg - if cfg is None or stage not in (None, "fit"): - return - if not hasattr(self.loss_fn, "class_weight"): - return - if isinstance(cfg, str): - if cfg.lower() != "auto_train": - raise ValueError("task.class_weights must be 'auto_train', a list, or null.") - labels = self._training_unit_labels() - if labels is None: - return - num_classes = int(self.hparams.graph_head.get("output_dim", 1)) - if num_classes <= 1: - return - weights = inverse_frequency_class_weights(labels, num_classes) - else: - weights = torch.as_tensor([float(v) for v in cfg], dtype=torch.float32) - self.loss_fn.class_weight = weights - self._class_weights = weights - - def _training_unit_labels(self) -> Optional[list[int]]: - try: - trainer = self.trainer - except RuntimeError: - return None - if trainer is None or trainer.datamodule is None: - return None - dataset = getattr(trainer.datamodule, "dataset_train", None) - if dataset is None: - return None - labels = collect_unit_labels(dataset, target_columns=self.task_cfg.get("target_columns")) - return labels or None - - def _configure_finetune_augmentation(self) -> None: - if not self.finetune_aug_cfg: - return - input_proj = getattr(self.encoder, "input_proj", None) - self._finetune_augment = build_finetune_augmentation( - self.finetune_aug_cfg, - unassigned_index=getattr(input_proj, "unassigned_index", None), - categorical_column=getattr(input_proj, "_column_index", None), - ) - - def train(self, mode: bool = True): - super().train(mode) - if mode and self._encoder_frozen and self.encoder is not None: - self.encoder.eval() - return self - - # ------------------------------------------------------------------ - # Forward helpers - # ------------------------------------------------------------------ - - def _forward_instances(self, batch): - if self.training and self._finetune_augment is not None: - batch = self._finetune_augment(batch) - node_emb, graph_emb = self.encoder( - batch.x, - batch.edge_index, - edge_attr=getattr(batch, "edge_attr", None), - batch=getattr(batch, "batch", None), - categorical_codes=getattr(batch, "categorical_codes", None), - return_graph_embedding=True, - ) - patch_logits = self.graph_head(graph_emb) - return patch_logits, graph_emb, node_emb - - def _to_optional_str_list( - self, values: Any, *, expected_length: int, field_name: str - ) -> List[Optional[str]]: - if values is None: - return [None] * expected_length - if isinstance(values, torch.Tensor): - values = values.detach().cpu().tolist() - if isinstance(values, (str, bytes)): - raw_values = [values] - else: - try: - raw_values = list(values) - except TypeError: - raw_values = [values] - if len(raw_values) == expected_length: - pass - elif len(raw_values) == 1 and expected_length > 1: - raw_values = raw_values * expected_length - else: - raise ValueError( - f"Batch field '{field_name}' has {len(raw_values)} values, expected " - f"{expected_length} for instance-level metadata alignment." - ) - normalized: List[Optional[str]] = [] - for value in raw_values: - if value is None: - normalized.append(None) - continue - text = str(value).strip() - if text.lower() in {"", "nan", "none", "null"}: - normalized.append(None) - else: - normalized.append(str(value)) - return normalized - - @staticmethod - def _resolve_categorical_column_index(batch: Any, label_name: str) -> Optional[int]: - slices = getattr(batch, "categorical_slices", None) - if isinstance(slices, dict): - value = slices.get(label_name) - return int(value) if value is not None else None - if isinstance(slices, (list, tuple)): - for item in slices: - if isinstance(item, dict) and label_name in item: - value = item[label_name] - return int(value) if value is not None else None - labels = getattr(batch, "categorical_labels", None) - if isinstance(labels, (list, tuple)): - for idx, value in enumerate(labels): - if str(value) == str(label_name): - return idx - return None - - @staticmethod - def _compute_instance_composition( - batch: Any, - *, - n_instances: int, - label_name: str, - device: torch.device, - ) -> Optional[torch.Tensor]: - categorical_codes = getattr(batch, "categorical_codes", None) - if not isinstance(categorical_codes, torch.Tensor) or categorical_codes.numel() == 0: - return None - categorical_codes = categorical_codes.to(device=device).long() - if categorical_codes.ndim != 2: - return None - col_idx = SupervisedModule._resolve_categorical_column_index(batch, label_name) - if col_idx is None or col_idx < 0 or col_idx >= int(categorical_codes.shape[1]): - return None - - node_labels = categorical_codes[:, int(col_idx)] - if node_labels.numel() == 0: - return None - batch_index = getattr(batch, "batch", None) - if isinstance(batch_index, torch.Tensor): - batch_index = batch_index.to(device=device).long() - if batch_index.numel() != node_labels.numel(): - return None - elif n_instances == 1: - batch_index = torch.zeros(node_labels.numel(), dtype=torch.long, device=device) - else: - return None - - n_classes = int(node_labels.max().item()) + 1 - if n_classes <= 0: - return None - composition = torch.zeros((n_instances, n_classes), dtype=torch.float32, device=device) - for inst_idx in range(n_instances): - mask = batch_index == inst_idx - labels = node_labels[mask] - if labels.numel() == 0: - continue - counts = torch.bincount(labels, minlength=n_classes).float() - total = counts.sum() - if total > 0: - counts = counts / total - composition[inst_idx, :] = counts - return composition - - @staticmethod - def _compute_instance_centroids( - batch: Any, - *, - n_instances: int, - device: torch.device, - ) -> Optional[torch.Tensor]: - positions = getattr(batch, "pos", None) - if not isinstance(positions, torch.Tensor) or positions.ndim != 2: - return None - if positions.numel() == 0: - return None - positions = positions.to(device=device, dtype=torch.float32) - if positions.shape[1] < 2: - return None - - batch_index = getattr(batch, "batch", None) - if isinstance(batch_index, torch.Tensor): - batch_index = batch_index.to(device=device).long() - if batch_index.numel() != positions.shape[0]: - return None - elif n_instances == 1: - batch_index = torch.zeros(positions.shape[0], dtype=torch.long, device=device) - else: - return None - - centroids = torch.zeros((n_instances, 2), dtype=torch.float32, device=device) - for inst_idx in range(n_instances): - mask = batch_index == inst_idx - if not torch.any(mask): - continue - centroids[inst_idx, :] = positions[mask][:, :2].mean(dim=0) - return centroids - - def _attention_width(self) -> int: - """Number of attention channels, i.e. 1 (shared) or one per class.""" - if self.attention is None: - return 1 - head = getattr(self.attention, "attention_c", None) - out_features = getattr(head, "out_features", None) - return int(out_features) if out_features is not None else 1 - - def _build_predict_group_metadata( - self, - batch, - *, - patch_logits: torch.Tensor, - graph_emb: torch.Tensor, - ordered_bag_ids: list[str], - bag_indices: list[list[int]], - include_instance_payload: bool = True, - ) -> Dict[str, Any]: - n_instances = int(patch_logits.shape[0]) - region_ids = self._to_optional_str_list( - getattr(batch, "region_id", None), - expected_length=n_instances, - field_name="region_id", - ) - sample_ids = self._to_optional_str_list( - getattr(batch, "sample_id", None), - expected_length=n_instances, - field_name="sample_id", - ) - row_region_ids: list[Optional[str]] = [] - row_sample_ids: list[Optional[str]] = [] - for indices in bag_indices: - row_region_ids.append(region_ids[indices[0]] if indices else None) - row_sample_ids.append(sample_ids[indices[0]] if indices else None) - - metadata: Dict[str, Any] = { - "row_region_ids": row_region_ids, - "row_sample_ids": row_sample_ids, - } - if not include_instance_payload: - return metadata - - patch_ids = self._to_optional_str_list( - getattr(batch, "patch_id", None), - expected_length=n_instances, - field_name="patch_id", - ) - instance_logits_chunks: list[torch.Tensor] = [] - instance_patch_ids: list[str] = [] - instance_bag_ids: list[str] = [] - instance_region_ids: list[Optional[str]] = [] - instance_sample_ids: list[Optional[str]] = [] - instance_attention_logits_chunks: list[torch.Tensor] = [] - instance_embedding_chunks: list[torch.Tensor] = [] - instance_graph_chunks: list[Any] = [] - composition_by_instance = ( - self._compute_instance_composition( - batch, - n_instances=n_instances, - label_name=str(getattr(self, "_predict_composition_label", "cell_type")), - device=patch_logits.device, - ) - if bool(getattr(self, "_predict_emit_instance_composition", False)) - else None - ) - centroids_by_instance = ( - self._compute_instance_centroids( - batch, - n_instances=n_instances, - device=patch_logits.device, - ) - if bool(getattr(self, "_predict_emit_instance_centroids", False)) - else None - ) - instance_composition_chunks: list[torch.Tensor] = [] - instance_centroid_chunks: list[torch.Tensor] = [] - compute_attention_logits = ( - include_instance_payload and self.use_attention and self.attention is not None - ) - # ShaDow batches are built directly and delimited by `ptr`, so - # `to_data_list()` raises on them. This previously sat behind a bare - # `except` that turned that into a silent None, and the spatial-table - # export then failed several steps later complaining about a missing - # payload field instead of naming the cause. - subgraph_data_list = extract_subgraphs(batch) - - for bag_id, indices in zip(ordered_bag_ids, bag_indices): - if not indices: - continue - idx = torch.tensor(indices, dtype=torch.long, device=patch_logits.device) - sub_logits = patch_logits.index_select(0, idx) - instance_logits_chunks.append(sub_logits) - instance_embedding_chunks.append(graph_emb.index_select(0, idx)) - instance_patch_ids.extend( - [patch_ids[i] if patch_ids[i] is not None else str(bag_id) for i in indices] - ) - instance_bag_ids.extend([str(bag_id)] * len(indices)) - instance_region_ids.extend([region_ids[i] for i in indices]) - instance_sample_ids.extend([sample_ids[i] for i in indices]) - if isinstance(subgraph_data_list, list): - for i in indices: - if 0 <= int(i) < len(subgraph_data_list): - instance_graph_chunks.append(subgraph_data_list[int(i)]) - if composition_by_instance is not None: - instance_composition_chunks.append(composition_by_instance.index_select(0, idx)) - if centroids_by_instance is not None: - instance_centroid_chunks.append(centroids_by_instance.index_select(0, idx)) - if compute_attention_logits: - sub_emb = graph_emb.index_select(0, idx) - attn_logits, _ = self.attention(sub_emb) - if attn_logits.dim() == 1: - attn_logits = attn_logits.unsqueeze(-1) - instance_attention_logits_chunks.append(attn_logits) - - metadata["instance_logits"] = ( - torch.cat(instance_logits_chunks, dim=0) - if instance_logits_chunks - else torch.empty((0, patch_logits.shape[-1]), device=patch_logits.device) - ) - metadata["instance_patch_ids"] = instance_patch_ids - metadata["instance_bag_ids"] = instance_bag_ids - metadata["instance_region_ids"] = instance_region_ids - metadata["instance_sample_ids"] = instance_sample_ids - metadata["instance_embeddings"] = ( - torch.cat(instance_embedding_chunks, dim=0) - if instance_embedding_chunks - else torch.empty((0, graph_emb.shape[-1]), device=graph_emb.device) - ) - if composition_by_instance is not None: - metadata["instance_composition"] = ( - torch.cat(instance_composition_chunks, dim=0) - if instance_composition_chunks - else torch.empty( - (0, composition_by_instance.shape[-1]), - device=composition_by_instance.device, - ) - ) - if centroids_by_instance is not None: - metadata["instance_centroids"] = ( - torch.cat(instance_centroid_chunks, dim=0) - if instance_centroid_chunks - else torch.empty((0, 2), device=patch_logits.device) - ) - if compute_attention_logits: - metadata["instance_attention_logits"] = ( - torch.cat(instance_attention_logits_chunks, dim=0) - if instance_attention_logits_chunks - else torch.empty((0, self._attention_width()), device=patch_logits.device) - ) - if instance_graph_chunks: - metadata["instance_graphs"] = instance_graph_chunks - return metadata - - def _build_predict_embedding_payload( - self, - batch, - *, - graph_emb: torch.Tensor, - node_emb: torch.Tensor, - ordered_bag_ids: list[str], - bag_indices: list[list[int]], - include_node_embeddings: bool, - ) -> Dict[str, Any]: - bag_graph_embeddings: list[torch.Tensor] = [] - bag_counts: list[int] = [] - for indices in bag_indices: - if not indices: - continue - idx = torch.tensor(indices, dtype=torch.long, device=graph_emb.device) - bag_graph_embeddings.append(graph_emb.index_select(0, idx).mean(dim=0, keepdim=True)) - bag_counts.append(int(idx.numel())) - - payload: Dict[str, Any] = { - "embedding_bag_ids": [str(v) for v in ordered_bag_ids], - "graph_embeddings": ( - torch.cat(bag_graph_embeddings, dim=0) - if bag_graph_embeddings - else torch.empty((0, graph_emb.shape[-1]), device=graph_emb.device) - ), - "embedding_bag_counts": bag_counts, - } - if include_node_embeddings: - bag_ids = extract_bag_ids( - batch, - bag_key=self.task_cfg.get("bag_key", "region_id"), - bag_fallback_key=self.task_cfg.get("bag_fallback_key", "sample_id"), - ) - batch_index = getattr(batch, "batch", None) - if batch_index is None: - batch_index = torch.zeros( - node_emb.shape[0], dtype=torch.long, device=node_emb.device - ) - node_bag_ids = [ - str(bag_ids[int(graph_idx)]) for graph_idx in batch_index.detach().cpu() - ] - payload["node_embeddings"] = node_emb - payload["node_bag_ids"] = node_bag_ids - return payload - - def _forward_bags(self, batch, *, allow_missing_targets: bool = False): - patch_logits, graph_emb, node_emb = self._forward_instances(batch) - - bag_ids = extract_bag_ids( - batch, - bag_key=self.task_cfg.get("bag_key", "region_id"), - bag_fallback_key=self.task_cfg.get("bag_fallback_key", "sample_id"), - ) - bag_groups = group_instance_indices_by_bag(bag_ids) - - if self.use_attention: - bag_logits, ordered_bag_ids, bag_attention, bag_indices = ( - aggregate_bag_logits_attention( - logits=patch_logits, - embeddings=graph_emb, - attention=self.attention, - bag_groups=bag_groups, - max_instances_per_bag=int(self.task_cfg.get("max_instances_per_bag", 0)), - instance_sampling=self.task_cfg.get("instance_sampling", "all"), - ) - ) - else: - bag_logits, ordered_bag_ids, bag_indices = aggregate_bag_logits_mean( - patch_logits, - bag_groups, - max_instances_per_bag=int(self.task_cfg.get("max_instances_per_bag", 0)), - instance_sampling=self.task_cfg.get("instance_sampling", "all"), - ) - bag_attention = None - - bag_targets = None - bag_weights = None - if hasattr(batch, "graph_y"): - bag_targets, bag_weights = gather_bag_targets( - batch=batch, - bag_indices=bag_indices, - target_columns=self.task_cfg.get("target_columns"), - ) - elif not allow_missing_targets: - raise ValueError("Batch is missing graph_y required for supervised loss computation.") - - return { - "patch_logits": patch_logits, - "graph_emb": graph_emb, - "node_emb": node_emb, - "bag_logits": bag_logits, - "ordered_bag_ids": ordered_bag_ids, - "bag_attention": bag_attention, - "bag_indices": bag_indices, - "bag_targets": bag_targets, - "bag_weights": bag_weights, - } - - # ------------------------------------------------------------------ - # Loss computation - # ------------------------------------------------------------------ - - def _compute_region_loss( - self, - *, - bag_logits: torch.Tensor, - bag_targets: torch.Tensor, - bag_weights: Optional[torch.Tensor], - ) -> torch.Tensor: - return compute_supervised_loss( - loss_fn=self.loss_fn, - task_cfg=self.task_cfg, - bag_logits=bag_logits, - bag_targets=bag_targets, - bag_weights=bag_weights, - ) - - def _compute_losses( - self, - *, - patch_logits: torch.Tensor, - bag_logits: torch.Tensor, - bag_targets: torch.Tensor, - bag_weights: Optional[torch.Tensor], - ordered_bag_ids: list[str], - bag_indices: list[list[int]], - bag_attention: Optional[dict[str, torch.Tensor]], - stage: str, - ) -> torch.Tensor: - """Region cross-entropy is the sole training objective.""" - return self._compute_region_loss( - bag_logits=bag_logits, - bag_targets=bag_targets, - bag_weights=bag_weights, - ) - - def _log_stage_metrics( - self, - *, - stage: str, - loss_value: torch.Tensor, - bag_logits: torch.Tensor, - bag_targets: torch.Tensor, - region_loss: Optional[torch.Tensor] = None, - ) -> None: - if stage not in {"train", "val", "test"}: - return - self.log( - f"{stage}/loss", - loss_value, - on_step=stage in {"train", "val"}, - on_epoch=True, - prog_bar=True, - ) - if self.use_attention and stage == "train" and region_loss is not None: - self.log( - "train/region_loss", - region_loss, - on_step=True, - on_epoch=True, - prog_bar=False, - ) - if stage in {"train", "test"} and self.task_cfg["target_type"] == "binary": - acc = compute_binary_accuracy(bag_logits, bag_targets) - self.log( - f"{stage}/acc", - acc, - on_step=False, - on_epoch=True, - prog_bar=stage != "train", - ) - elif stage in {"train", "test"} and self.task_cfg["target_type"] == "categorical": - acc = compute_categorical_accuracy(bag_logits, bag_targets) - self.log( - f"{stage}/acc", - acc, - on_step=False, - on_epoch=True, - prog_bar=stage != "train", - ) - - # ------------------------------------------------------------------ - # Shared step - # ------------------------------------------------------------------ - - def _shared_step(self, batch, stage: str): - payload = self._forward_bags(batch) - patch_logits = payload["patch_logits"] - bag_logits = payload["bag_logits"] - ordered_bag_ids = payload["ordered_bag_ids"] - bag_attention = payload["bag_attention"] - bag_indices = payload["bag_indices"] - bag_targets = payload["bag_targets"] - bag_weights = payload["bag_weights"] - - region_loss = self._compute_losses( - patch_logits=patch_logits, - bag_logits=bag_logits, - bag_targets=bag_targets, - bag_weights=bag_weights, - ordered_bag_ids=ordered_bag_ids, - bag_indices=bag_indices, - bag_attention=bag_attention, - stage=stage, - ) - self._log_stage_metrics( - stage=stage, - loss_value=region_loss, - bag_logits=bag_logits, - bag_targets=bag_targets, - region_loss=region_loss if self.use_attention else None, - ) - - result = { - "loss": region_loss, - "bag_ids": list(ordered_bag_ids), - "bag_logits": bag_logits, - "bag_targets": bag_targets, - "patch_logits": patch_logits, - "graph_emb": payload["graph_emb"], - "ordered_bag_ids": list(ordered_bag_ids), - "bag_indices": bag_indices, - "bag_attention": bag_attention, - } - if self.use_attention: - result["region_loss"] = region_loss - return result - - # ------------------------------------------------------------------ - # Training / validation / test / predict steps - # ------------------------------------------------------------------ - - def training_step(self, batch, batch_idx): - if not self.region_accum_enabled: - return self._shared_step(batch, "train")["loss"] - - # ---------- Region accumulation path (MIL only) --------------- - payload = self._forward_bags(batch) - bag_logits = payload["bag_logits"] - bag_targets = payload["bag_targets"] - bag_weights = payload["bag_weights"] - ordered_bag_ids = payload["ordered_bag_ids"] - bag_indices = payload["bag_indices"] - bag_attention = payload["bag_attention"] - patch_logits = payload["patch_logits"] - - region_losses: list[torch.Tensor] = [] - buffered_region_losses: list[torch.Tensor] = [] - for idx, bag_id in enumerate(ordered_bag_ids): - region_w = bag_weights[idx : idx + 1] if bag_weights is not None else None - region_loss = self._compute_losses( - patch_logits=patch_logits, - bag_logits=bag_logits[idx : idx + 1], - bag_targets=bag_targets[idx : idx + 1], - bag_weights=region_w, - bag_indices=[bag_indices[idx]], - ordered_bag_ids=[bag_id], - bag_attention={bag_id: bag_attention[bag_id]}, - stage="train", - ) - region_losses.append(region_loss.detach()) - buffered_region_losses.append(region_loss) - - if not buffered_region_losses: - return torch.zeros((), device=self.device) - - # Buffer one scalar per forward pass to avoid chunked backward passes over - # the same computation graph. - step_region_loss = torch.stack(buffered_region_losses).mean() - self._region_loss_buffer.append(step_region_loss) - - flush_losses = self._flush_region_buffer_if_needed(force=False) - for flush_loss in flush_losses: - self.log( - "train/flush_loss", - flush_loss, - on_step=True, - on_epoch=True, - prog_bar=False, - ) - self.log( - "train/region_buffer_size", - float(len(self._region_loss_buffer)), - on_step=True, - on_epoch=False, - prog_bar=False, - ) - mean_region_loss_value = step_region_loss.detach() - mean_region_loss = torch.stack(region_losses).mean() - self._log_stage_metrics( - stage="train", - loss_value=mean_region_loss_value, - bag_logits=bag_logits, - bag_targets=bag_targets, - region_loss=mean_region_loss, - ) - - return mean_region_loss_value - - def validation_step(self, batch, batch_idx): - return self._eval_step(batch, "val") - - def _eval_step(self, batch, stage: str): - """Shared val/test step: buffer instance-level outputs for region-level - aggregation at epoch end, so a region split over several sub-batches is - scored exactly once with one attention softmax over all of its cells.""" - out = self._shared_step(batch, stage) - metadata = self._build_predict_group_metadata( - batch, - patch_logits=out["patch_logits"], - graph_emb=out["graph_emb"], - ordered_bag_ids=out["ordered_bag_ids"], - bag_indices=out["bag_indices"], - ) - self._val_bag_ids.extend([str(v) for v in out["bag_ids"]]) - self._val_bag_logits_chunks.append(out["bag_logits"].detach().cpu()) - self._val_row_region_ids.extend(metadata["row_region_ids"]) - self._val_row_sample_ids.extend(metadata["row_sample_ids"]) - if out["bag_targets"] is not None: - self._val_bag_targets_chunks.append(out["bag_targets"].detach().cpu()) - self._val_instance_logits_chunks.append(metadata["instance_logits"].detach().cpu()) - if str(self.task_cfg.get("bag_key", "region_id")) == "patch_id": - # Patch-level bags: group instances by the bag they were pooled in - # (the tile), not by their ego-graph id, so the epoch-end - # re-aggregation reproduces the training bags exactly. - self._val_instance_patch_ids.extend([str(v) for v in metadata["instance_bag_ids"]]) - else: - self._val_instance_patch_ids.extend([str(v) for v in metadata["instance_patch_ids"]]) - self._val_instance_region_ids.extend(metadata["instance_region_ids"]) - self._val_instance_sample_ids.extend(metadata["instance_sample_ids"]) - if "instance_attention_logits" in metadata: - self._val_instance_attention_logits_chunks.append( - metadata["instance_attention_logits"].detach().cpu() - ) - return out - - def test_step(self, batch, batch_idx): - return self._eval_step(batch, "test") - - def on_test_epoch_start(self) -> None: - self._eval_stage = "test" - self._reset_val_epoch_buffers() - - def on_test_epoch_end(self) -> None: - self._finish_eval_epoch("test") - self._eval_stage = "val" - - def predict_step(self, batch, batch_idx, dataloader_idx=0): - out = self._forward_bags(batch, allow_missing_targets=True) - emit_instance_payload = bool(getattr(self, "_predict_emit_instance_payload", True)) - emit_embeddings_payload = bool(getattr(self, "_predict_emit_embeddings_payload", False)) - emit_node_embeddings = bool(getattr(self, "_predict_emit_node_embeddings", False)) - metadata = self._build_predict_group_metadata( - batch, - patch_logits=out["patch_logits"], - graph_emb=out.get("graph_emb", out["patch_logits"]), - ordered_bag_ids=out["ordered_bag_ids"], - bag_indices=out["bag_indices"], - include_instance_payload=emit_instance_payload, - ) - result: Dict[str, Any] = { - "bag_ids": list(out["ordered_bag_ids"]), - "bag_logits": out["bag_logits"].detach(), - "row_region_ids": metadata["row_region_ids"], - "row_sample_ids": metadata["row_sample_ids"], - } - if out["bag_targets"] is not None: - result["bag_targets"] = out["bag_targets"].detach() - if emit_instance_payload: - result["instance_logits"] = metadata["instance_logits"].detach() - result["instance_patch_ids"] = metadata["instance_patch_ids"] - result["instance_bag_ids"] = metadata["instance_bag_ids"] - result["instance_region_ids"] = metadata["instance_region_ids"] - result["instance_sample_ids"] = metadata["instance_sample_ids"] - if "instance_graphs" in metadata: - result["instance_graphs"] = metadata["instance_graphs"] - if bool(getattr(self, "_predict_emit_instance_embeddings", False)): - result["instance_embeddings"] = metadata["instance_embeddings"].detach() - if "instance_composition" in metadata: - result["instance_composition"] = metadata["instance_composition"].detach() - if "instance_centroids" in metadata: - result["instance_centroids"] = metadata["instance_centroids"].detach() - if self.use_attention: - result["bag_attention"] = out["bag_attention"] - if emit_instance_payload: - result["instance_attention_logits"] = metadata[ - "instance_attention_logits" - ].detach() - if emit_embeddings_payload: - embedding_payload = self._build_predict_embedding_payload( - batch, - graph_emb=out.get("graph_emb", out["patch_logits"]), - node_emb=out.get("node_emb", out.get("graph_emb", out["patch_logits"])), - ordered_bag_ids=out["ordered_bag_ids"], - bag_indices=out["bag_indices"], - include_node_embeddings=emit_node_embeddings, - ) - result["embedding_bag_ids"] = embedding_payload["embedding_bag_ids"] - result["graph_embeddings"] = embedding_payload["graph_embeddings"].detach() - result["embedding_bag_counts"] = embedding_payload["embedding_bag_counts"] - if emit_node_embeddings: - result["node_embeddings"] = embedding_payload["node_embeddings"].detach() - result["node_bag_ids"] = embedding_payload["node_bag_ids"] - return result - - # ------------------------------------------------------------------ - # Region accumulation helpers (MIL only) - # ------------------------------------------------------------------ - - def _flush_region_buffer_if_needed(self, *, force: bool) -> list[torch.Tensor]: - if not self._region_loss_buffer: - return [] - if not force and len(self._region_loss_buffer) < self.region_accum_hyperbatch_size: - return [] - - optimizer = self.optimizers() - if isinstance(optimizer, (list, tuple)): - optimizer = optimizer[0] - flush_losses: list[torch.Tensor] = [] - - while self._region_loss_buffer: - if force: - chunk = self._region_loss_buffer - else: - if len(self._region_loss_buffer) < self.region_accum_hyperbatch_size: - break - chunk = self._region_loss_buffer[: self.region_accum_hyperbatch_size] - - flush_loss = torch.stack(chunk, dim=0).mean() - self.manual_backward(flush_loss) - if self.grad_clip_norm is not None: - self.clip_gradients( - optimizer, - gradient_clip_val=float(self.grad_clip_norm), - gradient_clip_algorithm="norm", - ) - scaler = self._amp_scaler() - scale_before = float(scaler.get_scale()) if scaler is not None else None - optimizer.step() - optimizer.zero_grad(set_to_none=True) - # Under mixed precision a step with inf/nan gradients is skipped and the - # loss scale is reduced; the preprint only advances the schedule on - # steps that actually updated the weights. - stepped = scaler is None or float(scaler.get_scale()) >= scale_before - if stepped: - scheduler = self.lr_schedulers() - if scheduler is not None: - if isinstance(scheduler, (list, tuple)): - for sch in scheduler: - sch.step() - else: - scheduler.step() - else: - self._skipped_optimizer_steps += 1 - self._manual_optimizer_steps += 1 - self.log( - "train/manual_optimizer_steps", - float(self._manual_optimizer_steps), - on_step=not force, - on_epoch=force, - prog_bar=False, - ) - flush_losses.append(flush_loss.detach()) - - del self._region_loss_buffer[: len(chunk)] - if force: - break - - return flush_losses - - def _amp_scaler(self): - """The GradScaler of the active mixed-precision plugin, if any.""" - try: - plugin = self.trainer.precision_plugin - except RuntimeError: - return None - return getattr(plugin, "scaler", None) - - def on_train_start(self) -> None: - if self.region_accum_enabled: - self._region_loss_buffer.clear() - - def on_validation_epoch_start(self) -> None: - self._reset_val_epoch_buffers() - - def on_validation_epoch_end(self) -> None: - self._finish_eval_epoch("val") - - def _finish_eval_epoch(self, stage: str) -> None: - payload = self._build_validation_epoch_payload() - if payload is None or payload.bag_targets is None: - self._reset_val_epoch_buffers() - return - # Region-level metrics are computed per bag unit: the task's bag key - # decides whether a unit is a region (default), a patch/tile or a sample. - bag_key = str(self.task_cfg.get("bag_key", "region_id")) - bag_scope = {"patch_id": "patch", "sample_id": "sample"}.get(bag_key, "region") - aggregated = aggregate_group_logits( - payload, - mode="attention_weighted" if self.use_attention else "mean", - bag_scope=bag_scope, - subsample_fraction=1.0, - subsample_seed=None, - ) - if aggregated.bag_targets is None: - self._reset_val_epoch_buffers() - return - target_type = str(self.task_cfg.get("target_type", "binary")) - metrics = compute_task_metrics_from_tensors( - target_type=target_type, - logits=aggregated.bag_logits, - targets=aggregated.bag_targets, - threshold=0.5, - ) - region_metrics = self._region_level_metrics( - target_type, aggregated.bag_logits, aggregated.bag_targets - ) - - def _log(name: str, value: float, prog_bar: bool = False) -> None: - self.log( - f"{stage}/{name}", - torch.tensor(float(value), device=self.device), - on_step=False, - on_epoch=True, - prog_bar=prog_bar, - ) - - if target_type in {"binary", "categorical"}: - _log("acc", metrics["accuracy"], prog_bar=True) - for name, value in region_metrics.items(): - _log(name, value, prog_bar=name == "auroc") - elif target_type == "regression": - _log("mae", metrics["mae"], prog_bar=True) - _log("rmse", metrics["rmse"], prog_bar=True) - _log("r2", metrics["r2"]) - elif target_type == "survival": - _log("c_index", metrics["c_index"], prog_bar=True) - - if stage == "test": - self.test_predictions = { - "bag_ids": list(aggregated.bag_ids), - "logits": aggregated.bag_logits.detach().cpu(), - "targets": aggregated.bag_targets.detach().cpu(), - "region_ids": list(getattr(aggregated, "row_region_ids", None) or []), - "sample_ids": list(getattr(aggregated, "row_sample_ids", None) or []), - "target_type": target_type, - } - self._reset_val_epoch_buffers() - - @staticmethod - def _region_level_metrics( - target_type: str, logits: torch.Tensor, targets: torch.Tensor - ) -> Dict[str, float]: - """Unweighted region loss, AUROC and balanced accuracy on aggregated region logits. - - These are the quantities the preprint selects and early-stops on: the loss - of the region-level prediction (not of training sub-batches) and the - AUROC of the positive-class probability. - """ - import numpy as np - from sklearn.metrics import balanced_accuracy_score, roc_auc_score - - out: Dict[str, float] = {} - logits = logits.detach().float().cpu() - y = targets.detach().cpu() - if target_type == "categorical": - y_int = y.long().view(-1) - out["loss_region"] = float(torch.nn.functional.cross_entropy(logits, y_int)) - probs = torch.softmax(logits, dim=1).numpy() - yi = y_int.numpy() - n_classes = probs.shape[1] - try: - if n_classes == 2: - out["auroc"] = float(roc_auc_score(yi, probs[:, 1])) - else: - out["auroc"] = float( - roc_auc_score(yi, probs, multi_class="ovr", average="macro") - ) - except ValueError: - out["auroc"] = float("nan") - pred = probs.argmax(axis=1) - out["bal_acc"] = ( - float(balanced_accuracy_score(yi, pred)) - if len(np.unique(yi)) > 1 - else float("nan") - ) - out["pred_pos_frac"] = float((pred == 1).mean()) if n_classes == 2 else float("nan") - elif target_type == "binary": - yb = y.float().view(-1) - lg = logits.view(-1) - out["loss_region"] = float( - torch.nn.functional.binary_cross_entropy_with_logits(lg, yb) - ) - prob = torch.sigmoid(lg).numpy() - yi = yb.numpy().astype(int) - try: - out["auroc"] = float(roc_auc_score(yi, prob)) - except ValueError: - out["auroc"] = float("nan") - pred = (prob >= 0.5).astype(int) - out["bal_acc"] = ( - float(balanced_accuracy_score(yi, pred)) - if len(np.unique(yi)) > 1 - else float("nan") - ) - out["pred_pos_frac"] = float(pred.mean()) - return out - - def on_train_epoch_end(self) -> None: - if not self.region_accum_enabled: - return - flush_losses = self._flush_region_buffer_if_needed(force=True) - for flush_loss in flush_losses: - self.log( - "train/epoch_end_flush_loss", - flush_loss, - on_step=False, - on_epoch=True, - prog_bar=False, - ) - - # ------------------------------------------------------------------ - # Optimizer / scheduler - # ------------------------------------------------------------------ - - def configure_optimizers(self): - if not self._built: - self.setup("fit") - if self.optim_cfg is None: - raise ValueError( - "Missing optimizer config for training. Compose an optim group " - "(for example '+optim=adamw') or set model.optim explicitly." - ) - - params = self._build_optimizer_params() - optimizer = instantiate_optimizer(self.optim_cfg, params) - scheduler = instantiate_scheduler_with_warmup( - self.scheduler_cfg, - optimizer, - warmup_cfg=self.task_cfg.get("lr_warmup"), - t_max_auto=self._resolve_t_max_auto(), - ) - if scheduler is None: - return optimizer - return { - "optimizer": optimizer, - "lr_scheduler": { - "scheduler": scheduler, - "interval": self.scheduler_interval, - "frequency": 1, - }, - } - - def _build_optimizer_params(self): - """Parameters or parameter groups for the optimizer. - - ``task.optimization`` accepts either the per-component multipliers - ``encoder_lr_scale`` / ``head_lr_scale`` / ``attention_lr_scale`` (with - ``head_weight_decay``), or the absolute pair ``backbone_lr`` + - ``attention_lr``. Setting only one of the absolute pair is an error rather - than a silent no-op. - """ - opt = self.optimization_cfg - scale_keys = { - "encoder_lr_scale", - "head_lr_scale", - "attention_lr_scale", - "head_weight_decay", - } - has_scales = any(opt.get(k) is not None for k in scale_keys) - backbone_lr = opt.get("backbone_lr") - attention_lr = opt.get("attention_lr") - if (backbone_lr is None) != (attention_lr is None): - raise ValueError( - "task.optimization.backbone_lr and attention_lr must be set together " - "(or use encoder_lr_scale / head_lr_scale / attention_lr_scale)." - ) - if has_scales and backbone_lr is not None: - raise ValueError( - "task.optimization mixes absolute (backbone_lr/attention_lr) and scaled " - "learning rates; use one scheme." - ) - - base_cfg = ( - OmegaConf.to_container(self.optim_cfg, resolve=True) - if isinstance(self.optim_cfg, DictConfig) - else dict(self.optim_cfg) - ) - base_lr = float(base_cfg.get("lr", 1e-3)) - base_wd = float(base_cfg.get("weight_decay", 0.0) or 0.0) - - if has_scales or self._encoder_frozen or self.freeze_head: - groups = build_param_groups( - base_lr=base_lr, - base_weight_decay=base_wd, - encoder=self.encoder, - head=self.graph_head, - attention=self.attention if self.use_attention else None, - encoder_lr_scale=0.0 - if self._encoder_frozen - else float(opt.get("encoder_lr_scale", 1.0) or 0.0), - head_lr_scale=float(opt.get("head_lr_scale", 1.0) or 0.0), - attention_lr_scale=float(opt.get("attention_lr_scale", 1.0) or 0.0), - head_weight_decay=opt.get("head_weight_decay"), - freeze_head=self.freeze_head, - ) - return groups - - params = [p for p in self.parameters() if p.requires_grad] - if self.use_attention and backbone_lr is not None: - # Absolute scheme: one backbone group (encoder + head) and one attention group. - seen: set[int] = set() - - def _collect(module) -> list[torch.nn.Parameter]: - out: list[torch.nn.Parameter] = [] - for p in module.parameters(): - if p.requires_grad and id(p) not in seen: - seen.add(id(p)) - out.append(p) - return out - - backbone = _collect(self.encoder) + _collect(self.graph_head) - attention = _collect(self.attention) - groups: list[dict] = [] - if backbone: - groups.append({"params": backbone, "lr": float(backbone_lr)}) - if attention: - groups.append({"params": attention, "lr": float(attention_lr)}) - return groups or params - return params - - def _resolve_t_max_auto(self) -> Optional[int]: - """Optimizer steps for ``scheduler.T_max: auto``.""" - try: - trainer = self.trainer - except RuntimeError: - return None - if trainer is None: - return None - max_epochs = int(trainer.max_epochs or 1) - if self.region_accum_enabled: - n_units = None - dm = getattr(trainer, "datamodule", None) - ds = getattr(dm, "dataset_train", None) if dm is not None else None - if ds is not None: - try: - n_units = len(ds) - except TypeError: - n_units = None - if n_units is None: - return None - return resolve_t_max_auto( - num_train_units=int(n_units), - max_epochs=max_epochs, - hyperbatch_size=self.region_accum_hyperbatch_size, - ) - try: - return int(trainer.estimated_stepping_batches) - except Exception: - return None - - def _reset_val_epoch_buffers(self) -> None: - self._val_bag_ids: list[str] = [] - self._val_bag_logits_chunks: list[torch.Tensor] = [] - self._val_bag_targets_chunks: list[torch.Tensor] = [] - self._val_row_region_ids: list[Optional[str]] = [] - self._val_row_sample_ids: list[Optional[str]] = [] - self._val_instance_logits_chunks: list[torch.Tensor] = [] - self._val_instance_attention_logits_chunks: list[torch.Tensor] = [] - self._val_instance_patch_ids: list[str] = [] - self._val_instance_region_ids: list[Optional[str]] = [] - self._val_instance_sample_ids: list[Optional[str]] = [] - - @staticmethod - def _cat_or_none(chunks: list[torch.Tensor]) -> Optional[torch.Tensor]: - if not chunks: - return None - return torch.cat(chunks, dim=0) - - def _gather_validation_objects(self, obj: dict[str, Any]) -> list[dict[str, Any]]: - if not dist.is_available() or not dist.is_initialized(): - return [obj] - gathered: list[dict[str, Any]] = [None] * dist.get_world_size() # type: ignore[list-item] - dist.all_gather_object(gathered, obj) - return gathered - - def _build_validation_epoch_payload(self) -> Optional[BatchPredictionPayload]: - local_payload = { - "bag_ids": self._val_bag_ids, - "bag_logits": self._cat_or_none(self._val_bag_logits_chunks), - "bag_targets": self._cat_or_none(self._val_bag_targets_chunks), - "row_region_ids": self._val_row_region_ids, - "row_sample_ids": self._val_row_sample_ids, - "instance_logits": self._cat_or_none(self._val_instance_logits_chunks), - "instance_attention_logits": self._cat_or_none( - self._val_instance_attention_logits_chunks - ), - "instance_patch_ids": self._val_instance_patch_ids, - "instance_region_ids": self._val_instance_region_ids, - "instance_sample_ids": self._val_instance_sample_ids, - } - gathered = self._gather_validation_objects(local_payload) - - bag_ids: list[str] = [] - bag_logits_chunks: list[torch.Tensor] = [] - bag_targets_chunks: list[torch.Tensor] = [] - row_region_ids: list[Optional[str]] = [] - row_sample_ids: list[Optional[str]] = [] - instance_logits_chunks: list[torch.Tensor] = [] - instance_attention_logits_chunks: list[torch.Tensor] = [] - instance_patch_ids: list[str] = [] - instance_region_ids: list[Optional[str]] = [] - instance_sample_ids: list[Optional[str]] = [] - - for chunk in gathered: - chunk_bag_logits = chunk.get("bag_logits") - if chunk_bag_logits is not None and int(chunk_bag_logits.shape[0]) > 0: - bag_logits_chunks.append(chunk_bag_logits) - bag_ids.extend([str(v) for v in chunk["bag_ids"]]) - row_region_ids.extend( - [None if v is None else str(v) for v in chunk["row_region_ids"]] - ) - row_sample_ids.extend( - [None if v is None else str(v) for v in chunk["row_sample_ids"]] - ) - chunk_bag_targets = chunk.get("bag_targets") - if chunk_bag_targets is not None and int(chunk_bag_targets.shape[0]) > 0: - bag_targets_chunks.append(chunk_bag_targets) - chunk_instance_logits = chunk.get("instance_logits") - if chunk_instance_logits is not None and int(chunk_instance_logits.shape[0]) > 0: - instance_logits_chunks.append(chunk_instance_logits) - instance_patch_ids.extend([str(v) for v in chunk.get("instance_patch_ids", [])]) - instance_region_ids.extend( - [None if v is None else str(v) for v in chunk.get("instance_region_ids", [])] - ) - instance_sample_ids.extend( - [None if v is None else str(v) for v in chunk.get("instance_sample_ids", [])] - ) - chunk_instance_attn = chunk.get("instance_attention_logits") - if chunk_instance_attn is not None and int(chunk_instance_attn.shape[0]) > 0: - instance_attention_logits_chunks.append(chunk_instance_attn) - - if not bag_logits_chunks: - return None - bag_targets = self._cat_or_none(bag_targets_chunks) - instance_logits = self._cat_or_none(instance_logits_chunks) - instance_attention_logits = self._cat_or_none(instance_attention_logits_chunks) - return BatchPredictionPayload( - bag_ids=bag_ids, - bag_logits=torch.cat(bag_logits_chunks, dim=0), - bag_targets=bag_targets, - bag_attention=None, - row_region_ids=row_region_ids, - row_sample_ids=row_sample_ids, - instance_logits=instance_logits, - instance_attention_logits=instance_attention_logits, - instance_patch_ids=instance_patch_ids if instance_logits is not None else None, - instance_region_ids=instance_region_ids if instance_logits is not None else None, - instance_sample_ids=instance_sample_ids if instance_logits is not None else None, - ) diff --git a/src/grass_mil/models/training/__init__.py b/src/grass_mil/models/training/__init__.py deleted file mode 100644 index b839648..0000000 --- a/src/grass_mil/models/training/__init__.py +++ /dev/null @@ -1,80 +0,0 @@ -from .bagging import ( - aggregate_bag_logits_attention, - aggregate_bag_logits_mean, - extract_bag_ids, - gather_bag_targets, - group_instance_indices_by_bag, - maybe_sample_indices, - select_target_columns, -) -from .augmentations import ( - build_augmentation, - build_finetune_augmentation, - degree_importance, - drop_edges, - drop_features, - drop_importance, - feature_noise, -) -from .class_weights import collect_unit_labels, inverse_frequency_class_weights -from .builders import ( - build_supervised_components, - infer_categorical_binding, - infer_encoder_input_dim, - resolve_encoder_cfg, - validate_task_config, -) -from .checkpoint_init import ( - read_graph_head_bias, - load_state_dict_with_optional_mapping, - remap_encoder_keys, -) -from .loss_utils import ( - compute_binary_accuracy, - compute_categorical_accuracy, - compute_supervised_loss, -) -from .optimization import ( - build_param_groups, - instantiate_optimizer, - instantiate_scheduler, - instantiate_scheduler_with_warmup, - resolve_t_max_auto, -) -from .ssl_runtime import CosineWarmup - -__all__ = [ - "aggregate_bag_logits_attention", - "aggregate_bag_logits_mean", - "extract_bag_ids", - "gather_bag_targets", - "group_instance_indices_by_bag", - "maybe_sample_indices", - "select_target_columns", - "build_supervised_components", - "build_augmentation", - "build_finetune_augmentation", - "build_param_groups", - "collect_unit_labels", - "inverse_frequency_class_weights", - "resolve_t_max_auto", - "degree_importance", - "drop_edges", - "drop_features", - "drop_importance", - "feature_noise", - "infer_categorical_binding", - "infer_encoder_input_dim", - "resolve_encoder_cfg", - "validate_task_config", - "load_state_dict_with_optional_mapping", - "read_graph_head_bias", - "remap_encoder_keys", - "compute_binary_accuracy", - "compute_categorical_accuracy", - "compute_supervised_loss", - "instantiate_optimizer", - "instantiate_scheduler", - "instantiate_scheduler_with_warmup", - "CosineWarmup", -] diff --git a/src/grass_mil/models/training/augmentations.py b/src/grass_mil/models/training/augmentations.py deleted file mode 100644 index c157ec5..0000000 --- a/src/grass_mil/models/training/augmentations.py +++ /dev/null @@ -1,381 +0,0 @@ -"""Graph augmentations for self-supervised pretraining. - -BGRL needs two stochastic views of each ego-graph. The production views combine -a degree-importance-weighted node/edge drop with additive Gaussian noise on the -continuous cell-size feature. - -Ported from ``graph_augmentations_phenotype`` in the spatial-augmentations -tree. Two deviations from that code are deliberate and marked below; one -deviation from the manuscript is deliberate too -- the manuscript describes a -min-max degree normalisation, but the final run code uses mean-max, and the -code is authoritative. - -Every function is pure: it takes an explicit ``torch.Generator`` and returns a -new graph rather than mutating its input. -""" - -from __future__ import annotations - -from typing import Any, Callable, Mapping, Optional, Sequence - -import torch - -__all__ = [ - "drop_importance", - "drop_edges", - "drop_features", - "feature_noise", - "build_augmentation", - "degree_importance", -] - - -def _clone(data): - return data.clone() - - -def degree_importance(edge_index: torch.Tensor, num_nodes: int) -> torch.Tensor: - """Per-node importance from log degree, normalised to ``[0, 1]``. - - Uses the legacy **mean-max** normalisation ``(v - mean) / (max - mean)`` - clamped at zero, which maps every below-average-degree node to 0. The - manuscript describes min-max; the final run code does this, and the code is - the reference. - """ - from torch_geometric.utils import degree - - deg = degree(edge_index[0], num_nodes).float() - log_deg = torch.log1p(deg) - importance = (log_deg - log_deg.mean()) / (log_deg.max() - log_deg.mean() + 1e-8) - return torch.clamp(importance, min=0.0) - - -def _drop_probability(importance: torch.Tensor, *, mu: float, p_lambda: float) -> torch.Tensor: - """``min((1 - importance) * mu, p_lambda)``. - - Note the legacy name for this quantity was ``keep_prob``, but it is used as - ``rand > p`` to decide retention, so it is the *drop* probability. The - behaviour is preserved; only the name is corrected. - """ - ceiling = torch.full_like(importance, float(p_lambda)) - return torch.minimum((1.0 - importance) * float(mu), ceiling) - - -def _mutual_edge_mask(edge_index: torch.Tensor, keep: torch.Tensor) -> torch.Tensor: - """Restrict ``keep`` to edges whose reverse also survives. - - Dropping one direction of an undirected edge leaves an asymmetric graph and - makes message passing direction-dependent. - """ - if edge_index.numel() == 0: - return keep - num_nodes = int(edge_index.max()) + 1 - src, dst = edge_index[0], edge_index[1] - kept_keys = (src[keep] * num_nodes + dst[keep]).to(torch.int64) - reverse_keys = (dst * num_nodes + src).to(torch.int64) - return keep & torch.isin(reverse_keys, kept_keys) - - -def drop_importance( - data, - *, - mu: float, - p_lambda: float, - categorical_column: Optional[int] = 0, - unassigned_index: Optional[int] = None, - generator: Optional[torch.Generator] = None, -): - """Drop low-degree nodes and edges, keeping the graph undirected. - - Dropped nodes keep their position in the graph but have their features - zeroed and their cell type set to the reserved "unassigned" code, so the - encoder sees a masked node rather than a spurious cell type. - """ - if not 0.0 < mu <= 1.0: - raise ValueError(f"mu must be within (0, 1], got {mu}.") - if not 0.0 < p_lambda <= 1.0: - raise ValueError(f"p_lambda must be within (0, 1], got {p_lambda}.") - - out = _clone(data) - edge_index = out.edge_index - num_nodes = int(getattr(out, "num_nodes", out.x.size(0))) - importance = degree_importance(edge_index, num_nodes) - - node_drop_p = _drop_probability(importance, mu=mu, p_lambda=p_lambda) - draws = torch.rand(node_drop_p.shape, generator=generator, device=node_drop_p.device) - dropped = ~(draws > node_drop_p) - - if out.x is not None and out.x.numel(): - # Clone before the in-place write: `x` may be a view of the source - # graph, in which case masking here would corrupt the original. - x = out.x.clone() - x[dropped] = 0 - out.x = x - - codes = getattr(out, "categorical_codes", None) - if codes is not None and unassigned_index is not None and categorical_column is not None: - codes = codes.clone() - codes[dropped, int(categorical_column)] = int(unassigned_index) - out.categorical_codes = codes - - if edge_index.numel(): - edge_importance = (importance[edge_index[0]] + importance[edge_index[1]]) / 2.0 - edge_importance = (edge_importance - edge_importance.mean()) / ( - edge_importance.max() - edge_importance.mean() + 1e-8 - ) - edge_importance = torch.clamp(edge_importance, min=0.0) - edge_drop_p = _drop_probability(edge_importance, mu=mu, p_lambda=p_lambda) - edge_draws = torch.rand(edge_drop_p.shape, generator=generator, device=edge_drop_p.device) - keep = edge_draws > edge_drop_p - keep = _mutual_edge_mask(edge_index, keep) - - out.edge_index = edge_index[:, keep] - edge_attr = getattr(out, "edge_attr", None) - if edge_attr is not None: - out.edge_attr = edge_attr[keep] - return out - - -def drop_edges( - data, - *, - p: float, - force_undirected: bool = True, - generator: Optional[torch.Generator] = None, -): - """Uniformly drop edges. - - ``force_undirected`` keeps both directions of an edge together: the legacy - finetune transform left it off while pretraining had it on, so the two - regimes augmented differently and message passing became - direction-dependent. - - The draw is taken here rather than through ``torch_geometric.utils - .dropout_edge``, which offers no generator argument and would silently - consume the global RNG, making seeded runs irreproducible. - """ - if p <= 0.0: - return data - if not 0.0 <= p < 1.0: - raise ValueError(f"drop probability must be within [0, 1), got {p}.") - - out = _clone(data) - edge_index = out.edge_index - if edge_index.numel() == 0: - return out - - # Draw on the generator's device (CPU) and move the mask: a CPU generator - # cannot drive a CUDA draw, and fine-tuning batches live on the GPU. - draws = torch.rand((edge_index.size(1),), generator=generator).to(edge_index.device) - keep = draws >= float(p) - if force_undirected: - keep = _mutual_edge_mask(edge_index, keep) - - out.edge_index = edge_index[:, keep] - edge_attr = getattr(out, "edge_attr", None) - if edge_attr is not None: - out.edge_attr = edge_attr[keep] - return out - - -def drop_features( - data, - *, - p: float, - generator: Optional[torch.Generator] = None, - mode: str = "column", - categorical_column: Optional[int] = None, - unassigned_index: Optional[int] = None, -): - """Feature dropout. - - ``mode="column"`` zeroes whole feature columns with probability ``p`` (the - BGRL view generator). ``mode="elementwise"`` is the fine-tuning transform of - the preprint: every entry of ``x`` is zeroed independently with probability - ``p``, and a node whose draw falls below ``p`` has its categorical code - (column ``categorical_column`` of ``categorical_codes``) rewritten to - ``unassigned_index`` so the cell type is dropped as well. - """ - if p <= 0.0: - return data - if mode not in {"column", "elementwise"}: - raise ValueError(f"drop_features mode must be 'column' or 'elementwise', got {mode!r}.") - if mode == "column": - if data.x is None or data.x.numel() == 0: - return data - out = _clone(data) - x = out.x.clone() - draws = torch.rand((x.size(1),), generator=generator).to(x.device) - x[:, draws < float(p)] = 0 - out.x = x - return out - - out = _clone(data) - num_nodes = int(getattr(out, "num_nodes", 0) or (out.x.size(0) if out.x is not None else 0)) - if out.x is not None and out.x.numel() > 0: - x = out.x.clone() - draws = torch.rand(x.shape, generator=generator).to(x.device) - x[draws < float(p)] = 0 - out.x = x - codes = getattr(out, "categorical_codes", None) - if ( - unassigned_index is not None - and isinstance(codes, torch.Tensor) - and codes.numel() > 0 - and num_nodes > 0 - ): - col = 0 if categorical_column is None else int(categorical_column) - if col < codes.size(1): - node_draws = torch.rand((codes.size(0),), generator=generator).to(codes.device) - codes = codes.clone() - codes[node_draws < float(p), col] = int(unassigned_index) - out.categorical_codes = codes - return out - - -def build_finetune_augmentation( - cfg: Optional[Mapping[str, Any]], - *, - unassigned_index: Optional[int] = None, - categorical_column: Optional[int] = None, - seed: Optional[int] = None, -) -> Optional[Callable[[Any], Any]]: - """Training-time augmentation for supervised fine-tuning. - - ``cfg`` keys: ``drop_edge_p`` (undirected edge drop, both directions kept - together) and ``drop_feat_p`` (element-wise feature drop that also writes the - unassigned cell-type token). Returns ``None`` when both are zero, so the - default fine-tuning path is untouched. - """ - cfg = dict(cfg or {}) - drop_edge_p = float(cfg.get("drop_edge_p", 0.0) or 0.0) - drop_feat_p = float(cfg.get("drop_feat_p", 0.0) or 0.0) - if drop_edge_p <= 0.0 and drop_feat_p <= 0.0: - return None - generator = torch.Generator() - generator.manual_seed(int(cfg.get("seed", seed if seed is not None else 0))) - - def _augment(data): - out = data - if drop_edge_p > 0.0: - out = drop_edges(out, p=drop_edge_p, force_undirected=True, generator=generator) - if drop_feat_p > 0.0: - out = drop_features( - out, - p=drop_feat_p, - generator=generator, - mode="elementwise", - categorical_column=categorical_column, - unassigned_index=unassigned_index, - ) - return out - - return _augment - - -def feature_noise( - data, - *, - std: float, - feature_indices: Sequence[int], - generator: Optional[torch.Generator] = None, -): - """Add Gaussian noise to selected continuous feature columns.""" - if std <= 0.0 or data.x is None or data.x.numel() == 0: - return data - indices = [int(i) for i in feature_indices] - if not indices: - return data - out = _clone(data) - x = out.x.clone() - for index in indices: - if index < 0 or index >= x.size(1): - raise ValueError( - f"feature noise column {index} is out of range for x with " - f"{x.size(1)} column(s)." - ) - noise = torch.randn(x[:, index].shape, generator=generator, device=x.device) - x[:, index] = x[:, index] + noise * float(std) - out.x = x - return out - - -def resolve_feature_indices( - columns: Optional[Sequence[Any]], - *, - feature_names: Optional[Sequence[str]] = None, -) -> list[int]: - """Resolve feature columns by name, falling back to integer indices. - - Resolving by name matters: the legacy default noised a hardcoded column - index, which pointed at cell size in one cohort and at an unrelated feature - in another, feeding pure noise to the encoder. - """ - if not columns: - return [] - names = list(feature_names or []) - resolved: list[int] = [] - for column in columns: - if isinstance(column, int): - resolved.append(int(column)) - continue - name = str(column) - if name not in names: - raise ValueError( - f"Feature column '{name}' not found in molecular feature names {names}. " - "Set the column by integer index if the datamodule does not record names." - ) - resolved.append(names.index(name)) - return resolved - - -def build_augmentation( - cfg: Optional[Mapping[str, Any]], - *, - feature_names: Optional[Sequence[str]] = None, - unassigned_index: Optional[int] = None, - categorical_column: Optional[int] = 0, -) -> Callable[[Any, Optional[torch.Generator]], Any]: - """Build a view-generating transform from config. - - Returns a callable ``(data, generator) -> data``. - """ - cfg = dict(cfg or {}) - mode = str(cfg.get("mode", "uniform")).strip().lower() - - if mode == "importance": - mu = float(cfg.get("mu", 0.25)) - p_lambda = float(cfg.get("p_lambda", 0.45)) - std = float(cfg.get("feature_noise_std", 0.0)) - indices = resolve_feature_indices( - cfg.get("feature_noise_columns"), feature_names=feature_names - ) - - def _importance(data, generator: Optional[torch.Generator] = None): - out = drop_importance( - data, - mu=mu, - p_lambda=p_lambda, - categorical_column=categorical_column, - unassigned_index=unassigned_index, - generator=generator, - ) - if std > 0.0 and indices: - out = feature_noise(out, std=std, feature_indices=indices, generator=generator) - return out - - return _importance - - if mode == "uniform": - drop_edge_p = float(cfg.get("drop_edge_p", 0.0)) - drop_feat_p = float(cfg.get("drop_feat_p", 0.0)) - - def _uniform(data, generator: Optional[torch.Generator] = None): - out = drop_features(data, p=drop_feat_p, generator=generator) - return drop_edges(out, p=drop_edge_p, force_undirected=True, generator=generator) - - return _uniform - - raise ValueError( - f"Unsupported augmentation mode '{mode}'. Expected one of: 'importance', 'uniform'." - ) diff --git a/src/grass_mil/models/training/bagging.py b/src/grass_mil/models/training/bagging.py deleted file mode 100644 index d4135d8..0000000 --- a/src/grass_mil/models/training/bagging.py +++ /dev/null @@ -1,272 +0,0 @@ -from __future__ import annotations - -from typing import Any, Dict, List, Optional, Sequence - -import torch - -from grass_mil.contracts import validate_attention_width - -try: - from torch_geometric.data import Data -except Exception as exc: # pragma: no cover - raise ImportError("torch_geometric is required for training bagging helpers") from exc - - -def extract_bag_ids( - batch: Data, - *, - bag_key: str, - bag_fallback_key: str, -) -> List[str]: - bag_values = _to_list(getattr(batch, bag_key, None)) - fallback_values = _to_list(getattr(batch, bag_fallback_key, None)) - if bag_values is None and fallback_values is None: - raise ValueError( - f"Batch is missing both '{bag_key}' and '{bag_fallback_key}' for bagging." - ) - if bag_values is None: - return [str(v) for v in fallback_values] - if fallback_values is None: - return [str(v) for v in bag_values] - if len(bag_values) != len(fallback_values): - raise ValueError( - f"Batch key lengths mismatch for bagging: '{bag_key}' has {len(bag_values)} " - f"values while fallback '{bag_fallback_key}' has {len(fallback_values)}." - ) - region_values = _to_list(getattr(batch, "region_id", None)) - if region_values is not None and len(region_values) != len(bag_values): - if len(region_values) == 1: - region_values = region_values * len(bag_values) - else: - raise ValueError( - "Batch key lengths mismatch for bagging: 'region_id' has " - f"{len(region_values)} values while '{bag_key}' has {len(bag_values)}." - ) - - resolved: List[str] = [] - for idx, (primary, fallback) in enumerate(zip(bag_values, fallback_values)): - region_value = region_values[idx] if region_values is not None else None - use_fallback = _is_missing_value(primary) - if use_fallback and _is_missing_value(fallback): - raise ValueError( - f"Both '{bag_key}' and '{bag_fallback_key}' are missing for at least one instance." - ) - if bag_key == "patch_id": - if use_fallback: - if not _is_missing_value(region_value): - resolved.append(f"{str(fallback)}::{str(region_value)}") - else: - resolved.append(str(fallback)) - continue - if not _is_missing_value(fallback): - if not _is_missing_value(region_value): - resolved.append(f"{str(fallback)}::{str(region_value)}::{str(primary)}") - else: - resolved.append(f"{str(fallback)}::{str(primary)}") - continue - resolved.append(str(primary)) - continue - if use_fallback: - resolved.append(str(fallback)) - continue - # Namespacing primary bag ids by fallback ids avoids accidental cross-sample - # merges when primary ids (for example region_id) are only locally unique. - if bag_key != bag_fallback_key and not _is_missing_value(fallback): - resolved.append(f"{str(fallback)}::{str(primary)}") - continue - resolved.append(str(primary)) - return resolved - - -def group_instance_indices_by_bag(bag_ids: Sequence[str]) -> Dict[str, List[int]]: - grouped: Dict[str, List[int]] = {} - for idx, bag_id in enumerate(bag_ids): - grouped.setdefault(bag_id, []).append(idx) - return grouped - - -def maybe_sample_indices( - indices: List[int], - *, - max_instances_per_bag: int, - instance_sampling: str, -) -> List[int]: - if max_instances_per_bag <= 0 or len(indices) <= max_instances_per_bag: - return indices - if instance_sampling == "all": - return indices[:max_instances_per_bag] - rand_perm = torch.randperm(len(indices)) - keep = rand_perm[:max_instances_per_bag].tolist() - keep.sort() - return [indices[i] for i in keep] - - -def aggregate_bag_logits_mean( - logits: torch.Tensor, - bag_groups: Dict[str, List[int]], - *, - max_instances_per_bag: int, - instance_sampling: str, -) -> tuple[torch.Tensor, List[str], List[List[int]]]: - bag_logits = [] - bag_ids = [] - bag_indices = [] - for bag_id, indices in bag_groups.items(): - selected = maybe_sample_indices( - indices, - max_instances_per_bag=max_instances_per_bag, - instance_sampling=instance_sampling, - ) - bag_ids.append(bag_id) - bag_indices.append(selected) - bag_logits.append(logits[selected].mean(dim=0, keepdim=True)) - return torch.cat(bag_logits, dim=0), bag_ids, bag_indices - - -def aggregate_bag_logits_attention( - *, - logits: torch.Tensor, - embeddings: torch.Tensor, - attention: torch.nn.Module, - bag_groups: Dict[str, List[int]], - max_instances_per_bag: int, - instance_sampling: str, -) -> tuple[torch.Tensor, List[str], Dict[str, torch.Tensor], List[List[int]]]: - """Pool instance logits into bag logits with gated attention. - - Attention is normalised over the instances of a bag, **independently per - class**: - - .. math:: A_{i,c} = \\mathrm{softmax}_i(\\alpha_{i,c}), \\quad - L_c = \\sum_{i \\in B} A_{i,c}\\, \\ell_{i,c} - - Two attention widths are legal. ``n_classes == num_classes`` is the - published per-class form, and is what makes the attribution identity - ``sum_i A[i,c] * l[i,c] == L_c`` exact. ``n_classes == 1`` keeps a single - shared attention channel broadcast across the class logits. - - Returns: - Bag logits, bag ids, per-bag attention of shape ``[n_instances, - n_attention_classes]``, and the selected instance indices per bag. - """ - bag_logits = [] - bag_ids = [] - bag_attention = {} - bag_indices = [] - for bag_id, indices in bag_groups.items(): - selected = maybe_sample_indices( - indices, - max_instances_per_bag=max_instances_per_bag, - instance_sampling=instance_sampling, - ) - sub_emb = embeddings[selected] - sub_logits = logits[selected] - attn_logits, _ = attention(sub_emb) - if attn_logits.dim() == 1: - attn_logits = attn_logits.unsqueeze(-1) - validate_attention_width( - attention_width=int(attn_logits.size(-1)), - num_classes=int(sub_logits.size(-1)), - ) - # Softmax down the instance axis, per class column. - attn_scores = torch.softmax(attn_logits, dim=0) - weighted_logits = (sub_logits * attn_scores).sum(dim=0, keepdim=True) - bag_logits.append(weighted_logits) - bag_ids.append(bag_id) - bag_indices.append(selected) - bag_attention[bag_id] = attn_scores - return torch.cat(bag_logits, dim=0), bag_ids, bag_attention, bag_indices - - -def _find_label_index_by_name(label_names: Any, target_name: str) -> int: - if isinstance(label_names, (list, tuple)): - if target_name in label_names: - return int(label_names.index(target_name)) - raise ValueError(f"Target '{target_name}' not found in graph_label_names={label_names}.") - - -def select_target_columns( - graph_y: torch.Tensor, - batch: Data, - target_columns: Optional[Sequence[str]], -) -> torch.Tensor: - if graph_y.ndim == 1: - graph_y = graph_y.unsqueeze(-1) - if not target_columns: - return graph_y - label_names = getattr(batch, "graph_label_names", None) - # PyG collation may produce a list-of-lists (one per graph). Normalize to one - # shared label-name list, but do not collapse plain list[str] labels. - if ( - isinstance(label_names, list) - and len(label_names) == graph_y.shape[0] - and label_names - and isinstance(label_names[0], (list, tuple)) - ): - label_names = label_names[0] - indices = [_find_label_index_by_name(label_names, name) for name in target_columns] - return graph_y[:, indices] - - -def gather_bag_targets( - *, - batch: Data, - bag_indices: List[List[int]], - target_columns: Optional[Sequence[str]], -) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - if not hasattr(batch, "graph_y"): - raise ValueError("Batch is missing graph_y required for supervised training.") - graph_y = select_target_columns(batch.graph_y.float(), batch, target_columns) - graph_w = batch.graph_w.float() if hasattr(batch, "graph_w") else None - bag_targets: List[torch.Tensor] = [] - bag_weights = None - bag_weight_rows: List[torch.Tensor] = [] - for idxs in bag_indices: - bag_y = graph_y[idxs] - ref_y = bag_y[:1] - if ( - bag_y.shape[0] > 1 - and not torch.isclose(bag_y, ref_y.expand_as(bag_y), atol=0.0, rtol=0.0).all() - ): - raise ValueError( - "Inconsistent graph_y values detected within the same bag. " - "All instances in a bag must share one graph target." - ) - bag_targets.append(ref_y) - if graph_w is not None: - bag_w = graph_w[idxs] - ref_w = bag_w[:1] - if ( - bag_w.shape[0] > 1 - and not torch.isclose(bag_w, ref_w.expand_as(bag_w), atol=0.0, rtol=0.0).all() - ): - raise ValueError("Inconsistent graph_w values detected within the same bag.") - bag_weight_rows.append(ref_w) - bag_targets_t = torch.cat(bag_targets, dim=0) - if graph_w is not None: - bag_weights = torch.cat(bag_weight_rows, dim=0) - return bag_targets_t, bag_weights - - -def _to_list(values: Any) -> Optional[List[Any]]: - if values is None: - return None - if isinstance(values, torch.Tensor): - values = values.tolist() - if isinstance(values, (str, bytes)): - return [values] - try: - return list(values) - except TypeError: - return [values] - - -def _is_missing_value(value: Any) -> bool: - if value is None: - return True - if isinstance(value, str): - text = value.strip().lower() - return text in {"", "nan", "none", "null"} - if isinstance(value, float): - return value != value - return False diff --git a/src/grass_mil/models/training/builders.py b/src/grass_mil/models/training/builders.py deleted file mode 100644 index 604c24c..0000000 --- a/src/grass_mil/models/training/builders.py +++ /dev/null @@ -1,212 +0,0 @@ -from __future__ import annotations - -from dataclasses import asdict -from typing import Any, Dict, Optional - -import lightning as L -import torch - -from grass_mil.contracts import validate_attention_width -from grass_mil.models.components.embeddings import resolve_categorical_embedding_config -from grass_mil.models.components import ( - EncoderConfig, - build_attention, - build_encoder, - build_graph_head, - build_loss, - build_ssl, -) - - -def infer_encoder_input_dim(module: L.LightningModule, fallback: int = 0) -> int: - try: - trainer = module.trainer - except RuntimeError: - return fallback - if trainer is None or trainer.datamodule is None: - return fallback - datamodule = trainer.datamodule - if not hasattr(datamodule, "dataset_train") or datamodule.dataset_train is None: - # Datamodule may not have been set up yet; trigger setup so we can - # inspect the first sample and infer input_dim. - if hasattr(datamodule, "setup"): - datamodule.setup("fit") - if not hasattr(datamodule, "dataset_train") or datamodule.dataset_train is None: - return fallback - dataset = datamodule.dataset_train - if len(dataset) == 0: - return fallback - item = dataset[0] - x = getattr(item, "x", None) - if x is None: - return fallback - return int(x.shape[1]) - - -def infer_categorical_binding( - module: L.LightningModule, label: str -) -> tuple[Optional[int], Optional[int]]: - """Resolve ``(num_embeddings, column_index)`` for a categorical node label. - - The vocabulary size comes from the label map that precompute persists in - ``processed_index.json``, and the column index from the first sample's - ``categorical_slices``. Mirrors ``infer_encoder_input_dim``: best-effort, - returning ``None`` when the datamodule cannot be inspected, so the caller - can fall back to explicit config. - """ - try: - trainer = module.trainer - except RuntimeError: - return None, None - if trainer is None or trainer.datamodule is None: - return None, None - datamodule = trainer.datamodule - - dataset = getattr(datamodule, "dataset_train", None) - if dataset is None and hasattr(datamodule, "setup"): - datamodule.setup("fit") - dataset = getattr(datamodule, "dataset_train", None) - if dataset is None or len(dataset) == 0: - return None, None - - # `label_maps` lives on the graph dataset; unwrap any transform wrapper. - inner = dataset - while not hasattr(inner, "label_maps") and hasattr(inner, "dataset"): - inner = inner.dataset - - num_embeddings: Optional[int] = None - label_maps = getattr(inner, "label_maps", None) - if isinstance(label_maps, dict) and label in label_maps: - num_embeddings = len(label_maps[label]) - - column_index: Optional[int] = None - slices = getattr(dataset[0], "categorical_slices", None) - if isinstance(slices, dict) and label in slices: - column_index = int(slices[label]) - return num_embeddings, column_index - - -def validate_task_config(task_cfg: Dict[str, Any]) -> None: - target_type = task_cfg.get("target_type", "binary") - if target_type not in {"binary", "categorical", "regression", "survival"}: - raise ValueError( - f"Unsupported task.target_type='{target_type}'. " - "Expected one of: binary, categorical, regression, survival." - ) - instance_sampling = task_cfg.get("instance_sampling", "all") - if instance_sampling not in {"all", "random"}: - raise ValueError( - "task.instance_sampling must be one of ['all', 'random'], " - f"got '{instance_sampling}'." - ) - if target_type == "survival": - # Without these checks a mis-shaped survival config fails deep inside - # the Cox loss with an opaque shape error. - target_columns = task_cfg.get("target_columns") - if target_columns is not None and len(target_columns) != 2: - raise ValueError( - "Survival tasks need exactly two target columns, ordered " - f"[time, event]; got {list(target_columns)}." - ) - if task_cfg.get("loss") not in (None, "survival_coxsgd"): - raise ValueError( - "task.target_type='survival' requires task.loss='survival_coxsgd', " - f"got '{task_cfg.get('loss')}'." - ) - - -def resolve_encoder_cfg( - encoder_cfg: Dict[str, Any], - inferred_input_dim: int, - *, - module: Optional[L.LightningModule] = None, -) -> Dict[str, Any]: - cfg = dict(encoder_cfg) - categorical_cfg = cfg.get("categorical_embedding") - if cfg.get("input_dim", 0) in (None, 0): - if inferred_input_dim <= 0 and categorical_cfg is None: - raise ValueError( - "Could not infer encoder input_dim from datamodule; " - "set model.encoder.input_dim explicitly." - ) - # A cell-type-only cohort legitimately has zero continuous features; the - # categorical embedding then carries the whole input representation. - cfg["input_dim"] = max(int(inferred_input_dim), 0) - - if categorical_cfg is not None: - label = str(dict(categorical_cfg).get("label", "cell_type")) - num_embeddings, column_index = ( - infer_categorical_binding(module, label) if module is not None else (None, None) - ) - label_map = infer_categorical_label_map(module, label) if module is not None else None - resolved = resolve_categorical_embedding_config( - categorical_cfg, - cardinalities={label: num_embeddings} if num_embeddings is not None else None, - slices={label: column_index} if column_index is not None else None, - label_maps={label: label_map} if label_map else None, - ) - cfg["categorical_embedding"] = asdict(resolved) - return cfg - - -def infer_categorical_label_map( - module: Optional[L.LightningModule], label: str -) -> Optional[Dict[str, int]]: - """Return ``{category: code}`` for ``label`` from the attached datamodule, if any.""" - if module is None: - return None - try: - trainer = module.trainer - except RuntimeError: - return None - if trainer is None or trainer.datamodule is None: - return None - dataset = getattr(trainer.datamodule, "dataset_train", None) - if dataset is None: - return None - inner = dataset - while not hasattr(inner, "label_maps") and hasattr(inner, "dataset"): - inner = inner.dataset - label_maps = getattr(inner, "label_maps", None) - if isinstance(label_maps, dict) and label in label_maps: - return {str(k): int(v) for k, v in label_maps[label].items()} - return None - - -def build_supervised_components( - *, - module: L.LightningModule, - encoder_cfg: Dict[str, Any], - graph_head_cfg: Dict[str, Any], - attention_cfg: Optional[Dict[str, Any]], - use_attention: bool, - loss_cfg: Dict[str, Any], - use_ssl: bool = False, - ssl_cfg: Optional[Dict[str, Any]] = None, -) -> Dict[str, torch.nn.Module]: - inferred_input_dim = infer_encoder_input_dim(module) - resolved_encoder_cfg = resolve_encoder_cfg(encoder_cfg, inferred_input_dim, module=module) - encoder = build_encoder(EncoderConfig(**resolved_encoder_cfg)) - ssl_model = build_ssl(use_ssl, encoder, ssl_cfg) if use_ssl else None - graph_head = build_graph_head(graph_head_cfg) - attention = None - if use_attention: - if attention_cfg is None: - raise ValueError("attention config is required when use_attention=True.") - cfg = dict(attention_cfg) - attention_type = cfg.pop("attention_type", "gated_projected") - # Fail at build time rather than part-way through the first batch. - validate_attention_width( - attention_width=int(cfg.get("n_classes", 1)), - num_classes=int(graph_head_cfg.get("output_dim", 1)), - ) - attention = build_attention(True, attention_type=attention_type, **cfg) - loss_name = loss_cfg.pop("loss_type") - loss_fn = build_loss(loss_name, **loss_cfg) - return { - "encoder": encoder, - "ssl_model": ssl_model, - "graph_head": graph_head, - "attention": attention, - "loss_fn": loss_fn, - } diff --git a/src/grass_mil/models/training/checkpoint_init.py b/src/grass_mil/models/training/checkpoint_init.py deleted file mode 100644 index 8a9adfc..0000000 --- a/src/grass_mil/models/training/checkpoint_init.py +++ /dev/null @@ -1,136 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from typing import Dict, List, Optional - -import torch - - -def load_state_dict_with_optional_mapping( - module: torch.nn.Module, - *, - init_from_ckpt: Optional[str], - init_strict: bool, - encoder_init_map: str, -) -> Dict[str, List[str]]: - if not init_from_ckpt: - return {"missing_keys": [], "unexpected_keys": []} - ckpt_path = Path(init_from_ckpt) - if not ckpt_path.exists(): - raise FileNotFoundError(f"Checkpoint not found: {init_from_ckpt}") - # First-party Lightning checkpoint; see the note in data/components/datasets.py - # on why weights_only must be set explicitly under torch>=2.6. - raw = torch.load(str(ckpt_path), map_location="cpu", weights_only=False) - state_dict = raw["state_dict"] if isinstance(raw, dict) and "state_dict" in raw else raw - if not isinstance(state_dict, dict): - raise ValueError("Checkpoint must resolve to a state_dict dictionary.") - remapped = remap_encoder_keys(state_dict, encoder_init_map=encoder_init_map) - missing, unexpected = module.load_state_dict(remapped, strict=init_strict) - return {"missing_keys": list(missing), "unexpected_keys": list(unexpected)} - - -def remap_encoder_keys( - state_dict: Dict[str, torch.Tensor], *, encoder_init_map: str -) -> Dict[str, torch.Tensor]: - if encoder_init_map == "identity": - return dict(state_dict) - if encoder_init_map == "spacegm_bgrl": - return remap_spacegm_bgrl_keys(state_dict) - if encoder_init_map != "auto_bgrl_or_identity": - raise ValueError( - "encoder_init_map must be one of ['identity', 'auto_bgrl_or_identity', 'spacegm_bgrl']" - ) - remapped: Dict[str, torch.Tensor] = {} - for key, value in state_dict.items(): - new_key = key - if key.startswith("online_encoder."): - new_key = f"encoder.{key[len('online_encoder.') :]}" - elif key.startswith("ssl_model.online_encoder."): - new_key = f"encoder.{key[len('ssl_model.online_encoder.') :]}" - elif key.startswith("model.online_encoder."): - new_key = f"encoder.{key[len('model.online_encoder.') :]}" - remapped[new_key] = value - return remapped - - -def remap_spacegm_bgrl_keys(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - """Import an encoder pretrained with the original SPACE-GM style BGRL code. - - Maps ``online_encoder.x_embedding`` -> ``encoder.input_proj.embedding``, - ``online_encoder.feat_embedding`` -> ``encoder.input_proj.linear``, - ``online_encoder.gnns.{i}.gine.{nn.0,nn.2,lin,eps}`` -> ``encoder.layers.{i}.*`` - and ``online_encoder.batch_norms.{i}`` -> ``encoder.norms.{i}``. The unused - GIN MLP copies, edge-type embeddings, the target encoder and the predictor - are dropped. The continuous input width must match the public model's - (``input_proj.linear.weight`` has one column per continuous node feature); - select columns beforehand if the legacy run carried masked extra columns. - """ - import re - - out: Dict[str, torch.Tensor] = {} - layer_re = re.compile( - r"^online_encoder\.gnns\.(\d+)\.gine\.(nn\.0|nn\.2|lin|eps)(\.(weight|bias))?$" - ) - norm_re = re.compile( - r"^online_encoder\.batch_norms\.(\d+)\.(weight|bias|running_mean|running_var|num_batches_tracked)$" - ) - for key, value in state_dict.items(): - if key == "online_encoder.x_embedding.weight": - out["encoder.input_proj.embedding.weight"] = value - continue - if key in ("online_encoder.feat_embedding.weight", "online_encoder.feat_embedding.bias"): - out[key.replace("online_encoder.feat_embedding", "encoder.input_proj.linear")] = value - continue - m = layer_re.match(key) - if m: - suffix = m.group(2) + (m.group(3) or "") - out[f"encoder.layers.{m.group(1)}.{suffix}"] = value - continue - m = norm_re.match(key) - if m: - out[f"encoder.norms.{m.group(1)}.{m.group(2)}"] = value - continue - # gnns.{i}.mlp.* (unused duplicate), edge_embedding, target_encoder, predictor: dropped - if not out: - raise ValueError( - "No SPACE-GM BGRL encoder keys (online_encoder.*) found in the checkpoint." - ) - return out - - -def read_graph_head_bias( - checkpoint_path: str, - *, - head_prefix: str = "graph_head", -) -> Optional[List[float]]: - """Read the graph head's output bias from a training checkpoint. - - Attribution subtracts this bias so the reported margin reflects the - instance-driven part of the decision rather than the head's prior. The bias - lives on the final ``Linear`` of the head MLP, which is the highest-numbered - ``net..bias`` entry under ``head_prefix``. - - Returns: - One value per class, or ``None`` when the checkpoint carries no head - (an encoder-only checkpoint, for instance). - """ - import re - - import torch - - payload = torch.load(checkpoint_path, map_location="cpu", weights_only=False) - state = payload.get("state_dict", payload) if isinstance(payload, dict) else payload - if not isinstance(state, dict): - raise ValueError(f"Checkpoint {checkpoint_path!r} holds no state dict.") - - pattern = re.compile(rf"^{re.escape(head_prefix)}\.net\.(\d+)\.bias$") - candidates = {} - for key, value in state.items(): - match = pattern.match(str(key)) - if match is not None: - candidates[int(match.group(1))] = value - if not candidates: - return None - - bias = candidates[max(candidates)] - return [float(v) for v in bias.detach().cpu().reshape(-1)] diff --git a/src/grass_mil/models/training/class_weights.py b/src/grass_mil/models/training/class_weights.py deleted file mode 100644 index 4957e0f..0000000 --- a/src/grass_mil/models/training/class_weights.py +++ /dev/null @@ -1,67 +0,0 @@ -"""Class weights for the region cross-entropy, computed from the TRAINING split only. - -The preprint weights class ``c`` by inverse frequency on the training regions, -``w_c = n / (k * count_c)``, normalised to mean 1, so the weights change the -balance between classes but not the overall loss scale. Computing them on the -training split alone keeps the test fold's label prior out of the objective. -""" - -from __future__ import annotations - -from typing import Any, Iterable, Optional, Sequence - -import torch - -__all__ = ["inverse_frequency_class_weights", "collect_unit_labels"] - - -def inverse_frequency_class_weights( - labels: Sequence[int] | torch.Tensor, - num_classes: int, - *, - normalize_mean: bool = True, -) -> torch.Tensor: - """``w_c = n / (k * count_c)`` over ``labels`` (absent classes get weight 1).""" - if num_classes < 1: - raise ValueError("num_classes must be >= 1") - y = torch.as_tensor(list(labels) if not isinstance(labels, torch.Tensor) else labels) - y = y.long().view(-1) - if y.numel() == 0: - return torch.ones(num_classes, dtype=torch.float32) - counts = torch.bincount(y, minlength=num_classes).float() - weights = torch.ones(num_classes, dtype=torch.float32) - present = counts > 0 - weights[present] = float(y.numel()) / (float(int(present.sum())) * counts[present]) - if normalize_mean: - weights = weights / weights.mean() - return weights - - -def collect_unit_labels( - dataset: Iterable[Any], - *, - target_columns: Optional[Sequence[str]] = None, -) -> list[int]: - """Read the integer class label of every unit of a graph dataset. - - Units carry ``graph_y`` (shape ``[1, T]`` or ``[T]``); the first target column - is used unless ``target_columns`` names one and the unit exposes - ``graph_label_names``. - """ - labels: list[int] = [] - for data in dataset: - y = getattr(data, "graph_y", None) - if y is None: - continue - y = torch.as_tensor(y).view(-1) - col = 0 - names = getattr(data, "graph_label_names", None) - if target_columns and isinstance(names, (list, tuple)): - flat = list(names[0]) if names and isinstance(names[0], (list, tuple)) else list(names) - if target_columns[0] in flat: - col = int(flat.index(target_columns[0])) - value = float(y[col]) - if value != value: # NaN: unlabelled unit - continue - labels.append(int(round(value))) - return labels diff --git a/src/grass_mil/models/training/loss_utils.py b/src/grass_mil/models/training/loss_utils.py deleted file mode 100644 index 8d50aef..0000000 --- a/src/grass_mil/models/training/loss_utils.py +++ /dev/null @@ -1,51 +0,0 @@ -from __future__ import annotations - -from typing import Any, Dict, Optional - -import torch - - -def compute_supervised_loss( - *, - loss_fn: torch.nn.Module, - task_cfg: Dict[str, Any], - bag_logits: torch.Tensor, - bag_targets: torch.Tensor, - bag_weights: Optional[torch.Tensor], -) -> torch.Tensor: - target_type = task_cfg["target_type"] - if target_type == "binary": - if bag_targets.ndim == 1: - bag_targets = bag_targets.unsqueeze(-1) - return loss_fn(bag_logits, bag_targets, bag_weights) - if target_type == "categorical": - if bag_logits.ndim != 2: - raise ValueError( - "Categorical task expects logits with shape [N, C] for cross-entropy." - ) - targets = bag_targets.long().view(-1) - return loss_fn(bag_logits, targets, bag_weights) - if target_type == "regression": - return loss_fn(bag_logits, bag_targets, bag_weights) - if bag_targets.ndim != 2 or bag_targets.shape[1] < 2: - raise ValueError("Survival task expects at least two target columns [time, event].") - return loss_fn( - bag_logits.squeeze(-1), - bag_targets[:, 0].float(), - bag_targets[:, 1].float(), - ) - - -def compute_binary_accuracy(logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor: - if target.ndim == 1: - target = target.unsqueeze(-1) - pred = (torch.sigmoid(logits) >= 0.5).float() - return (pred == target.float()).float().mean() - - -def compute_categorical_accuracy(logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor: - if logits.ndim != 2: - raise ValueError("Categorical accuracy expects logits with shape [N, C].") - pred = torch.argmax(logits, dim=1) - true = target.long().view(-1) - return (pred == true).float().mean() diff --git a/src/grass_mil/models/training/optimization.py b/src/grass_mil/models/training/optimization.py deleted file mode 100644 index ebaae01..0000000 --- a/src/grass_mil/models/training/optimization.py +++ /dev/null @@ -1,168 +0,0 @@ -from __future__ import annotations - -from typing import Any, Dict, Iterable, Optional - -import hydra -import torch -from omegaconf import DictConfig, OmegaConf - - -def instantiate_optimizer( - cfg: Dict[str, Any] | DictConfig, - params: Iterable[torch.nn.Parameter], -) -> torch.optim.Optimizer: - opt_cfg = cfg if isinstance(cfg, DictConfig) else OmegaConf.create(cfg) - # When *params* is a list of param-group dicts (used for per-component - # learning rates), Hydra's instantiate would wrap them into OmegaConf - # containers and lose tensor references. Fall back to direct construction. - if isinstance(params, list) and params and isinstance(params[0], dict): - resolved = OmegaConf.to_container(opt_cfg, resolve=True) - target = resolved.pop("_target_") - resolved.pop("_recursive_", None) - resolved.pop("_convert_", None) - cls = hydra.utils.get_class(target) - return cls(params, **resolved) - return hydra.utils.instantiate(opt_cfg, params=params) - - -def instantiate_scheduler( - cfg: Optional[Dict[str, Any] | DictConfig], - optimizer: torch.optim.Optimizer, - *, - t_max_auto: Optional[int] = None, -): - """Instantiate a scheduler; ``T_max: auto`` is replaced by ``t_max_auto`` steps.""" - if cfg is None: - return None - sch_cfg = cfg if isinstance(cfg, DictConfig) else OmegaConf.create(cfg) - resolved = OmegaConf.to_container(sch_cfg, resolve=True) - if isinstance(resolved, dict) and str(resolved.get("T_max", "")).lower() == "auto": - if t_max_auto is None or int(t_max_auto) < 1: - raise ValueError( - "scheduler.T_max='auto' needs the number of optimizer steps; the module " - "computes it from max_epochs, the training set size and the hyperbatch." - ) - resolved["T_max"] = int(t_max_auto) - sch_cfg = OmegaConf.create(resolved) - return hydra.utils.instantiate(sch_cfg, optimizer=optimizer) - - -def resolve_t_max_auto(*, num_train_units: int, max_epochs: int, hyperbatch_size: int) -> int: - """Optimizer steps of a run: ``max_epochs * ceil(num_train_units / hyperbatch)``. - - The preprint anneals the cosine schedule over every optimizer step of the - run (``T_max = epochs * n_train // hyperbatch``), so the learning rate never - restarts. - """ - hb = max(1, int(hyperbatch_size)) - steps_per_epoch = max(1, (int(num_train_units) + hb - 1) // hb) - return max(1, int(max_epochs)) * steps_per_epoch - - -def build_param_groups( - *, - base_lr: float, - base_weight_decay: float, - encoder: Optional[torch.nn.Module], - head: Optional[torch.nn.Module], - attention: Optional[torch.nn.Module], - encoder_lr_scale: float = 1.0, - head_lr_scale: float = 1.0, - attention_lr_scale: float = 1.0, - head_weight_decay: Optional[float] = None, - freeze_head: bool = False, -) -> list[dict]: - """Parameter groups with per-component learning-rate multipliers. - - A scale of 0 removes the component from the optimizer (the caller also sets - ``requires_grad=False`` and keeps it in eval mode). Groups are ordered - encoder, head, attention; ``head_weight_decay`` overrides the optimizer's - weight decay for the head group only (the preprint uses 1e-3 on a pre-fitted - linear head). - """ - seen: set[int] = set() - - def _params(module: Optional[torch.nn.Module]) -> list[torch.nn.Parameter]: - out: list[torch.nn.Parameter] = [] - if module is None: - return out - for p in module.parameters(): - if not p.requires_grad or id(p) in seen: - continue - seen.add(id(p)) - out.append(p) - return out - - groups: list[dict] = [] - if encoder is not None and float(encoder_lr_scale) > 0.0: - enc = _params(encoder) - if enc: - groups.append( - {"params": enc, "lr": float(base_lr) * float(encoder_lr_scale), "name": "encoder"} - ) - if head is not None and not freeze_head and float(head_lr_scale) > 0.0: - hp = _params(head) - if hp: - group = {"params": hp, "lr": float(base_lr) * float(head_lr_scale), "name": "head"} - if head_weight_decay is not None: - group["weight_decay"] = float(head_weight_decay) - groups.append(group) - if attention is not None and float(attention_lr_scale) > 0.0: - ap = _params(attention) - if ap: - groups.append( - { - "params": ap, - "lr": float(base_lr) * float(attention_lr_scale), - "name": "attention", - } - ) - if not groups: - raise ValueError( - "No trainable parameters: every component is frozen or has a zero learning-rate scale." - ) - for g in groups: - g.setdefault("weight_decay", float(base_weight_decay)) - return groups - - -def instantiate_scheduler_with_warmup( - cfg: Optional[Dict[str, Any] | DictConfig], - optimizer: torch.optim.Optimizer, - warmup_cfg: Optional[Dict[str, Any] | DictConfig] = None, - *, - t_max_auto: Optional[int] = None, -): - warm_cfg = ( - warmup_cfg if isinstance(warmup_cfg, DictConfig) else OmegaConf.create(warmup_cfg or {}) - ) - warmup_enabled = bool(warm_cfg.get("enabled", False)) - if not warmup_enabled: - return instantiate_scheduler(cfg, optimizer, t_max_auto=t_max_auto) - - warmup_steps = int(warm_cfg.get("warmup_steps", 0)) - if warmup_steps <= 0: - return instantiate_scheduler(cfg, optimizer, t_max_auto=t_max_auto) - - start_factor = float(warm_cfg.get("start_factor", 0.1)) - if start_factor <= 0: - start_factor = 1e-4 - if start_factor > 1: - start_factor = 1.0 - - warmup_scheduler = torch.optim.lr_scheduler.LinearLR( - optimizer=optimizer, - start_factor=start_factor, - end_factor=1.0, - total_iters=warmup_steps, - ) - - base_scheduler = instantiate_scheduler(cfg, optimizer, t_max_auto=t_max_auto) - if base_scheduler is None: - return warmup_scheduler - - return torch.optim.lr_scheduler.SequentialLR( - optimizer=optimizer, - schedulers=[warmup_scheduler, base_scheduler], - milestones=[warmup_steps], - ) diff --git a/src/grass_mil/models/training/probe.py b/src/grass_mil/models/training/probe.py deleted file mode 100644 index db376c1..0000000 --- a/src/grass_mil/models/training/probe.py +++ /dev/null @@ -1,275 +0,0 @@ -"""Linear probe pre-fit on a frozen encoder (preprint "probe head"). - -The preprint initialises the linear instance head of the probehead / ftALL -regimes from a logistic regression fitted on ego-graph embeddings of the -frozen (BGRL-pretrained) encoder: every ego-graph inherits the label of its -region, a ``StandardScaler`` + ``LogisticRegression`` is fitted on the training -regions, and the scaler is folded into the weights so the head acts on raw -embeddings. For a binary task the class-0 row stays zero (the logistic decision -function becomes the class-1 logit); for more classes the multinomial -coefficients are used directly. - -Nothing here is specific to the preprint sizes: the number of roots per region, -the seed, the regulariser and the ego-graph settings are all arguments (the CLI -reads them from ``configs/probe.yaml`` and from ``data.sampler``). -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Dict, List, Optional, Sequence - -import numpy as np -import torch - -__all__ = [ - "ProbeEmbeddings", - "ProbeFit", - "attach_datamodule", - "collect_frozen_embeddings", - "fit_probe", - "fold_probe_into_linear_head", - "probe_sampler_config", - "region_mean_auroc", - "save_probe_head", - "load_probe_head", -] - - -@dataclass -class ProbeEmbeddings: - """Ego-graph embeddings with their region ids and inherited region labels.""" - - embeddings: np.ndarray # (n_instances, dim) - labels: np.ndarray # (n_instances,) int - region_ids: np.ndarray # (n_instances,) str - label_by_region: Dict[str, int] = field(default_factory=dict) - - -@dataclass -class ProbeFit: - weight: torch.Tensor # (num_classes, dim), scaler folded in - bias: torch.Tensor # (num_classes,) - metrics: Dict[str, float] = field(default_factory=dict) - settings: Dict[str, Any] = field(default_factory=dict) - - -def attach_datamodule(model: Any, datamodule: Any, trainer: Any) -> None: - """Bind a datamodule to a model outside ``Trainer.fit`` so ``setup`` can inspect it. - - ``SupervisedModule.setup`` infers the encoder input width and the categorical - vocabulary from ``trainer.datamodule``; this mirrors what Lightning does at - the start of ``fit``/``test``. - """ - trainer._data_connector.attach_data(model, datamodule=datamodule) - model.trainer = trainer - datamodule.trainer = trainer - - -def probe_sampler_config( - base: Any, - *, - roots_per_unit: Optional[int], - seed: int, - batch_size: int, -) -> Any: - """Evaluation sampler for the probe: fixed uniform roots per region, no bagging. - - ``base`` is the datamodule's training ``SamplerConfig``; its ego-graph settings - (depth, neighbour limit, radius cap, interior seeds) are kept. - """ - from grass_mil.data.components.samplers import SamplerConfig - - if base is None or getattr(base, "name", "identity") == "identity": - raise ValueError( - "The probe needs an ego-graph sampler (data.sampler.name=shadow_custom with " - "runtime.enabled=true, e.g. data/sampler=ego_graph_mil)." - ) - runtime = dict(base.runtime or {}) - runtime.update( - { - "enabled": True, - "subsample_fraction": 1.0, - "subsample_count": int(roots_per_unit) if roots_per_unit else None, - "subsample_seed": None, - "fixed_per_unit_seed": int(seed), - "proportional_root_sampling": False, - "unit_as_bag": False, - "merge_units": False, - "shuffle_subgraphs": False, - "subgraph_batch_size": int(batch_size), - } - ) - return SamplerConfig(name=base.name, kwargs=dict(base.kwargs or {}), runtime=runtime) - - -def _region_ids_of(batch: Any, n: int) -> List[str]: - from grass_mil.models.training.bagging import extract_bag_ids - - ids = extract_bag_ids(batch, bag_key="region_id", bag_fallback_key="sample_id") - if len(ids) == 1 and n > 1: - ids = ids * n - if len(ids) != n: - raise ValueError(f"Batch carries {len(ids)} region ids for {n} ego-graphs.") - return [str(v) for v in ids] - - -def _labels_of(batch: Any, n: int, target_columns: Optional[Sequence[str]]) -> np.ndarray: - from grass_mil.models.training.bagging import select_target_columns - - if not hasattr(batch, "graph_y"): - raise ValueError("Batch has no graph_y; the probe needs region labels.") - y = select_target_columns(batch.graph_y.float(), batch, target_columns) - y = y.view(y.shape[0], -1)[:, 0] - if y.shape[0] == 1 and n > 1: - y = y.expand(n) - return np.rint(y.detach().cpu().numpy()).astype(int) - - -@torch.no_grad() -def collect_frozen_embeddings( - model: Any, - dataset: Any, - *, - sampler_config: Any, - target_columns: Optional[Sequence[str]] = None, - device: Optional[torch.device | str] = None, - num_workers: int = 0, -) -> ProbeEmbeddings: - """Encode ``roots_per_unit`` ego-graphs per region with the frozen encoder.""" - from grass_mil.data.components.samplers import get_sampler_strategy - - strategy = get_sampler_strategy(sampler_config) - loader = strategy.build_dataset_loader( - dataset=dataset, batch_size=1, num_workers=num_workers, pin_memory=False, shuffle=False - ) - device = torch.device(device) if device is not None else next(model.parameters()).device - was_training = model.training - model.eval() - embs: List[np.ndarray] = [] - labels: List[np.ndarray] = [] - rids: List[np.ndarray] = [] - for batch in loader: - batch = batch.to(device) - _, graph_emb, _ = model._forward_instances(batch) - n = int(graph_emb.shape[0]) - embs.append(graph_emb.detach().float().cpu().numpy()) - labels.append(_labels_of(batch, n, target_columns)) - rids.append(np.asarray(_region_ids_of(batch, n), dtype=object)) - if was_training: - model.train() - if not embs: - raise ValueError("No ego-graphs were produced; is the dataset split empty?") - e = np.concatenate(embs) - y = np.concatenate(labels) - r = np.concatenate(rids) - by_region = {str(k): int(v) for k, v in zip(r, y)} - return ProbeEmbeddings(embeddings=e, labels=y, region_ids=r, label_by_region=by_region) - - -def region_mean_auroc(prob: np.ndarray, region_ids: np.ndarray, label_by_region: Dict[str, int]): - """AUROC of the per-region mean probability (how the preprint scores the probe).""" - from sklearn.metrics import roc_auc_score - - regions = np.unique(region_ids) - means = np.array([float(prob[region_ids == r].mean()) for r in regions]) - y = np.array([label_by_region[str(r)] for r in regions]) - if len(np.unique(y)) < 2: - return float("nan") - return float(roc_auc_score(y, means)) - - -def fold_probe_into_linear_head( - coef: np.ndarray, - intercept: np.ndarray, - scaler_mean: np.ndarray, - scaler_scale: np.ndarray, - *, - num_classes: int, -) -> tuple[torch.Tensor, torch.Tensor]: - """Fold ``StandardScaler`` + logistic coefficients into ``Linear(dim, num_classes)``. - - Binary (``coef`` has one row): the class-1 row carries the decision function and - the class-0 row is zero, so ``softmax`` over the two logits equals the logistic - probability. Multinomial: rows map one-to-one onto classes. - """ - coef = np.asarray(coef, dtype=np.float64) - intercept = np.asarray(intercept, dtype=np.float64).reshape(-1) - w = coef / scaler_scale[None, :] - b = intercept - (coef * scaler_mean[None, :] / scaler_scale[None, :]).sum(axis=1) - dim = coef.shape[1] - W = torch.zeros(num_classes, dim, dtype=torch.float32) - B = torch.zeros(num_classes, dtype=torch.float32) - if coef.shape[0] == 1: - if num_classes < 2: - raise ValueError("A binary probe needs a 2-logit head (graph_head.output_dim=2).") - W[1] = torch.from_numpy(w[0].astype(np.float32)) - B[1] = float(b[0]) - else: - if coef.shape[0] != num_classes: - raise ValueError( - f"Probe has {coef.shape[0]} classes but the head has {num_classes} outputs." - ) - W[:] = torch.from_numpy(w.astype(np.float32)) - B[:] = torch.from_numpy(b.astype(np.float32)) - return W, B - - -def fit_probe( - train: ProbeEmbeddings, - *, - num_classes: int, - C: float = 1.0, - class_weight: Optional[str] = "balanced", - max_iter: int = 3000, - eval_sets: Optional[Dict[str, ProbeEmbeddings]] = None, -) -> ProbeFit: - """Fit the scaler + logistic regression and fold it into head weights.""" - from sklearn.linear_model import LogisticRegression - from sklearn.metrics import roc_auc_score - from sklearn.preprocessing import StandardScaler - - if len(np.unique(train.labels)) < 2: - raise ValueError("The probe training set has a single class.") - scaler = StandardScaler().fit(train.embeddings) - clf = LogisticRegression(max_iter=int(max_iter), C=float(C), class_weight=class_weight) - clf.fit(scaler.transform(train.embeddings), train.labels) - W, B = fold_probe_into_linear_head( - clf.coef_, clf.intercept_, scaler.mean_, scaler.scale_, num_classes=num_classes - ) - metrics: Dict[str, float] = {"n_train_instances": float(len(train.labels))} - for name, emb in (eval_sets or {}).items(): - prob = clf.predict_proba(scaler.transform(emb.embeddings)) - if prob.shape[1] == 2: - p1 = prob[:, 1] - if len(np.unique(emb.labels)) >= 2: - metrics[f"{name}_instance_auroc"] = float(roc_auc_score(emb.labels, p1)) - metrics[f"{name}_region_auroc"] = region_mean_auroc( - p1, emb.region_ids, emb.label_by_region - ) - metrics[f"n_{name}_instances"] = float(len(emb.labels)) - settings = {"C": float(C), "class_weight": class_weight, "max_iter": int(max_iter)} - return ProbeFit(weight=W, bias=B, metrics=metrics, settings=settings) - - -def save_probe_head( - fit: ProbeFit, path: str | Path, extra: Optional[Dict[str, Any]] = None -) -> Path: - """Write the folded head as ``graph_head.net.0.{weight,bias}`` (+ metrics).""" - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - payload: Dict[str, Any] = { - "graph_head.net.0.weight": fit.weight.clone(), - "graph_head.net.0.bias": fit.bias.clone(), - "metrics": dict(fit.metrics), - "settings": dict(fit.settings), - } - if extra: - payload.update(extra) - torch.save(payload, path) - return path - - -def load_probe_head(path: str | Path) -> Dict[str, Any]: - return torch.load(Path(path), map_location="cpu", weights_only=False) diff --git a/src/grass_mil/models/training/ssl_runtime.py b/src/grass_mil/models/training/ssl_runtime.py deleted file mode 100644 index f1bad8c..0000000 --- a/src/grass_mil/models/training/ssl_runtime.py +++ /dev/null @@ -1,22 +0,0 @@ -from __future__ import annotations - -import math -from dataclasses import dataclass - - -@dataclass -class CosineWarmup: - base_value: float - warmup_steps: int - total_steps: int - min_value: float = 0.0 - - def value(self, step: int) -> float: - if self.total_steps <= 0: - return self.base_value - step = max(0, min(step, self.total_steps)) - if self.warmup_steps > 0 and step < self.warmup_steps: - return self.base_value * float(step + 1) / float(self.warmup_steps) - progress = (step - self.warmup_steps) / max(1, self.total_steps - self.warmup_steps) - cosine = 0.5 * (1.0 + math.cos(math.pi * progress)) - return self.min_value + (self.base_value - self.min_value) * cosine diff --git a/src/grass_mil/nested_cv.py b/src/grass_mil/nested_cv.py index 4edd848..2c868ed 100644 --- a/src/grass_mil/nested_cv.py +++ b/src/grass_mil/nested_cv.py @@ -1,531 +1,174 @@ -"""``grass-mil-nested-cv``: nested, patient-grouped cross-validation (preprint protocol). - -Per outer fold ``i`` (``folds_file``): development = all other folds, split into -``inner_k`` stratified patient-grouped inner folds; every inner run trains the -selected regime and reports its best epoch; the refit trains on the whole -development set for the median best epoch and is scored once on fold ``i`` with -every cell of every region. One out-of-fold prediction per region and a mean -per patient are written alongside the per-fold summary. +"""``grass-mil-nested-cv``: nested, patient-grouped cross-validation. + +For every outer fold the remaining folds form the development set, which is +split into group-aware inner folds. Every inner run reports the epoch with the +best validation metric (``task.selection``); the model is refitted on the whole +development set for the reduced number of epochs, without validation, and +scored once on the held-out fold. The summary reports the metric of every +outer fold, its mean and standard deviation, and the pooled out-of-fold metric +at unit and patient level. """ from __future__ import annotations -import copy -import hashlib import json +import logging import statistics from pathlib import Path -from typing import Any, Dict, List, Optional, Sequence +from typing import Any, Sequence import hydra -import lightning as L import numpy as np import pandas as pd from omegaconf import DictConfig, open_dict -from grass_mil._paths import setup_project_root - -setup_project_root() - -from grass_mil.data.components.splits import ( # noqa: E402 - LabelTable, - load_fold_json, - load_labels, - make_inner_folds, -) -from grass_mil.train import train # noqa: E402 -from grass_mil.utils import RankedLogger, extras # noqa: E402 +from grass_mil.data.labels import class_index +from grass_mil.data.splits import inner_folds, load_folds +from grass_mil.inference.predictions import frame_metrics, patient_frame +from grass_mil.runtime import read_epoch_metrics, start_run +from grass_mil.train import train -log = RankedLogger(__name__, rank_zero_only=True) +__all__ = ["run_nested_cv", "best_epoch", "main"] -__all__ = ["run_nested_cv", "best_epoch_from_jsonl", "main"] +log = logging.getLogger(__name__) -# ---------------------------------------------------------------------------- -# helpers -# ---------------------------------------------------------------------------- +def best_epoch(rows: Sequence[dict[str, Any]], metric: str, mode: str) -> tuple[int, str]: + """Number of epochs (1-based) reaching the best ``metric``, earliest on ties. - -def best_epoch_from_jsonl( - path: str | Path, - *, - metric: str, - mode: str = "max", - fallback_metric: Optional[str] = "val/loss_region", -) -> Optional[tuple[int, str]]: - """1-based number of epochs reaching the best ``metric`` (earliest on ties). - - Returns ``(epochs, metric_used)``. When ``metric`` was never logged (e.g. the - validation fold holds a single class, so AUROC is undefined) the - ``fallback_metric`` is minimised instead, and failing that the last epoch is - returned with ``"last"``. + Falls back to the lowest ``val/loss`` when ``metric`` is never finite (for + example AUROC on a single-class validation fold), and to the last epoch + without any validation metric. Returns ``(epochs, metric_used)``. """ - path = Path(path) - if not path.exists(): - return None - rows = [json.loads(line) for line in path.read_text().splitlines() if line.strip()] - if not rows: - return None - usable = [r for r in rows if r.get(metric) is not None] - if usable: - sign = 1.0 if mode == "max" else -1.0 - best = max(usable, key=lambda r: (sign * float(r[metric]), -int(r["epoch"]))) - return int(best["epoch"]) + 1, metric - if fallback_metric: - usable = [r for r in rows if r.get(fallback_metric) is not None] + sign = 1.0 if mode == "max" else -1.0 + for name, s in ((metric, sign), ("val/loss", -1.0)): + usable = [r for r in rows if r.get(name) is not None and np.isfinite(r[name])] if usable: - best = min(usable, key=lambda r: (float(r[fallback_metric]), int(r["epoch"]))) - return int(best["epoch"]) + 1, fallback_metric - return int(max(int(r["epoch"]) for r in rows)) + 1, "last" + best = max(usable, key=lambda r: (s * float(r[name]), -int(r["epoch"]))) + return int(best["epoch"]) + 1, name + return (max(int(r["epoch"]) for r in rows) + 1 if rows else 1), "last" -def _reduce_epochs(epochs: Sequence[int], how: str, minimum: int) -> int: - if not epochs: - raise ValueError("No inner-fold epochs to reduce.") - how = str(how).lower() - if how == "median": - value = statistics.median(epochs) - elif how == "mean": - value = statistics.mean(epochs) - elif how == "max": - value = max(epochs) - else: +def _reduce(epochs: Sequence[int], how: str, minimum: int) -> int: + reducers = {"median": statistics.median, "mean": statistics.mean, "max": max} + if how not in reducers: raise ValueError("nested_cv.epoch_reduce must be median, mean or max.") - return max(int(minimum), int(value)) - - -def _resolve_label_column(cfg: DictConfig) -> str: - explicit = cfg.nested_cv.get("label_column") - if explicit: - return str(explicit) - targets = cfg.task.get("target_columns") - if targets: - return str(list(targets)[0]) - tasks = cfg.data.graph_labels.get("tasks") - if tasks: - return str(list(tasks)[0]) - raise ValueError( - "Set nested_cv.label_column (or task.target_columns / data.graph_labels.tasks) " - "so the driver knows which label column stratifies the inner folds." - ) - - -def _probe_mode(cfg: DictConfig) -> str: - """'always' | 'never' | 'auto' -> resolved to whether a probe is needed.""" - enabled = cfg.nested_cv.probe.get("enabled", "auto") - if isinstance(enabled, bool): - return "always" if enabled else "never" - enabled = str(enabled).lower() - if enabled not in {"auto", "true", "false"}: - raise ValueError("nested_cv.probe.enabled must be auto, true or false.") - if enabled == "auto": - return "auto" - return "always" if enabled == "true" else "never" - - -def _needs_probe(cfg: DictConfig) -> bool: - mode = _probe_mode(cfg) - if mode == "always": - return True - if mode == "never": - return False - return str(cfg.task.get("head_init_path") or "").strip().lower() == "auto" - - -def _split_hash(train_ids: Sequence[str], cfg: DictConfig) -> str: - h = hashlib.sha1() - h.update("|".join(sorted(str(r) for r in train_ids)).encode()) - h.update(str(cfg.model.get("init_from_ckpt")).encode()) - p = cfg.nested_cv.probe - h.update(f"{p.roots_per_unit}|{p.seed}|{p.C}|{p.class_weight}".encode()) - h.update(f"{cfg.data.sampler.runtime.depth}|{cfg.data.sampler.runtime.num_neighbors}".encode()) - return h.hexdigest()[:10] - - -def _id_kind(cfg: DictConfig) -> str: - kind = str(cfg.nested_cv.get("id_kind", "region")).lower() - if kind not in {"region", "sample", "patch"}: - raise ValueError("nested_cv.id_kind must be region, sample or patch") - return kind - - -def _with_split(cfg: DictConfig, *, train_ids, val_ids, test_ids, output_dir: Path) -> DictConfig: - fold_cfg = copy.deepcopy(cfg) - kind = _id_kind(cfg) - with open_dict(fold_cfg): - fold_cfg.data.split.explicit = { - "train": [str(r) for r in train_ids], - "val": [str(r) for r in val_ids], - "test": [str(r) for r in test_ids], - "id_kind": kind, - } - if kind == "patch": - fold_cfg.task.bag_key = "patch_id" - fold_cfg.data.graph_labels.scope = "patch" - elif kind == "sample": - fold_cfg.task.bag_key = "sample_id" - fold_cfg.data.graph_labels.scope = "sample" - fold_cfg.paths.output_dir = str(output_dir) - fold_cfg.trainer.default_root_dir = str(output_dir) - return fold_cfg - - -def _fit_probe_for(cfg: DictConfig, *, train_ids, val_ids, tag: str) -> Path: - from grass_mil.probe import run_probe - - cache_dir = Path(str(cfg.nested_cv.probe.cache_dir)) - cache_dir.mkdir(parents=True, exist_ok=True) - out = cache_dir / f"probe_{tag}_{_split_hash(train_ids, cfg)}.pt" - if out.exists() and bool(cfg.nested_cv.get("resume", True)): - log.info(f"[{tag}] probe cached: {out}") - return out - probe_cfg = _with_split( - cfg, train_ids=train_ids, val_ids=val_ids, test_ids=val_ids, output_dir=cache_dir / tag - ) - with open_dict(probe_cfg): - probe_cfg.task.head_init_path = None - probe_cfg.probe = { - "roots_per_unit": cfg.nested_cv.probe.roots_per_unit, - "seed": cfg.nested_cv.probe.seed, - "batch_size": cfg.nested_cv.probe.batch_size, - "C": cfg.nested_cv.probe.C, - "class_weight": cfg.nested_cv.probe.class_weight, - "max_iter": cfg.nested_cv.probe.max_iter, - "eval_splits": ["val"], - "output_path": str(out), - "device": cfg.nested_cv.probe.device, - } - run_probe(probe_cfg) - return out - - -def _train_run( - cfg: DictConfig, - *, - train_ids, - val_ids, - test_ids, - output_dir: Path, - head_init: Optional[Path], - max_epochs: Optional[int] = None, - refit: bool = False, - test: bool = False, -) -> Dict[str, Any]: - done_path = output_dir / "done.json" - if done_path.exists() and bool(cfg.nested_cv.get("resume", True)): - if not test or _predictions_cover(output_dir, test_ids, _id_kind(cfg)): - log.info(f"resume: {output_dir} already finished") - return json.loads(done_path.read_text()) - log.warning(f"resume: {output_dir} has stale test predictions; rerunning") - fold_cfg = _with_split( - cfg, train_ids=train_ids, val_ids=val_ids, test_ids=test_ids, output_dir=output_dir - ) - with open_dict(fold_cfg): - fold_cfg.task.head_init_path = str(head_init) if head_init is not None else None - fold_cfg.train = True - fold_cfg.test = bool(test) - fold_cfg.task_name = f"{cfg.task_name}_{output_dir.name}" + return max(int(minimum), int(round(reducers[how](epochs)))) + + +def _run(cfg: DictConfig, *, train_ids, val_ids, test_ids, out: Path, max_epochs: int | None = None) -> dict[str, Any]: + """One training run on explicit unit lists (resumable through ``done.json``).""" + done = out / "done.json" + if done.exists() and cfg.nested_cv.resume: + return json.loads(done.read_text()) + run_cfg = cfg.copy() + with open_dict(run_cfg): + run_cfg.data.split.train = list(train_ids) + run_cfg.data.split.val = list(val_ids) + run_cfg.data.split.test = list(test_ids) + run_cfg.paths.output_dir = str(out) + run_cfg.trainer.default_root_dir = str(out) + run_cfg.test = bool(test_ids) if max_epochs is not None: - fold_cfg.trainer.max_epochs = int(max_epochs) - fold_cfg.trainer.min_epochs = int(max_epochs) - if refit and fold_cfg.get("callbacks"): - # Fixed-length refit: no early stopping, no metric-based selection; - # the last epoch's weights are the model. - cbs = fold_cfg.callbacks - if "early_stopping" in cbs: - del cbs["early_stopping"] - if "model_checkpoint" in cbs and cbs.model_checkpoint is not None: - cbs.model_checkpoint.monitor = None - cbs.model_checkpoint.save_top_k = 0 - cbs.model_checkpoint.save_last = True - cbs.model_checkpoint.filename = "last_epoch" - metrics, _objects = train(fold_cfg) - numeric = {} - for k, v in metrics.items(): - try: - numeric[k] = float(v.item() if hasattr(v, "item") else v) - except Exception: - continue - result = { - "output_dir": str(output_dir), - "metrics": numeric, - "n_train": len(train_ids), - "n_val": len(val_ids), - "n_test": len(test_ids), - "max_epochs": int(fold_cfg.trainer.max_epochs), - "head_init_path": str(head_init) if head_init is not None else None, - } - best = best_epoch_from_jsonl( - output_dir / "metrics_per_epoch.jsonl", - metric=str(cfg.nested_cv.selection_metric), - mode=str(cfg.nested_cv.selection_mode), - ) - result["best_epoch"] = best[0] if best else None - result["selection_metric_used"] = best[1] if best else None - done_path.write_text(json.dumps(result, indent=1)) + run_cfg.trainer.max_epochs = run_cfg.trainer.min_epochs = int(max_epochs) + out.mkdir(parents=True, exist_ok=True) + metrics, _ = train(run_cfg) + epochs, used = best_epoch(read_epoch_metrics(out / "metrics_per_epoch.jsonl"), cfg.task.selection.metric, cfg.task.selection.mode) + result = {"metrics": metrics, "best_epoch": epochs, "selection_metric_used": used, "output_dir": str(out)} + done.write_text(json.dumps(result, indent=1)) return result -def _unit_ids(df: pd.DataFrame, kind: str) -> pd.Series: - """Unit id of every prediction row. - - Region bags carry their region id in ``region_id``; patch and sample bags - are identified by the bag id, which the bagging code namespaces as - ``sample::region::patch`` (or ``sample::patch``) for patches, so the unit is - its last segment. - """ - if kind == "region": - return df["region_id"].astype(str) - return df["bag_id"].astype(str).str.split("::").str[-1] - - -def _predictions_cover(output_dir: Path, test_ids, kind: str) -> bool: - """True when the run's test predictions hold exactly the expected unit ids.""" - path = output_dir / "test_predictions.csv" - if not path.exists(): - return False - df = pd.read_csv(path) - if ("region_id" if kind == "region" else "bag_id") not in df.columns: - return False - return set(_unit_ids(df, kind)) == {str(r) for r in test_ids} - - -def _read_test_predictions(output_dir: Path) -> pd.DataFrame: - path = output_dir / "test_predictions.csv" - if not path.exists(): - raise FileNotFoundError(f"Refit run wrote no test predictions: {path}") - return pd.read_csv(path) - - -def _auroc(y: Sequence[int], score: Sequence[float]) -> float: - from sklearn.metrics import roc_auc_score - - y = np.asarray(y) - if len(np.unique(y)) < 2: - return float("nan") - return float(roc_auc_score(y, np.asarray(score, dtype=float))) - - -# ---------------------------------------------------------------------------- -# driver -# ---------------------------------------------------------------------------- - - -def run_nested_cv(cfg: DictConfig) -> Dict[str, Any]: - if cfg.get("seed"): - L.seed_everything(cfg.seed, workers=True) +def run_nested_cv(cfg: DictConfig) -> dict[str, Any]: + start_run(cfg) ncfg = cfg.nested_cv - base_out = Path(str(cfg.paths.output_dir)) - base_out.mkdir(parents=True, exist_ok=True) + kind = str(cfg.task.target_type) + out_dir = Path(str(cfg.paths.output_dir)) + out_dir.mkdir(parents=True, exist_ok=True) - label_column = _resolve_label_column(cfg) - patient_column = ncfg.get("patient_column") - table: LabelTable = load_labels( - str(ncfg.label_file), - id_column=str(ncfg.id_column), - label_column=label_column, - patient_column=str(patient_column) if patient_column else None, - ) - folds = load_fold_json(str(ncfg.folds_file)) - labelled = set(table.labels.index) - dropped = sum(len([r for r in f if r not in labelled]) for f in folds) - if dropped: - log.warning(f"{dropped} region(s) in the folds have no '{label_column}' label; dropped") - folds = [[r for r in f if r in labelled] for f in folds] - if any(len(f) == 0 for f in folds): - raise ValueError("An outer fold is empty after dropping unlabelled regions.") - - # Precompute once (the inner runs re-partition the same processed graphs). - datamodule = hydra.utils.instantiate(cfg.data) + datamodule = hydra.utils.instantiate(cfg.data, _recursive_=False) datamodule.prepare_data() - - needs_probe = _needs_probe(cfg) - if needs_probe and int(cfg.model.graph_head.get("num_layers", 2)) != 1: - raise ValueError( - "A probe pre-fit targets a linear instance head; set model.graph_head.num_layers=1 " - "(the linear_* regimes do) or nested_cv.probe.enabled=false." - ) - if not needs_probe and str(cfg.task.get("head_init_path") or "").lower() == "auto": - raise ValueError( - "task.head_init_path='auto' needs a probe but nested_cv.probe.enabled=false; " - "pass an explicit task.head_init_path or enable the probe." - ) - fixed_head = None - if not needs_probe and cfg.task.get("head_init_path"): - fixed_head = Path(str(cfg.task.head_init_path)) - log.warning(f"Using one fixed head init for every fold: {fixed_head}") - - selected = ncfg.get("outer_folds") - outer_indices = list(range(len(folds))) if not selected else [int(i) for i in selected] - inner_k = int(ncfg.inner_k) - regime = str(cfg.get("regime_name", "custom")) - log.info( - f"Nested CV [{regime}]: {len(folds)} outer folds ({len(outer_indices)} selected), " - f"inner k={inner_k} seed={int(ncfg.inner_seed)}, probe={'yes' if needs_probe else 'no'}" - ) - - fold_results: List[Dict[str, Any]] = [] - oof_frames: List[pd.DataFrame] = [] - for i in outer_indices: - outer_dir = base_out / f"outer_{i}" - outer_dir.mkdir(parents=True, exist_ok=True) - result_path = outer_dir / "result.json" - if result_path.exists() and bool(ncfg.get("resume", True)): - res = json.loads(result_path.read_text()) - log.info(f"[outer {i}] resume: AUROC {res['test_auroc']:.4f} (E={res['E']})") - fold_results.append(res) - oof_frames.append(pd.read_csv(outer_dir / "oof_predictions.csv")) + units = datamodule.all_units() + if datamodule.read_labels() is None: + raise ValueError("Cross-validation needs labels (data.labels.file).") + meta = datamodule.unit_table(units) + labelled = {e.unit_id: units.label_id(e) for e in units.entries if units.label_id(e) is not None} + groups = {u: (meta.loc[u, "patient_id"] if "patient_id" in meta.columns else meta.loc[u, "region_id"]) for u in labelled} + classes = {u: class_index(units.labels.target(i)[0]) for u, i in labelled.items()} if kind == "classification" else None + + def _units_of(ids: Sequence[str]) -> list[str]: + wanted = set(map(str, ids)) + return sorted(u for u in labelled if {u, meta.loc[u, "region_id"], meta.loc[u, "sample_id"]} & wanted) + + folds = [_units_of(f) for f in load_folds(str(ncfg.folds_file))] + if any(not f for f in folds): + raise ValueError("An outer fold holds no labelled unit.") + selected = list(ncfg.outer_folds) if ncfg.outer_folds is not None else range(len(folds)) + metric = str(cfg.task.selection.metric).split("/", 1)[-1] + + results, oof = [], [] + for i in selected: + fold_dir = out_dir / f"outer_{i}" + result_path = fold_dir / "result.json" + if result_path.exists() and ncfg.resume: + results.append(json.loads(result_path.read_text())) + oof.append(pd.read_csv(fold_dir / "oof_predictions.csv")) continue - - test_ids = list(folds[i]) - dev_ids = sorted(r for j, f in enumerate(folds) if j != i for r in f) - inner = make_inner_folds(dev_ids, table, k=inner_k, seed=int(ncfg.inner_seed)) - - inner_epochs: List[int] = [] - inner_summaries: List[Dict[str, Any]] = [] + test = folds[i] + dev = sorted(u for j, f in enumerate(folds) if j != i for u in f) + inner = inner_folds(dev, groups=groups, classes=classes, k=int(ncfg.inner_k), seed=int(ncfg.inner_seed)) + epochs = [] for j, (tr, va) in enumerate(inner): - tag = f"o{i}_i{j}" - head = fixed_head - if needs_probe: - head = _fit_probe_for(cfg, train_ids=tr, val_ids=va, tag=tag) - res = _train_run( - cfg, - train_ids=tr, - val_ids=va, - test_ids=test_ids, - output_dir=outer_dir / f"inner_{j}", - head_init=head, - test=False, - ) - if res.get("best_epoch") is None: - raise RuntimeError( - f"[{tag}] no '{ncfg.selection_metric}' in metrics_per_epoch.jsonl; " - "is validation enabled with region-level metrics?" - ) - used = res.get("selection_metric_used") - if used != str(ncfg.selection_metric): - log.warning( - f"[{tag}] '{ncfg.selection_metric}' was never logged (single-class " - f"validation fold?); epoch chosen by '{used}'" - ) - inner_epochs.append(int(res["best_epoch"])) - inner_summaries.append( - { - "inner": j, - "best_epoch": res["best_epoch"], - "selection_metric_used": used, - "n_train": len(tr), - "n_val": len(va), - } - ) - log.info(f"[{tag}] best epoch {res['best_epoch']} ({used})") - - E = _reduce_epochs(inner_epochs, str(ncfg.epoch_reduce), int(ncfg.min_refit_epochs)) - tag = f"o{i}_refit" - head = fixed_head - if needs_probe: - head = _fit_probe_for(cfg, train_ids=dev_ids, val_ids=test_ids, tag=tag) - res = _train_run( - cfg, - train_ids=dev_ids, - val_ids=test_ids, - test_ids=test_ids, - output_dir=outer_dir / "refit", - head_init=head, - max_epochs=E, - refit=True, - test=True, - ) - preds = _read_test_predictions(outer_dir / "refit") + res = _run(cfg, train_ids=tr, val_ids=va, test_ids=[], out=fold_dir / f"inner_{j}") + if res["selection_metric_used"] != cfg.task.selection.metric: + log.warning(f"[outer {i} inner {j}] epoch chosen by {res['selection_metric_used']}") + epochs.append(int(res["best_epoch"])) + n_epochs = _reduce(epochs, str(ncfg.epoch_reduce), int(ncfg.min_refit_epochs)) + refit = _run(cfg, train_ids=dev, val_ids=[], test_ids=test, out=fold_dir / "refit", max_epochs=n_epochs) + preds = pd.read_csv(fold_dir / "refit" / "test_predictions.csv") preds.insert(0, "outer", i) - kind = _id_kind(cfg) - if kind != "region": - # keep the unit id (patch or sample) in region_id so every - # downstream table has one id column - preds["region_id"] = _unit_ids(preds, kind) - preds["patient_id"] = [table.group_of(r) for r in preds["region_id"].astype(str)] - preds.to_csv(outer_dir / "oof_predictions.csv", index=False) - score_col = "prob_class_1" if "prob_class_1" in preds.columns else None - test_auroc = res["metrics"].get("test/auroc") - if test_auroc is None and score_col is not None: - test_auroc = _auroc(preds["label"], preds[score_col]) - fold_res = { + preds.to_csv(fold_dir / "oof_predictions.csv", index=False) + fold = { "outer": i, - "E": E, - "inner_epochs": inner_epochs, - "inner": inner_summaries, - "test_auroc": float(test_auroc) if test_auroc is not None else float("nan"), - "test_bal_acc": res["metrics"].get("test/bal_acc"), - "n_test": len(test_ids), - "n_test_evaluated": int(preds["region_id"].nunique()), - "n_dev": len(dev_ids), - "run_dir": str(outer_dir / "refit"), - "head_init_path": res.get("head_init_path"), + "epochs": n_epochs, + "inner_best_epochs": epochs, + "n_dev": len(dev), + "n_test": len(test), + "metrics": frame_metrics(preds, kind), + "run_dir": refit["output_dir"], } - if fold_res["n_test_evaluated"] != len(test_ids): - log.warning( - f"[outer {i}] evaluated {fold_res['n_test_evaluated']} of {len(test_ids)} regions" - ) - result_path.write_text(json.dumps(fold_res, indent=1)) - fold_results.append(fold_res) - oof_frames.append(preds) - log.info(f"[outer {i}] E={E} held-out AUROC {fold_res['test_auroc']:.4f}") - - # ------------------------------------------------------------------ summary - oof = pd.concat(oof_frames, ignore_index=True) if oof_frames else pd.DataFrame() - oof.to_csv(base_out / "oof_predictions.csv", index=False) - aucs = [r["test_auroc"] for r in fold_results if r["test_auroc"] == r["test_auroc"]] - summary: Dict[str, Any] = { - "regime": regime, - "task": label_column, - "n_outer": len(fold_results), - "n_outer_total": len(folds), - "complete": len(fold_results) == len(folds), - "inner_k": inner_k, - "inner_seed": int(ncfg.inner_seed), - "epoch_reduce": str(ncfg.epoch_reduce), - "selection_metric": str(ncfg.selection_metric), - "max_epochs": int(cfg.trainer.max_epochs), - "probe": needs_probe, - "encoder_ckpt": cfg.model.get("init_from_ckpt"), - "auc_mean": float(np.mean(aucs)) if aucs else float("nan"), - "auc_sd": float(np.std(aucs, ddof=1)) if len(aucs) > 1 else 0.0, - "auc_folds": [round(a, 4) for a in aucs], - "folds": fold_results, + result_path.write_text(json.dumps(fold, indent=1)) + results.append(fold) + oof.append(preds) + log.info(f"[outer {i}] {n_epochs} epochs, held-out {metric} {fold['metrics'].get(metric, float('nan')):.4f}") + + oof_frame = pd.concat(oof, ignore_index=True) + oof_frame.to_csv(out_dir / "oof_predictions.csv", index=False) + values = [r["metrics"].get(metric, float("nan")) for r in results] + finite = [v for v in values if np.isfinite(v)] + summary: dict[str, Any] = { + "regime": cfg.get("regime_name"), + "target_type": kind, + "metric": metric, + "fold_values": values, + "mean": float(np.mean(finite)) if finite else float("nan"), + "sd": float(np.std(finite, ddof=1)) if len(finite) > 1 else float("nan"), + "pooled": frame_metrics(oof_frame, kind), + "complete": len(results) == len(folds), + "folds": results, } - if not oof.empty and "prob_class_1" in oof.columns: - summary["pooled_region_auroc"] = _auroc(oof["label"], oof["prob_class_1"]) - agg = str(ncfg.get("patient_aggregation", "mean")).lower() - if agg != "mean": - raise ValueError( - "nested_cv.patient_aggregation supports 'mean' here; use grass-mil-predict for merged bags." - ) - pat = ( - oof.groupby("patient_id") - .agg( - prob_class_1=("prob_class_1", "mean"), - label=("label", "max"), - n_regions=("region_id", "nunique"), - outer=("outer", "first"), - ) - .reset_index() - ) - pat.to_csv(base_out / "oof_patient_predictions.csv", index=False) - summary["patient_auroc"] = _auroc(pat["label"], pat["prob_class_1"]) - summary["n_patients"] = int(len(pat)) - (base_out / str(ncfg.summary_filename)).write_text(json.dumps(summary, indent=1, default=str)) - log.info( - f"== [{regime}] AUROC {summary['auc_mean']:.4f} +/- {summary['auc_sd']:.4f} over " - f"{len(aucs)} outer folds; patient-level {summary.get('patient_auroc', float('nan')):.4f}" - ) + if "patient_id" in oof_frame.columns: + patients = patient_frame(oof_frame, kind) + patients.to_csv(out_dir / "oof_patient_predictions.csv", index=False) + summary["patient"] = frame_metrics(patients, kind) + (out_dir / "nested_cv_summary.json").write_text(json.dumps(summary, indent=1, default=float)) + log.info(f"{metric}: {summary['mean']:.4f} ± {summary['sd']:.4f} over {len(finite)} outer folds") return summary @hydra.main(version_base="1.3", config_path="configs", config_name="nested_cv.yaml") -def main(cfg: DictConfig) -> Optional[float]: - extras(cfg) - summary = run_nested_cv(cfg) - return summary.get("auc_mean") +def main(cfg: DictConfig) -> None: + run_nested_cv(cfg) if __name__ == "__main__": diff --git a/src/grass_mil/predict.py b/src/grass_mil/predict.py new file mode 100644 index 0000000..bb377f7 --- /dev/null +++ b/src/grass_mil/predict.py @@ -0,0 +1,110 @@ +"""``grass-mil-predict``: unit (and patient) predictions of a trained model on a cohort.""" + +from __future__ import annotations + +import json +import logging +from pathlib import Path + +import hydra +import pandas as pd +import torch +from omegaconf import DictConfig + +from grass_mil.data.datamodule import GraphUnitDataset +from grass_mil.data.samplers import EgoGraphLoader +from grass_mil.data.splits import read_ids +from grass_mil.inference.cell_cache import sampler_of_checkpoint +from grass_mil.inference.predictions import frame_metrics, patient_frame, prediction_frame +from grass_mil.models import SupervisedModule, attention_pool +from grass_mil.runtime import resolve_device, start_run + +__all__ = ["run_predict", "main"] + +log = logging.getLogger(__name__) + + +@torch.no_grad() +def run_predict(cfg: DictConfig) -> dict: + start_run(cfg) + pcfg = cfg.predict + if pcfg.patient_aggregation not in ("mean", "merged_bag"): + raise ValueError("predict.patient_aggregation must be 'mean' or 'merged_bag'.") + out_dir = Path(str(cfg.paths.output_dir)) + out_dir.mkdir(parents=True, exist_ok=True) + device = resolve_device(pcfg.device) + ckpt = torch.load(str(cfg.ckpt_path), map_location="cpu", weights_only=False) + model = SupervisedModule.from_checkpoint(str(cfg.ckpt_path)).to(device) + kind = model.task["target_type"] + datamodule = hydra.utils.instantiate(cfg.data, _recursive_=False) + datamodule.prepare_data() + if tuple(datamodule.input_spec.cell_types) != tuple(model.input_spec.cell_types): + raise ValueError("The cohort's cell types differ from the ones the model was trained on.") + units = datamodule.all_units() + wanted = read_ids(pcfg.get("units")) + if wanted is not None: + wanted = set(wanted) + units = GraphUnitDataset( + [e for e in units.entries if {e.unit_id, e.region_id, e.sample_id} & wanted], units.labels + ) + sampler = sampler_of_checkpoint(ckpt, int(pcfg.batch_size)) + meta = datamodule.unit_table(units) + + unit_ids, bag_logits, targets, instances = [], [], [], {} + loader = EgoGraphLoader(units, sampler, mode="chunks", roots_per_unit=None) + current: dict[str, list] = {} + for batch in loader: + unit = batch.unit_id[0] + _, logits, scores = model.forward_instances(batch.to(device)) + current.setdefault(unit, [[], [], getattr(batch, "graph_y", None)]) + current[unit][0].append(logits.float().cpu()) + current[unit][1].append(scores.float().cpu()) + for unit, (logits, scores, y) in current.items(): + logits, scores = torch.cat(logits), torch.cat(scores) + unit_ids.append(unit) + bag_logits.append(attention_pool(logits, scores)[0]) + targets.append(None if y is None else y[0].cpu()) + instances[unit] = (logits, scores) + labelled = all(t is not None for t in targets) + frame = prediction_frame( + unit_ids, + torch.stack(bag_logits), + target_type=kind, + targets=torch.stack(targets) if labelled else None, + meta=meta, + ) + frame.to_csv(out_dir / "predictions.csv", index=False) + summary: dict = {"n_units": len(frame), "ckpt_path": str(cfg.ckpt_path)} + if labelled: + summary["unit_metrics"] = frame_metrics(frame, kind) + + if "patient_id" in frame.columns: + if pcfg.patient_aggregation == "mean": + patients = patient_frame(frame, kind) + else: + rows = [] + for patient, group in frame.groupby("patient_id"): + logits = torch.cat([instances[u][0] for u in group["unit_id"]]) + scores = torch.cat([instances[u][1] for u in group["unit_id"]]) + rows.append((patient, attention_pool(logits, scores)[0], group.index[0])) + patients = prediction_frame( + [r[0] for r in rows], + torch.stack([r[1] for r in rows]), + target_type=kind, + targets=torch.stack([targets[i] for _, _, i in rows]) if labelled else None, + ).rename(columns={"unit_id": "patient_id"}) + patients.to_csv(out_dir / "patient_predictions.csv", index=False) + if labelled: + summary["patient_metrics"] = frame_metrics(patients, kind) + (out_dir / "metrics.json").write_text(json.dumps(summary, indent=1, default=float)) + log.info(f"Predictions for {len(frame)} units written to {out_dir}") + return summary + + +@hydra.main(version_base="1.3", config_path="configs", config_name="predict.yaml") +def main(cfg: DictConfig) -> None: + run_predict(cfg) + + +if __name__ == "__main__": + main() diff --git a/src/grass_mil/probe.py b/src/grass_mil/probe.py deleted file mode 100644 index f257ff1..0000000 --- a/src/grass_mil/probe.py +++ /dev/null @@ -1,133 +0,0 @@ -"""``grass-mil-probe``: fit the linear probe head on a frozen encoder.""" - -from __future__ import annotations - -import json -from pathlib import Path -from typing import Any, Dict, Optional - -import hydra -import lightning as L -import torch -from omegaconf import DictConfig - -from grass_mil._paths import setup_project_root - -setup_project_root() - -from grass_mil.models.training.probe import ( # noqa: E402 - attach_datamodule, - collect_frozen_embeddings, - fit_probe, - probe_sampler_config, - save_probe_head, -) -from grass_mil.utils import RankedLogger, extras # noqa: E402 - -log = RankedLogger(__name__, rank_zero_only=True) - - -def _resolve_device(name: str) -> torch.device: - name = str(name or "auto").lower() - if name == "auto": - return torch.device("cuda" if torch.cuda.is_available() else "cpu") - return torch.device(name) - - -def run_probe(cfg: DictConfig) -> Dict[str, Any]: - if cfg.get("seed"): - L.seed_everything(cfg.seed, workers=True) - pcfg = cfg.probe - - datamodule = hydra.utils.instantiate(cfg.data) - datamodule.prepare_data() - datamodule.setup("fit") - - model = hydra.utils.instantiate(cfg.model) - trainer = hydra.utils.instantiate(cfg.trainer, logger=False, enable_checkpointing=False) - attach_datamodule(model, datamodule, trainer) - model.setup("fit") - device = _resolve_device(pcfg.get("device", "auto")) - model.to(device) - model.eval() - for p in model.parameters(): - p.requires_grad_(False) - - sampler_cfg = probe_sampler_config( - datamodule.sampler_config, - roots_per_unit=pcfg.get("roots_per_unit"), - seed=int(pcfg.get("seed", 0)), - batch_size=int(pcfg.get("batch_size", 256)), - ) - target_columns = cfg.task.get("target_columns") - target_columns = list(target_columns) if target_columns else None - num_classes = int(cfg.model.graph_head.output_dim) - - log.info( - f"Collecting frozen embeddings: {pcfg.get('roots_per_unit')} roots/region, " - f"seed {pcfg.get('seed', 0)}, {len(datamodule.dataset_train)} training regions" - ) - train_emb = collect_frozen_embeddings( - model, - datamodule.dataset_train, - sampler_config=sampler_cfg, - target_columns=target_columns, - device=device, - num_workers=int(cfg.data.get("num_workers", 0)), - ) - eval_sets = {} - for split in list(pcfg.get("eval_splits") or []): - ds = getattr(datamodule, f"dataset_{split}", None) - if ds is None or len(ds) == 0: - log.warning(f"probe.eval_splits: split '{split}' is empty; skipped") - continue - eval_sets[str(split)] = collect_frozen_embeddings( - model, - ds, - sampler_config=sampler_cfg, - target_columns=target_columns, - device=device, - num_workers=int(cfg.data.get("num_workers", 0)), - ) - - class_weight = pcfg.get("class_weight") - class_weight = None if class_weight in (None, "null", "none", "None") else str(class_weight) - fit = fit_probe( - train_emb, - num_classes=num_classes, - C=float(pcfg.get("C", 1.0)), - class_weight=class_weight, - max_iter=int(pcfg.get("max_iter", 3000)), - eval_sets=eval_sets, - ) - fit.settings.update( - { - "roots_per_unit": pcfg.get("roots_per_unit"), - "seed": int(pcfg.get("seed", 0)), - "sampler_runtime": dict(sampler_cfg.runtime), - "init_from_ckpt": cfg.model.get("init_from_ckpt"), - "n_train_regions": len(datamodule.dataset_train), - } - ) - out = save_probe_head(fit, str(pcfg.output_path)) - summary = {"output_path": str(out), "metrics": fit.metrics, "settings": fit.settings} - summary_path = Path(str(out)).with_suffix(".json") - summary_path.write_text(json.dumps(summary, indent=1, default=str)) - log.info( - f"Probe head written to {out}: " - + ", ".join( - f"{k}={v:.4f}" if isinstance(v, float) else f"{k}={v}" for k, v in fit.metrics.items() - ) - ) - return summary - - -@hydra.main(version_base="1.3", config_path="configs", config_name="probe.yaml") -def main(cfg: DictConfig) -> Optional[float]: - extras(cfg) - summary = run_probe(cfg) - return summary["metrics"].get("val_region_auroc") - - -if __name__ == "__main__": - main() diff --git a/src/grass_mil/runtime.py b/src/grass_mil/runtime.py new file mode 100644 index 0000000..5979376 --- /dev/null +++ b/src/grass_mil/runtime.py @@ -0,0 +1,82 @@ +"""Shared plumbing of the command-line entry points.""" + +from __future__ import annotations + +import json +import logging +import math +from pathlib import Path +from typing import Any + +import hydra +import lightning as L +import torch +from omegaconf import DictConfig, OmegaConf + +__all__ = [ + "start_run", + "instantiate_group", + "resolve_device", + "EpochMetricsWriter", + "read_epoch_metrics", +] + +log = logging.getLogger("grass_mil") + + +def start_run(cfg: DictConfig) -> None: + """Seed every RNG and optionally print the composed config.""" + if cfg.get("seed") is not None: + L.seed_everything(int(cfg.seed), workers=True, verbose=False) + if cfg.get("print_config"): + log.info("\n" + OmegaConf.to_yaml(cfg, resolve=True)) + + +def instantiate_group(group: DictConfig | None) -> list[Any]: + """Instantiate every entry of a ``{name: {_target_: ...}}`` group (callbacks, loggers).""" + if not group: + return [] + return [hydra.utils.instantiate(conf) for conf in group.values() if conf and "_target_" in conf] + + +def resolve_device(name: str | None) -> torch.device: + """``auto`` (CUDA when available), ``cpu``, ``cuda`` or ``cuda:``.""" + name = str(name or "auto").lower() + if name == "auto": + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + return torch.device(name) + + +class EpochMetricsWriter(L.Callback): + """Append the logged ``train/*`` and ``val/*`` metrics of every epoch to a JSON-lines file. + + Cross-validation reads it to find the best epoch of every inner run. + """ + + def __init__(self, path: str | Path) -> None: + self.path = Path(path) + + def on_train_start(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None: + if trainer.is_global_zero: + self.path.parent.mkdir(parents=True, exist_ok=True) + self.path.unlink(missing_ok=True) + + def on_validation_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None: + # Runs after the module logged its epoch-level validation metrics. + if trainer.sanity_checking or not trainer.is_global_zero: + return + row: dict[str, Any] = {"epoch": int(trainer.current_epoch), "step": int(trainer.global_step)} + for key, value in trainer.callback_metrics.items(): + if key.startswith(("train/", "val/")) and not key.endswith("_step"): + value = float(value) + if not math.isnan(value): + row[key] = value + with self.path.open("a") as fh: + fh.write(json.dumps(row) + "\n") + + +def read_epoch_metrics(path: str | Path) -> list[dict[str, Any]]: + path = Path(path) + if not path.exists(): + return [] + return [json.loads(line) for line in path.read_text().splitlines() if line.strip()] diff --git a/src/grass_mil/train.py b/src/grass_mil/train.py index 502b1aa..2be226d 100644 --- a/src/grass_mil/train.py +++ b/src/grass_mil/train.py @@ -1,170 +1,83 @@ -from typing import Any, Dict, List, Optional, Tuple +"""``grass-mil-train``: one training run (fine-tuning, or ``regime=pretrain``).""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import Any import hydra import lightning as L -from lightning import Callback, LightningDataModule, LightningModule, Trainer -from lightning.pytorch.loggers import Logger -from omegaconf import DictConfig, OmegaConf - -from grass_mil._paths import setup_project_root - -setup_project_root() - -from grass_mil.utils import ( - RankedLogger, - extras, - get_metric_value, - instantiate_callbacks, - instantiate_loggers, - log_hyperparameters, - task_wrapper, -) # noqa: E402 - -log = RankedLogger(__name__, rank_zero_only=True) - - -def _requires_zero_validation(cfg: DictConfig) -> bool: - split_cfg = cfg.get("data", {}).get("split") - if not split_cfg: - return False - - loocv_cfg = split_cfg.get("loocv") - if loocv_cfg and bool(loocv_cfg.get("enabled", False)): - return float(loocv_cfg.get("val_ratio", 0.0)) == 0.0 - - ratios = split_cfg.get("train_val_test_split") - if ratios is None: - return False - try: - return float(ratios[1]) == 0.0 - except Exception: - return False - - -def _filter_val_monitor_callbacks( - callbacks_cfg: Optional[DictConfig], -) -> Tuple[Optional[DictConfig], List[str]]: - if not callbacks_cfg or not isinstance(callbacks_cfg, DictConfig): - return callbacks_cfg, [] - - # Resolve interpolations (``${paths.output_dir}`` and friends) first: the - # filtered copy below is a new root, so unresolved references into the - # full config would fail at instantiation. - try: - callbacks_cfg = OmegaConf.create(OmegaConf.to_container(callbacks_cfg, resolve=True)) - except Exception: # pragma: no cover - unresolvable outside a Hydra run - pass - filtered = OmegaConf.create({}) - removed: List[str] = [] - for name, cb_conf in callbacks_cfg.items(): - monitor = cb_conf.get("monitor") if isinstance(cb_conf, DictConfig) else None - if isinstance(monitor, str) and monitor.startswith("val/"): - removed.append(str(name)) - continue - filtered[name] = cb_conf - return filtered, removed - - -@task_wrapper -def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: - """Train (and optionally test) using Hydra-composed config.""" - if cfg.get("seed"): - L.seed_everything(cfg.seed, workers=True) - - log.info(f"Instantiating datamodule <{cfg.data._target_}>") - datamodule: LightningDataModule = hydra.utils.instantiate(cfg.data) - - if cfg.get("model") is None: - raise ValueError("No model config provided. Set `model=...` in your Hydra config.") - log.info(f"Instantiating model <{cfg.model._target_}>") - model: LightningModule = hydra.utils.instantiate(cfg.model) - - log.info("Instantiating callbacks...") - callbacks_cfg = cfg.get("callbacks") - if _requires_zero_validation(cfg): - callbacks_cfg, removed = _filter_val_monitor_callbacks(callbacks_cfg) - if removed: - log.info( - "Validation is disabled by configuration; skipping val-monitored callbacks: " - f"{removed}" - ) - callbacks: List[Callback] = instantiate_callbacks(callbacks_cfg) - output_dir = cfg.get("paths", {}).get("output_dir") if cfg.get("paths") else None - if output_dir and not _requires_zero_validation(cfg): - from grass_mil.callbacks import EpochMetricsWriter - - callbacks.append(EpochMetricsWriter(str(output_dir))) - - log.info("Instantiating loggers...") - logger: List[Logger] = instantiate_loggers(cfg.get("logger")) - - log.info(f"Instantiating trainer <{cfg.trainer._target_}>") - trainer: Trainer = hydra.utils.instantiate(cfg.trainer, callbacks=callbacks, logger=logger) - - object_dict = { - "cfg": cfg, - "datamodule": datamodule, - "model": model, - "callbacks": callbacks, - "logger": logger, - "trainer": trainer, - } - - if logger: - log.info("Logging hyperparameters!") - log_hyperparameters(object_dict) +from omegaconf import DictConfig, OmegaConf, open_dict + +from grass_mil.inference.predictions import prediction_frame +from grass_mil.runtime import EpochMetricsWriter, instantiate_group, start_run + +__all__ = ["train", "main"] + +log = logging.getLogger(__name__) + + +def _without_validation(cfg: DictConfig) -> DictConfig: + """Drop the callbacks that monitor validation metrics and disable the loop.""" + cfg = cfg.copy() + with open_dict(cfg): + for name in list(cfg.callbacks or {}): + monitor = (cfg.callbacks[name] or {}).get("monitor") + if isinstance(monitor, str) and monitor.startswith("val/"): + if name == "model_checkpoint": + cfg.callbacks[name].monitor = None + cfg.callbacks[name].save_top_k = 0 + else: + del cfg.callbacks[name] + cfg.trainer.limit_val_batches = 0 + return cfg + + +def train(cfg: DictConfig) -> tuple[dict[str, float], dict[str, Any]]: + """Train (and test) the configured model; returns the final metrics and the run objects.""" + start_run(cfg) + output_dir = Path(str(cfg.paths.output_dir)) + datamodule: L.LightningDataModule = hydra.utils.instantiate(cfg.data, _recursive_=False) + datamodule.prepare_data() + datamodule.setup("fit") + if len(datamodule.dataset_val) == 0 and cfg.trainer.get("limit_val_batches") != 0: + log.info("The validation split is empty: training without validation.") + cfg = _without_validation(cfg) + + model: L.LightningModule = hydra.utils.instantiate(cfg.model, _recursive_=False) + callbacks = instantiate_group(cfg.get("callbacks")) + if cfg.trainer.get("limit_val_batches") != 0: + callbacks.append(EpochMetricsWriter(output_dir / "metrics_per_epoch.jsonl")) + trainer: L.Trainer = hydra.utils.instantiate( + cfg.trainer, callbacks=callbacks, logger=instantiate_group(cfg.get("logger")) or False + ) if cfg.get("train"): - log.info("Starting training!") - trainer.fit( - model=model, - datamodule=datamodule, - ckpt_path=cfg.get("ckpt_path"), + trainer.fit(model, datamodule=datamodule, ckpt_path=cfg.get("ckpt_path")) + metrics = {k: float(v) for k, v in trainer.callback_metrics.items()} + + if cfg.get("test") and len(datamodule.dataset_test): + checkpoint = trainer.checkpoint_callback + best = checkpoint.best_model_path if checkpoint is not None and checkpoint.monitor else "" + trainer.test(model, datamodule=datamodule, ckpt_path=best or None) + metrics.update({k: float(v) for k, v in trainer.callback_metrics.items()}) + preds = model.test_predictions + frame = prediction_frame( + preds["unit_ids"], + preds["logits"], + target_type=model.task["target_type"], + targets=preds["targets"], + meta=datamodule.unit_table(datamodule.dataset_test), ) - - train_metrics = trainer.callback_metrics - - if cfg.get("test"): - log.info("Starting testing!") - ckpt_path = None - checkpoint_callback = getattr(trainer, "checkpoint_callback", None) - if checkpoint_callback is not None: - ckpt_path = checkpoint_callback.best_model_path - if ckpt_path == "": - log.warning("Best ckpt not found! Using current weights for testing...") - ckpt_path = None - else: - log.warning( - "Checkpoint callback is not configured. Using current weights for testing..." - ) - trainer.test(model=model, datamodule=datamodule, ckpt_path=ckpt_path) - log.info(f"Best ckpt path: {ckpt_path}") - if output_dir: - from grass_mil.callbacks import write_test_predictions - - written = write_test_predictions(model, str(output_dir)) - if written is not None: - log.info(f"Test predictions written to {written}") - - test_metrics = trainer.callback_metrics - - metric_dict = {**train_metrics, **test_metrics} - - return metric_dict, object_dict + frame.to_csv(output_dir / "test_predictions.csv", index=False) + log.info(f"Test predictions written to {output_dir / 'test_predictions.csv'}") + return metrics, {"cfg": cfg, "datamodule": datamodule, "model": model, "trainer": trainer} @hydra.main(version_base="1.3", config_path="configs", config_name="train.yaml") -def main(cfg: DictConfig) -> Optional[float]: - """Entry point for training.""" - extras(cfg) - - metric_dict, _ = train(cfg) - - metric_value = get_metric_value( - metric_dict=metric_dict, metric_name=cfg.get("optimized_metric") - ) - - return metric_value +def main(cfg: DictConfig) -> None: + train(cfg) if __name__ == "__main__": diff --git a/src/grass_mil/utils/__init__.py b/src/grass_mil/utils/__init__.py deleted file mode 100644 index 4e8148f..0000000 --- a/src/grass_mil/utils/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -from grass_mil.utils.instantiators import instantiate_callbacks, instantiate_loggers -from grass_mil.utils.logging_utils import log_hyperparameters -from grass_mil.utils.pylogger import RankedLogger -from grass_mil.utils.rich_utils import enforce_tags, print_config_tree -from grass_mil.utils.utils import extras, get_metric_value, task_wrapper - -# Declared so these read as a deliberate re-export surface rather than unused -# imports, and so `from grass_mil.utils import *` stays well defined. -__all__ = [ - "instantiate_callbacks", - "instantiate_loggers", - "log_hyperparameters", - "RankedLogger", - "enforce_tags", - "print_config_tree", - "extras", - "get_metric_value", - "task_wrapper", -] diff --git a/src/grass_mil/utils/instantiators.py b/src/grass_mil/utils/instantiators.py deleted file mode 100644 index b823cfb..0000000 --- a/src/grass_mil/utils/instantiators.py +++ /dev/null @@ -1,56 +0,0 @@ -from typing import List - -import hydra -from lightning import Callback -from lightning.pytorch.loggers import Logger -from omegaconf import DictConfig - -from grass_mil.utils import pylogger - -log = pylogger.RankedLogger(__name__, rank_zero_only=True) - - -def instantiate_callbacks(callbacks_cfg: DictConfig) -> List[Callback]: - """Instantiates callbacks from config. - - :param callbacks_cfg: A DictConfig object containing callback configurations. - :return: A list of instantiated callbacks. - """ - callbacks: List[Callback] = [] - - if not callbacks_cfg: - log.warning("No callback configs found! Skipping..") - return callbacks - - if not isinstance(callbacks_cfg, DictConfig): - raise TypeError("Callbacks config must be a DictConfig!") - - for _, cb_conf in callbacks_cfg.items(): - if isinstance(cb_conf, DictConfig) and "_target_" in cb_conf: - log.info(f"Instantiating callback <{cb_conf._target_}>") - callbacks.append(hydra.utils.instantiate(cb_conf)) - - return callbacks - - -def instantiate_loggers(logger_cfg: DictConfig) -> List[Logger]: - """Instantiates loggers from config. - - :param logger_cfg: A DictConfig object containing logger configurations. - :return: A list of instantiated loggers. - """ - logger: List[Logger] = [] - - if not logger_cfg: - log.warning("No logger configs found! Skipping...") - return logger - - if not isinstance(logger_cfg, DictConfig): - raise TypeError("Logger config must be a DictConfig!") - - for _, lg_conf in logger_cfg.items(): - if isinstance(lg_conf, DictConfig) and "_target_" in lg_conf: - log.info(f"Instantiating logger <{lg_conf._target_}>") - logger.append(hydra.utils.instantiate(lg_conf)) - - return logger diff --git a/src/grass_mil/utils/logging_utils.py b/src/grass_mil/utils/logging_utils.py deleted file mode 100644 index 01d6cdd..0000000 --- a/src/grass_mil/utils/logging_utils.py +++ /dev/null @@ -1,57 +0,0 @@ -from typing import Any, Dict - -from lightning_utilities.core.rank_zero import rank_zero_only -from omegaconf import OmegaConf - -from grass_mil.utils import pylogger - -log = pylogger.RankedLogger(__name__, rank_zero_only=True) - - -@rank_zero_only -def log_hyperparameters(object_dict: Dict[str, Any]) -> None: - """Controls which config parts are saved by Lightning loggers. - - Additionally saves: - - Number of model parameters - - :param object_dict: A dictionary containing the following objects: - - `"cfg"`: A DictConfig object containing the main config. - - `"model"`: The Lightning model. - - `"trainer"`: The Lightning trainer. - """ - hparams = {} - - cfg = OmegaConf.to_container(object_dict["cfg"]) - model = object_dict["model"] - trainer = object_dict["trainer"] - - if not trainer.logger: - log.warning("Logger not found! Skipping hyperparameter logging...") - return - - hparams["model"] = cfg["model"] - - # save number of model parameters - hparams["model/params/total"] = sum(p.numel() for p in model.parameters()) - hparams["model/params/trainable"] = sum( - p.numel() for p in model.parameters() if p.requires_grad - ) - hparams["model/params/non_trainable"] = sum( - p.numel() for p in model.parameters() if not p.requires_grad - ) - - hparams["data"] = cfg["data"] - hparams["trainer"] = cfg["trainer"] - - hparams["callbacks"] = cfg.get("callbacks") - hparams["extras"] = cfg.get("extras") - - hparams["task_name"] = cfg.get("task_name") - hparams["tags"] = cfg.get("tags") - hparams["ckpt_path"] = cfg.get("ckpt_path") - hparams["seed"] = cfg.get("seed") - - # send hparams to all loggers - for logger in trainer.loggers: - logger.log_hyperparams(hparams) diff --git a/src/grass_mil/utils/pylogger.py b/src/grass_mil/utils/pylogger.py deleted file mode 100644 index c4ee867..0000000 --- a/src/grass_mil/utils/pylogger.py +++ /dev/null @@ -1,51 +0,0 @@ -import logging -from typing import Mapping, Optional - -from lightning_utilities.core.rank_zero import rank_prefixed_message, rank_zero_only - - -class RankedLogger(logging.LoggerAdapter): - """A multi-GPU-friendly python command line logger.""" - - def __init__( - self, - name: str = __name__, - rank_zero_only: bool = False, - extra: Optional[Mapping[str, object]] = None, - ) -> None: - """Initializes a multi-GPU-friendly python command line logger that logs on all processes - with their rank prefixed in the log message. - - :param name: The name of the logger. Default is ``__name__``. - :param rank_zero_only: Whether to force all logs to only occur on the rank zero process. Default is `False`. - :param extra: (Optional) A dict-like object which provides contextual information. See `logging.LoggerAdapter`. - """ - logger = logging.getLogger(name) - super().__init__(logger=logger, extra=extra) - self.rank_zero_only = rank_zero_only - - def log(self, level: int, msg: str, rank: Optional[int] = None, *args, **kwargs) -> None: - """Delegate a log call to the underlying logger, after prefixing its message with the rank - of the process it's being logged from. If `'rank'` is provided, then the log will only - occur on that rank/process. - - :param level: The level to log at. Look at `logging.__init__.py` for more information. - :param msg: The message to log. - :param rank: The rank to log at. - :param args: Additional args to pass to the underlying logging function. - :param kwargs: Any additional keyword args to pass to the underlying logging function. - """ - if self.isEnabledFor(level): - msg, kwargs = self.process(msg, kwargs) - current_rank = getattr(rank_zero_only, "rank", None) - if current_rank is None: - raise RuntimeError("The `rank_zero_only.rank` needs to be set before use") - msg = rank_prefixed_message(msg, current_rank) - if self.rank_zero_only: - if current_rank == 0: - self.logger.log(level, msg, *args, **kwargs) - else: - if rank is None: - self.logger.log(level, msg, *args, **kwargs) - elif current_rank == rank: - self.logger.log(level, msg, *args, **kwargs) diff --git a/src/grass_mil/utils/rich_utils.py b/src/grass_mil/utils/rich_utils.py deleted file mode 100644 index 778b53d..0000000 --- a/src/grass_mil/utils/rich_utils.py +++ /dev/null @@ -1,101 +0,0 @@ -from pathlib import Path -from typing import Sequence - -import rich -import rich.syntax -import rich.tree -from hydra.core.hydra_config import HydraConfig -from lightning_utilities.core.rank_zero import rank_zero_only -from omegaconf import DictConfig, OmegaConf, open_dict -from rich.prompt import Prompt - -from grass_mil.utils import pylogger - -log = pylogger.RankedLogger(__name__, rank_zero_only=True) - - -@rank_zero_only -def print_config_tree( - cfg: DictConfig, - print_order: Sequence[str] = ( - "data", - "model", - "callbacks", - "logger", - "trainer", - "paths", - "extras", - ), - resolve: bool = False, - save_to_file: bool = False, -) -> None: - """Prints the contents of a DictConfig as a tree structure using the Rich library. - - :param cfg: A DictConfig composed by Hydra. - :param print_order: Determines in what order config components are printed. Default is ``("data", "model", - "callbacks", "logger", "trainer", "paths", "extras")``. - :param resolve: Whether to resolve reference fields of DictConfig. Default is ``False``. - :param save_to_file: Whether to export config to the hydra output folder. Default is ``False``. - """ - style = "dim" - tree = rich.tree.Tree("CONFIG", style=style, guide_style=style) - - queue = [] - - # add fields from `print_order` to queue - for field in print_order: - queue.append(field) if field in cfg else log.warning( - f"Field '{field}' not found in config. Skipping '{field}' config printing..." - ) - - # add all the other fields to queue (not specified in `print_order`) - for field in cfg: - if field not in queue: - queue.append(field) - - # generate config tree from queue - for field in queue: - branch = tree.add(field, style=style, guide_style=style) - - config_group = cfg[field] - if isinstance(config_group, DictConfig): - branch_content = OmegaConf.to_yaml(config_group, resolve=resolve) - else: - branch_content = str(config_group) - - branch.add(rich.syntax.Syntax(branch_content, "yaml")) - - # print config tree - rich.print(tree) - - # save config tree to file - if save_to_file: - Path(cfg.paths.output_dir).mkdir(parents=True, exist_ok=True) - with open(Path(cfg.paths.output_dir, "config_tree.log"), "w") as file: - rich.print(tree, file=file) - - -@rank_zero_only -def enforce_tags(cfg: DictConfig, save_to_file: bool = False) -> None: - """Prompts user to input tags from command line if no tags are provided in config. - - :param cfg: A DictConfig composed by Hydra. - :param save_to_file: Whether to export tags to the hydra output folder. Default is ``False``. - """ - if not cfg.get("tags"): - if "id" in HydraConfig().cfg.hydra.job: - raise ValueError("Specify tags before launching a multirun!") - - log.warning("No tags provided in config. Prompting user to input tags...") - tags = Prompt.ask("Enter a list of comma separated tags", default="dev") - tags = [t.strip() for t in tags.split(",") if t != ""] - - with open_dict(cfg): - cfg.tags = tags - - log.info(f"Tags: {cfg.tags}") - - if save_to_file: - Path(cfg.paths.output_dir).mkdir(parents=True, exist_ok=True) - with open(Path(cfg.paths.output_dir, "tags.log"), "w") as file: - rich.print(cfg.tags, file=file) diff --git a/src/grass_mil/utils/utils.py b/src/grass_mil/utils/utils.py deleted file mode 100644 index 5dbdbc1..0000000 --- a/src/grass_mil/utils/utils.py +++ /dev/null @@ -1,119 +0,0 @@ -import warnings -from importlib.util import find_spec -from typing import Any, Callable, Dict, Optional, Tuple - -from omegaconf import DictConfig - -from grass_mil.utils import pylogger, rich_utils - -log = pylogger.RankedLogger(__name__, rank_zero_only=True) - - -def extras(cfg: DictConfig) -> None: - """Applies optional utilities before the task is started. - - Utilities: - - Ignoring python warnings - - Setting tags from command line - - Rich config printing - - :param cfg: A DictConfig object containing the config tree. - """ - # return if no `extras` config - if not cfg.get("extras"): - log.warning("Extras config not found! ") - return - - # disable python warnings - if cfg.extras.get("ignore_warnings"): - log.info("Disabling python warnings! ") - warnings.filterwarnings("ignore") - - # prompt user to input tags from command line if none are provided in the config - if cfg.extras.get("enforce_tags"): - log.info("Enforcing tags! ") - rich_utils.enforce_tags(cfg, save_to_file=True) - - # pretty print config tree using Rich library - if cfg.extras.get("print_config"): - log.info("Printing config tree with Rich! ") - rich_utils.print_config_tree(cfg, resolve=True, save_to_file=True) - - -def task_wrapper(task_func: Callable) -> Callable: - """Optional decorator that controls the failure behavior when executing the task function. - - This wrapper can be used to: - - make sure loggers are closed even if the task function raises an exception (prevents multirun failure) - - save the exception to a `.log` file - - mark the run as failed with a dedicated file in the `logs/` folder (so we can find and rerun it later) - - etc. (adjust depending on your needs) - - Example: - ``` - @utils.task_wrapper - def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: - ... - return metric_dict, object_dict - ``` - - :param task_func: The task function to be wrapped. - - :return: The wrapped task function. - """ - - def wrap(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: - # execute the task - try: - metric_dict, object_dict = task_func(cfg=cfg) - - # things to do if exception occurs - except Exception as ex: - # save exception to `.log` file - log.exception("") - - # some hyperparameter combinations might be invalid or cause out-of-memory errors - # so when using hparam search plugins like Optuna, you might want to disable - # raising the below exception to avoid multirun failure - raise ex - - # things to always do after either success or exception - finally: - # display output dir path in terminal - log.info(f"Output dir: {cfg.paths.output_dir}") - - # always close wandb run (even if exception occurs so multirun won't fail) - if find_spec("wandb"): # check if wandb is installed - import wandb - - if wandb.run: - log.info("Closing wandb!") - wandb.finish() - - return metric_dict, object_dict - - return wrap - - -def get_metric_value(metric_dict: Dict[str, Any], metric_name: Optional[str]) -> Optional[float]: - """Safely retrieves value of the metric logged in LightningModule. - - :param metric_dict: A dict containing metric values. - :param metric_name: If provided, the name of the metric to retrieve. - :return: If a metric name was provided, the value of the metric. - """ - if not metric_name: - log.info("Metric name is None! Skipping metric value retrieval...") - return None - - if metric_name not in metric_dict: - raise Exception( - f"Metric value not found! \n" - "Make sure metric name logged in LightningModule is correct!\n" - "Make sure `optimized_metric` name in `hparams_search` config is correct!" - ) - - metric_value = metric_dict[metric_name].item() - log.info(f"Retrieved metric value! <{metric_name}={metric_value}>") - - return metric_value From 801f1972ab63c5d0cb815a91330f1c89afbd3152 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 15 Oct 2026 08:51:57 +0200 Subject: [PATCH 2/9] Replace the analysis layer: niches as the primary results, metaclusters optional - grass_mil.niches: niche discovery in the attention-head space (cosine HDBSCAN on a sample, k-NN assignment), niches numbered by margin with Background for unassigned cells; per-niche composition, empirical association and attention-weighted margin with a bootstrap CI over units; per-niche Moran's I, neighbourhood enrichment, cross-L at r* and contact enrichment; niche co-localisation; optional survival (Cox per niche share, Kaplan-Meier of score tertiles); figures following the preprint panels - grass_mil.niches.replication: grass-mil-metaclusters on the same discovery - grass_mil.spatial: one implementation of every statistic, shared by the niche descriptors and grass-mil-spatial-stats - remove the report/plugin stack, grass-mil-scores, the instance export and the evaluation package - packaging: 1.1.0, seven entry points, dependencies trimmed to imports Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NqExFvaq8RhDYBj6RShCMk --- .gitignore | 16 +- .project-root | 2 - .vscode/settings.json | 13 - CITATION.cff | 4 +- pyproject.toml | 93 +-- src/grass_mil/__init__.py | 6 - src/grass_mil/_compat.py | 56 -- src/grass_mil/configs/metaclusters.yaml | 60 +- src/grass_mil/configs/niches.yaml | 82 ++- src/grass_mil/configs/spatial_stats.yaml | 41 +- src/grass_mil/evaluation/__init__.py | 39 - src/grass_mil/evaluation/cli.py | 271 ------- src/grass_mil/evaluation/operating_points.py | 136 ---- src/grass_mil/evaluation/stats.py | 187 ----- src/grass_mil/extract.py | 8 +- src/grass_mil/inference/cell_cache.py | 19 +- src/grass_mil/inference/predictions.py | 10 +- src/grass_mil/interpretability/__init__.py | 51 -- .../interpretability/core/__init__.py | 27 - .../interpretability/core/agreement.py | 103 --- .../interpretability/core/attribution.py | 347 --------- .../interpretability/core/biomarkers.py | 66 -- .../interpretability/core/clustering.py | 309 -------- src/grass_mil/interpretability/core/data.py | 73 -- src/grass_mil/interpretability/core/niches.py | 167 ----- .../interpretability/core/reduction.py | 46 -- .../interpretability/core/transfer.py | 217 ------ .../interpretability/metaclusters/__init__.py | 22 - .../interpretability/metaclusters/basis.py | 252 ------- .../interpretability/metaclusters/match.py | 420 ----------- .../interpretability/metaclusters/records.py | 183 ----- .../interpretability/metaclusters/transfer.py | 231 ------ .../interpretability/metaclusters_cli.py | 187 ----- .../interpretability/niches/__init__.py | 16 - .../interpretability/niches/colocalisation.py | 34 - .../interpretability/niches/figures.py | 180 ----- .../interpretability/niches/margins.py | 108 --- .../interpretability/niches/pipeline.py | 372 ---------- .../interpretability/niches/tables.py | 131 ---- src/grass_mil/interpretability/niches_cli.py | 105 --- src/grass_mil/interpretability/pipeline.py | 302 -------- .../interpretability/plugins/__init__.py | 15 - .../interpretability/plugins/base.py | 21 - .../interpretability/plugins/builtin.py | 689 ------------------ .../interpretability/plugins/registry.py | 40 - src/grass_mil/interpretability/report_cli.py | 81 -- .../interpretability/reporting/__init__.py | 3 - .../interpretability/reporting/html.py | 74 -- .../interpretability/reporting/pdf.py | 33 - .../reporting/plotly_builders.py | 492 ------------- .../interpretability/reporting/render.py | 61 -- .../interpretability/reporting/sections.py | 32 - .../interpretability/reporting/snapshot.py | 27 - .../interpretability/tier2/__init__.py | 17 - .../interpretability/tier2/autocorrelation.py | 213 ------ .../interpretability/tier2/cell_level.py | 311 -------- .../interpretability/tier2/filtration.py | 110 --- .../interpretability/tier2/neighborhood.py | 421 ----------- .../interpretability/tier2/ripley.py | 227 ------ .../interpretability/tier2/tissue_graph.py | 207 ------ src/grass_mil/models/augmentations.py | 12 +- src/grass_mil/models/encoder.py | 26 +- src/grass_mil/models/mil.py | 17 +- src/grass_mil/models/objectives.py | 52 +- src/grass_mil/models/optimization.py | 8 +- src/grass_mil/models/pretraining.py | 34 +- src/grass_mil/models/probe.py | 8 +- src/grass_mil/models/supervised.py | 91 ++- src/grass_mil/nested_cv.py | 63 +- src/grass_mil/niches/__init__.py | 18 + src/grass_mil/niches/analysis.py | 301 ++++++++ src/grass_mil/niches/attribution.py | 50 ++ src/grass_mil/niches/cells.py | 127 ++++ src/grass_mil/niches/cli.py | 62 ++ src/grass_mil/niches/descriptors.py | 254 +++++++ src/grass_mil/niches/discovery.py | 147 ++++ src/grass_mil/niches/figures.py | 283 +++++++ src/grass_mil/niches/replication/__init__.py | 1 + src/grass_mil/niches/replication/cli.py | 170 +++++ src/grass_mil/niches/replication/stages.py | 367 ++++++++++ src/grass_mil/niches/survival.py | 138 ++++ src/grass_mil/niches/tables.py | 93 +++ src/grass_mil/predict.py | 4 +- src/grass_mil/runtime.py | 9 +- src/grass_mil/spatial/__init__.py | 2 + src/grass_mil/spatial/cli.py | 126 ++++ src/grass_mil/spatial/graph_stats.py | 114 +++ src/grass_mil/spatial/neighbourhoods.py | 33 + src/grass_mil/spatial/stats.py | 73 ++ src/grass_mil/spatial_stats/__init__.py | 24 - .../spatial_stats/cellular_neighbourhoods.py | 51 -- src/grass_mil/spatial_stats/cli.py | 61 -- src/grass_mil/spatial_stats/contacts.py | 56 -- src/grass_mil/spatial_stats/cross_k.py | 48 -- .../spatial_stats/nhood_enrichment.py | 103 --- src/grass_mil/spatial_stats/run.py | 224 ------ src/grass_mil/spatial_stats/tests.py | 64 -- src/grass_mil/train.py | 2 +- 98 files changed, 2755 insertions(+), 8627 deletions(-) delete mode 100644 .project-root delete mode 100644 .vscode/settings.json delete mode 100644 src/grass_mil/_compat.py delete mode 100644 src/grass_mil/evaluation/__init__.py delete mode 100644 src/grass_mil/evaluation/cli.py delete mode 100644 src/grass_mil/evaluation/operating_points.py delete mode 100644 src/grass_mil/evaluation/stats.py delete mode 100644 src/grass_mil/interpretability/__init__.py delete mode 100644 src/grass_mil/interpretability/core/__init__.py delete mode 100644 src/grass_mil/interpretability/core/agreement.py delete mode 100644 src/grass_mil/interpretability/core/attribution.py delete mode 100644 src/grass_mil/interpretability/core/biomarkers.py delete mode 100644 src/grass_mil/interpretability/core/clustering.py delete mode 100644 src/grass_mil/interpretability/core/data.py delete mode 100644 src/grass_mil/interpretability/core/niches.py delete mode 100644 src/grass_mil/interpretability/core/reduction.py delete mode 100644 src/grass_mil/interpretability/core/transfer.py delete mode 100644 src/grass_mil/interpretability/metaclusters/__init__.py delete mode 100644 src/grass_mil/interpretability/metaclusters/basis.py delete mode 100644 src/grass_mil/interpretability/metaclusters/match.py delete mode 100644 src/grass_mil/interpretability/metaclusters/records.py delete mode 100644 src/grass_mil/interpretability/metaclusters/transfer.py delete mode 100644 src/grass_mil/interpretability/metaclusters_cli.py delete mode 100644 src/grass_mil/interpretability/niches/__init__.py delete mode 100644 src/grass_mil/interpretability/niches/colocalisation.py delete mode 100644 src/grass_mil/interpretability/niches/figures.py delete mode 100644 src/grass_mil/interpretability/niches/margins.py delete mode 100644 src/grass_mil/interpretability/niches/pipeline.py delete mode 100644 src/grass_mil/interpretability/niches/tables.py delete mode 100644 src/grass_mil/interpretability/niches_cli.py delete mode 100644 src/grass_mil/interpretability/pipeline.py delete mode 100644 src/grass_mil/interpretability/plugins/__init__.py delete mode 100644 src/grass_mil/interpretability/plugins/base.py delete mode 100644 src/grass_mil/interpretability/plugins/builtin.py delete mode 100644 src/grass_mil/interpretability/plugins/registry.py delete mode 100644 src/grass_mil/interpretability/report_cli.py delete mode 100644 src/grass_mil/interpretability/reporting/__init__.py delete mode 100644 src/grass_mil/interpretability/reporting/html.py delete mode 100644 src/grass_mil/interpretability/reporting/pdf.py delete mode 100644 src/grass_mil/interpretability/reporting/plotly_builders.py delete mode 100644 src/grass_mil/interpretability/reporting/render.py delete mode 100644 src/grass_mil/interpretability/reporting/sections.py delete mode 100644 src/grass_mil/interpretability/reporting/snapshot.py delete mode 100644 src/grass_mil/interpretability/tier2/__init__.py delete mode 100644 src/grass_mil/interpretability/tier2/autocorrelation.py delete mode 100644 src/grass_mil/interpretability/tier2/cell_level.py delete mode 100644 src/grass_mil/interpretability/tier2/filtration.py delete mode 100644 src/grass_mil/interpretability/tier2/neighborhood.py delete mode 100644 src/grass_mil/interpretability/tier2/ripley.py delete mode 100644 src/grass_mil/interpretability/tier2/tissue_graph.py create mode 100644 src/grass_mil/niches/__init__.py create mode 100644 src/grass_mil/niches/analysis.py create mode 100644 src/grass_mil/niches/attribution.py create mode 100644 src/grass_mil/niches/cells.py create mode 100644 src/grass_mil/niches/cli.py create mode 100644 src/grass_mil/niches/descriptors.py create mode 100644 src/grass_mil/niches/discovery.py create mode 100644 src/grass_mil/niches/figures.py create mode 100644 src/grass_mil/niches/replication/__init__.py create mode 100644 src/grass_mil/niches/replication/cli.py create mode 100644 src/grass_mil/niches/replication/stages.py create mode 100644 src/grass_mil/niches/survival.py create mode 100644 src/grass_mil/niches/tables.py create mode 100644 src/grass_mil/spatial/__init__.py create mode 100644 src/grass_mil/spatial/cli.py create mode 100644 src/grass_mil/spatial/graph_stats.py create mode 100644 src/grass_mil/spatial/neighbourhoods.py create mode 100644 src/grass_mil/spatial/stats.py delete mode 100644 src/grass_mil/spatial_stats/__init__.py delete mode 100644 src/grass_mil/spatial_stats/cellular_neighbourhoods.py delete mode 100644 src/grass_mil/spatial_stats/cli.py delete mode 100644 src/grass_mil/spatial_stats/contacts.py delete mode 100644 src/grass_mil/spatial_stats/cross_k.py delete mode 100644 src/grass_mil/spatial_stats/nhood_enrichment.py delete mode 100644 src/grass_mil/spatial_stats/run.py delete mode 100644 src/grass_mil/spatial_stats/tests.py diff --git a/.gitignore b/.gitignore index 69d0e65..d446e47 100644 --- a/.gitignore +++ b/.gitignore @@ -127,12 +127,8 @@ dmypy.json # Pyre type checker .pyre/ -### VisualStudioCode -.vscode/* -!.vscode/settings.json -!.vscode/tasks.json -!.vscode/launch.json -!.vscode/extensions.json +### Editors +.vscode/ *.code-workspace # JetBrains @@ -143,17 +139,11 @@ dmypy.json *.tar *.tar.gz -# Lightning-Hydra-Template -configs/local/default.yaml -src/grass_mil/configs/local/*.yaml -!src/grass_mil/configs/local/.gitkeep +# Run outputs and local data /data/ /logs/ .env -# Aim logging -.aim - # Vignette scratch space (downloaded data, graphs, runs) vignettes/vignette_work/ diff --git a/.project-root b/.project-root deleted file mode 100644 index 63eab77..0000000 --- a/.project-root +++ /dev/null @@ -1,2 +0,0 @@ -# this file is required for inferring the project root directory -# do not delete diff --git a/.vscode/settings.json b/.vscode/settings.json deleted file mode 100644 index 893f70e..0000000 --- a/.vscode/settings.json +++ /dev/null @@ -1,13 +0,0 @@ -{ - "editor.formatOnSave": true, - "editor.codeActionsOnSave": { - "source.fixAll.ruff": "always", - "source.organizeImports.ruff": "always" - }, - "[python]": { - "editor.defaultFormatter": "charliermarsh.ruff" - }, - "ruff.format.args": [ - "--line-length=99" - ] -} diff --git a/CITATION.cff b/CITATION.cff index e12d63c..eabc2db 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -2,8 +2,8 @@ cff-version: 1.2.0 message: "If you use GRASS-MIL, please cite the preprint." title: "GRASS-MIL: graph-based representation and discovery of phenotype-associated spatial structures with multiple instance learning" type: software -version: 1.0.0 -date-released: 2026-10-13 +version: 1.1.0 +date-released: 2026-10-15 license: MIT repository-code: https://github.com/BoevaLab/grass-mil authors: diff --git a/pyproject.toml b/pyproject.toml index a4a1e31..40ec871 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,97 +4,70 @@ build-backend = "hatchling.build" [project] name = "grass-mil" -version = "1.0.0" +version = "1.1.0" description = "GRASS-MIL: graph-based representation and discovery of phenotype-associated spatial structures with multiple instance learning" readme = "README.md" requires-python = ">=3.11" license = { file = "LICENSE" } +authors = [ + { name = "Lovro Rabuzin" }, + { name = "Josephine Yates" }, + { name = "Mina Xu" }, + { name = "Valentina Boeva" }, +] dependencies = [ - "torch>=2.0", - "lightning>=2.0", - # Imported directly by utils/{pylogger,rich_utils,logging_utils}.py rather - # than relied on transitively through lightning. - "lightning-utilities>=0.10", - "torchmetrics>=0.11.4", + "torch>=2.1", + "lightning>=2.1", "torch-geometric>=2.5", "hydra-core>=1.3.5,<1.4", "hydra-colorlog==1.2.0", - "omegaconf", + "omegaconf>=2.3", "numpy", - "pandas", + "pandas>=2.0", "scipy>=1.11", - "scikit-learn", - # Imported directly by interpretability/core/transfer.py to persist niche - # transfer bundles; not relied on transitively through scikit-learn. + "scikit-learn>=1.3", "joblib", + "matplotlib>=3.7", "rich", - "rootutils", ] [project.optional-dependencies] -interpretability = ["plotly", "umap-learn", "jinja2", "matplotlib"] -stats = ["statsmodels"] -plots = ["matplotlib", "umap-learn"] -report = ["kaleido", "playwright"] -spatial-stats = ["squidpy", "scanpy", "anndata"] -io = ["anndata", "shapely"] -# SingleCellExperiment / .rds inputs. Also needs an R runtime with Bioconductor; -# rpy2 alone is not sufficient. -sce = ["rpy2"] +# h5ad inputs and polygon units. +io = ["anndata>=0.10", "shapely>=2.0"] +# SingleCellExperiment (.rds) inputs; also needs R with Bioconductor. +sce = ["rpy2>=3.5"] +test = ["pytest>=7", "pytest-cov"] + +[project.urls] +Homepage = "https://github.com/BoevaLab/grass-mil" +Preprint = "https://www.biorxiv.org/content/10.64898/2026.10.08.757722v1" [project.scripts] grass-mil-train = "grass_mil.train:main" -grass-mil-eval = "grass_mil.eval:main" -grass-mil-loocv = "grass_mil.loocv:main" -grass-mil-probe = "grass_mil.probe:main" grass-mil-nested-cv = "grass_mil.nested_cv:main" +grass-mil-predict = "grass_mil.predict:main" grass-mil-extract = "grass_mil.extract:main" -grass-mil-scores = "grass_mil.evaluation.cli:main" -grass-mil-niches = "grass_mil.interpretability.niches_cli:main" -grass-mil-metaclusters = "grass_mil.interpretability.metaclusters_cli:main" -grass-mil-spatial-stats = "grass_mil.spatial_stats.cli:main" -grass-mil-predict = "grass_mil.inference.predict:main" -grass-mil-report = "grass_mil.interpretability.report_cli:main" +grass-mil-niches = "grass_mil.niches.cli:main" +grass-mil-metaclusters = "grass_mil.niches.replication.cli:main" +grass-mil-spatial-stats = "grass_mil.spatial.cli:main" [tool.hatch.build.targets.wheel] packages = ["src/grass_mil"] [tool.pytest.ini_options] pythonpath = ["src", "."] -addopts = [ - "--color=yes", - "--durations=0", - "--strict-markers", - "--doctest-modules", -] +addopts = ["--strict-markers"] filterwarnings = [ - "ignore::DeprecationWarning", - "ignore::UserWarning", -] -log_cli = "True" -markers = [ - "slow: end-to-end chains that take more than a minute (deselect with -m 'not slow')", -] -minversion = "6.0" -testpaths = "tests/" - -[tool.coverage.report] -exclude_lines = [ - "pragma: nocover", - "raise NotImplementedError", - "raise NotImplementedError()", - "if __name__ == .__main__.:", + # Lightning's hints about dataloader workers and logging intervals. + "ignore::lightning.fabric.utilities.warnings.PossibleUserWarning", ] +markers = ["slow: the end-to-end chain on synthetic data (deselect with -m 'not slow')"] +testpaths = ["tests"] [tool.ruff] line-length = 99 -target-version = "py311" # matches requires-python floor +target-version = "py311" exclude = ["logs", "data"] [tool.ruff.lint] -# E402: train.py/eval.py must call rootutils.setup_root() before importing the -# package, so module-level imports legitimately follow executable code. -# E501: a handful of long string literals the formatter will not wrap. -# F401/F841/E203 were previously ignored too. They are now clean and enforced, -# so unused imports and dead locals cannot creep back in. -ignore = ["E402", "E501"] +ignore = ["E501"] diff --git a/src/grass_mil/__init__.py b/src/grass_mil/__init__.py index 650186b..4fc10ba 100644 --- a/src/grass_mil/__init__.py +++ b/src/grass_mil/__init__.py @@ -3,12 +3,6 @@ from importlib.metadata import PackageNotFoundError, version -from grass_mil._compat import allow_omegaconf_in_checkpoints - -# Applied at import so every entrypoint and any user code that loads a -# checkpoint gets it. See _compat for why this is needed under torch>=2.6. -allow_omegaconf_in_checkpoints() - try: __version__ = version("grass-mil") except PackageNotFoundError: # source checkout without an installed distribution diff --git a/src/grass_mil/_compat.py b/src/grass_mil/_compat.py deleted file mode 100644 index 9409419..0000000 --- a/src/grass_mil/_compat.py +++ /dev/null @@ -1,56 +0,0 @@ -"""Cross-version compatibility shims applied at package import. - -Currently only one: allowlisting the OmegaConf types that appear inside -Lightning checkpoints. - -PyTorch 2.6 flipped the default of ``torch.load(weights_only=...)`` from False -to True. Our own loaders pass ``weights_only=False`` explicitly, because they -read first-party artifacts. Lightning's checkpoint loader -(``lightning.fabric.utilities.cloud_io._load``) is not ours to change, and it -fails on any checkpoint whose ``hyper_parameters`` hold an OmegaConf container --- which is every checkpoint this package writes, since the modules save their -Hydra config. - -Allowlisting is preferable to disabling the check: the entries below are inert -container and metadata types whose construction cannot execute arbitrary code, -so the protection still applies to everything else in the file. -""" - -from __future__ import annotations - -import collections -import typing - - -def allow_omegaconf_in_checkpoints() -> None: - """Register OmegaConf container types as safe for ``torch.load``. - - A no-op on torch versions without ``add_safe_globals`` (< 2.4) and on any - OmegaConf layout that no longer exposes these names, so an upstream rename - degrades to the previous behaviour rather than breaking package import. - """ - try: - import torch.serialization as ts - except Exception: # torch missing or partially initialised - return - - add = getattr(ts, "add_safe_globals", None) - if add is None: - return - - safe: list = [dict, list, int, typing.Any, collections.defaultdict] - try: - from omegaconf.base import ContainerMetadata, Metadata - from omegaconf.dictconfig import DictConfig - from omegaconf.listconfig import ListConfig - from omegaconf.nodes import AnyNode - - safe += [DictConfig, ListConfig, ContainerMetadata, Metadata, AnyNode] - except Exception: - return - - try: - add(safe) - except Exception: - # Never let a compatibility shim break importing the package. - return diff --git a/src/grass_mil/configs/metaclusters.yaml b/src/grass_mil/configs/metaclusters.yaml index 37f7e83..d866236 100644 --- a/src/grass_mil/configs/metaclusters.yaml +++ b/src/grass_mil/configs/metaclusters.yaml @@ -1,69 +1,31 @@ -# grass-mil-metaclusters: cross-fold niche replication. +# grass-mil-metaclusters (optional): replication of niches across the models of +# a cross-validation (see docs/niche_analysis.md). # -# grass-mil-metaclusters cells_dir=/cells labels.label_file=labels.csv \ -# labels.label_column=STAS labels.patient_column=PATIENT stage=all -# -# Stages: cluster (each fold model's held-out cells in its own space, fixed -# settings) -> match (Hellinger distance on composition, cannot-link within -# fold, threshold scanned or fixed) -> transfer (every model labels every -# region; train vs held-out, pairwise ARI, consensus patient shares) -> -# basis (shared label-free k-means basis with per-model attribution). +# grass-mil-metaclusters cells_dir=/cells defaults: - _self_ - paths: default - hydra: default task_name: metaclusters -cells_dir: ??? # directory holding one cache subdirectory per fold model -models: null # subset of model names; null = every cache with held-out regions +cells_dir: ??? # one cache per fold model (written by grass-mil-extract from a nested CV) output_dir: ${paths.output_dir} stage: all # all | cluster | match | transfer | basis -labels: - label_file: ??? # region labels (and patient column) - id_column: REGION_ID - label_column: ??? - patient_column: null positive_class: 1 positive_name: positive - -# Stage 1 clustering (same keys as grass-mil-niches) -reduction: - method: pca - n_components: 10 - random_state: 42 -clustering: - method: hdbscan - n_clusters: null - random_state: 42 - params: - min_cluster_size: 300 - min_samples: 1 - cluster_selection_method: eom -sample_size: 100000 -l2_normalise: true -assign_k: 15 -seed: 0 +# Niche discovery of every model; same keys as `clustering` of grass-mil-niches. +clustering: {} n_boot: 2000 -# Stage 2 matching match: - distance: hellinger # hellinger | aitchison - threshold: null # fixed threshold (e.g. 0.40); null = scan - scan: {min: 0.15, max: 0.6, n: 46} - min_folds: 3 # "replicated" = members from at least this many folds + threshold: null # Hellinger distance threshold; null = scan + scan: [0.15, 0.6, 46] # start, stop, number of thresholds + min_models: 3 # a metacluster is replicated when found by this many models -# Stage 3 transfer -transfer: - min_cells: 100 - present_threshold: 0.02 - -# Stage 4 shared basis basis: enabled: true - space_model: null # model name whose embedding defines the basis (null = first); or a cache dir + model: null # model whose space defines the partition (null = the first) + space: embedding # attention | embedding k: 20 sample_size: 300000 - n_init: 10 seed: 0 - -figures: true diff --git a/src/grass_mil/configs/niches.yaml b/src/grass_mil/configs/niches.yaml index 57a9c57..5b33c47 100644 --- a/src/grass_mil/configs/niches.yaml +++ b/src/grass_mil/configs/niches.yaml @@ -1,54 +1,58 @@ -# grass-mil-niches: niche discovery on a trained model's cell cache. +# grass-mil-niches: niches of one model's cells (see docs/niche_analysis.md). # -# grass-mil-niches cells_dir=/cells/fold2 output_dir=niches_fold2 -# grass-mil-niches cells_dir=... clustering.method=kmeans clustering.n_clusters=12 -# grass-mil-niches cells_dir=... space=attention attention_ckpt=/checkpoints/last.ckpt +# grass-mil-niches cells_dir=/cells/fold0 # -# The defaults are the preprint's recipe; every reduction / clustering key is -# a plain override (any backend of grass_mil.interpretability.core.clustering). +# Cells are clustered in the model's attention space, each niche is described by +# its composition, empirical association and attention-weighted margin, and by +# spatial descriptors computed within the niche. defaults: - _self_ - paths: default - hydra: default task_name: niches -cells_dir: ??? # one model's cache directory written by grass-mil-extract +cells_dir: ??? # one model's cache, written by grass-mil-extract output_dir: ${paths.output_dir} -regions: held_out # held_out | all | list of region ids | path to .json/.txt -labels: - # Optional label CSV overriding the labels stored in the cache. - label_file: null - id_column: REGION_ID - label_column: null -space: gnn # gnn (encoder embedding) | attention (8-d pre-tanh attention space) -attention_ckpt: null # checkpoint of the model; required for space=attention +units: held_out # held_out (units the model did not train on) | all | list / file of unit ids positive_class: 1 positive_name: positive -title: null # report title (default from the cache name) -held_out_auroc: null # optional, printed in the report header +title: null -reduction: - method: pca # pca | none | umap - n_components: 10 - random_state: 42 clustering: - method: hdbscan # hdbscan | kmeans | agglomerative | dbscan | optics | spectral - n_clusters: null # for kmeans / agglomerative / spectral - random_state: 42 - params: - min_cluster_size: 300 - min_samples: 1 - cluster_selection_method: eom -sample_size: 100000 # cells used to fit the reducer and the clustering (null = all) -l2_normalise: true # cluster on L2-normalised coordinates (cosine geometry) -assign_k: 15 # k-NN assignment of unsampled cells on the reduced coordinates -seed: 0 + space: attention # attention | embedding + pca_components: null # e.g. 10 with space=embedding + cosine: true + method: hdbscan # hdbscan | kmeans + min_cluster_size: 300 + min_samples: 1 + selection: eom + n_clusters: 10 # kmeans only + sample_size: 100000 # cells used to fit the clustering (null = all) + assign_k: 15 # neighbours assigning the remaining cells + seed: 0 + +present_threshold: 0.02 # a niche is present in a unit above this share +enrichment_threshold: 1.5 # cell types over-represented by more than this factor +n_boot: 2000 # bootstrap over units for the margin interval + +descriptors: + analyses: [moran, enrichment, cross_l, contacts] + min_cells: 20 # cells of a niche in a unit needed to describe it + n_perms: 100 + n_radii: 20 + max_radius: null # null = a quarter of the median unit diagonal + min_contacts: 100 + max_edge_length: null + seed: 0 + +# Optional: niche shares versus survival (one row per patient). +survival: + file: null + id_column: patient_id + time_column: time + event_column: event + pseudo_share: 0.001 + score: [] # e.g. ["Niche 1", "-Niche 4"] for Kaplan-Meier curves of score tertiles -present_threshold: 0.02 # a niche is "present" in a region above this share -n_boot: 2000 # bootstrap over regions for the margin CI figures: true -umap: - enabled: true - n_neighbors: 30 - min_dist: 0.1 - metric: cosine +render_units: 4 # maps of the most positive and most negative units diff --git a/src/grass_mil/configs/spatial_stats.yaml b/src/grass_mil/configs/spatial_stats.yaml index 7b55016..03f2808 100644 --- a/src/grass_mil/configs/spatial_stats.yaml +++ b/src/grass_mil/configs/spatial_stats.yaml @@ -1,11 +1,10 @@ -# grass-mil-spatial-stats: classical spatial statistics per region, compared -# between classes at the patient level. +# grass-mil-spatial-stats (optional): model-free spatial statistics of every unit, +# compared between positive and other patients. # -# grass-mil-spatial-stats cells_dir=/cells/ \ -# labels.label_file=labels.csv labels.label_column=STAS labels.patient_column=PATIENT \ -# reference=[tumor] partners=[myeloid,CD8] infiltrating=[myeloid,CD8] +# grass-mil-spatial-stats cells_dir=/cells/fold0 reference=[tumor] # -# Any model's cache works: only cell types, positions and the cell graph are used. +# Only the cell types, positions and the cell graph of the cache are used, so +# any model's cache of the cohort gives the same result. defaults: - _self_ - paths: default @@ -14,27 +13,17 @@ defaults: task_name: spatial_stats cells_dir: ??? output_dir: ${paths.output_dir} -regions: null # list / .json / .txt of region ids; null = all cached regions -labels: - label_file: ??? - id_column: REGION_ID - label_column: ??? - patient_column: null -positive_name: positive -analyses: [nhood, cn, contacts, crossk] +units: all # all | held_out | list / file of unit ids +positive_class: 1 +analyses: [enrichment, neighbourhoods, contacts, cross_l] # Optional fine -> coarse cell-type grouping ({fine: group}); identity by default. group_map: {} -reference: [] # reference group(s) for contacts and cross-K (e.g. [tumor]); empty skips both -partners: null # partner groups (null = every other group) -infiltrating: [] # groups whose "intra-reference" fraction is reported -nhood: - n_perms: 200 - backend: numpy # numpy | squidpy -cn: - window_k: 10 - n_cn: 10 +reference: [] # group(s) whose contacts and cross-L with the other groups are measured +n_perms: 200 # enrichment permutations +neighbourhoods: + window: 10 # nearest neighbours per window + n: 10 # neighbourhood clusters sample_size: 200000 -crossk: - radii: [20, 50, 100] +radii: [20.0, 50.0, 100.0] # cross-L radii, coordinate units +min_patients: 10 # per class, for a test seed: 0 -min_per_group: 10 # patients per class needed for a test diff --git a/src/grass_mil/evaluation/__init__.py b/src/grass_mil/evaluation/__init__.py deleted file mode 100644 index 31908c7..0000000 --- a/src/grass_mil/evaluation/__init__.py +++ /dev/null @@ -1,39 +0,0 @@ -"""Evaluation statistics for held-out scores (AUROC CIs, DeLong, operating points).""" - -from .operating_points import ( - confusion_counts, - flag_top_within_fold, - likelihood_ratios, - metrics_from_confusion, - operating_point, - operating_point_bootstrap, - ppv_npv_at_prevalence, - threshold_points, -) -from .stats import ( - auroc, - bootstrap_auroc_ci, - delong_test, - per_fold_auroc, - rank_within_fold, - stratified_auroc, - z_within_fold, -) - -__all__ = [ - "auroc", - "bootstrap_auroc_ci", - "confusion_counts", - "delong_test", - "flag_top_within_fold", - "likelihood_ratios", - "metrics_from_confusion", - "operating_point", - "operating_point_bootstrap", - "per_fold_auroc", - "ppv_npv_at_prevalence", - "rank_within_fold", - "stratified_auroc", - "threshold_points", - "z_within_fold", -] diff --git a/src/grass_mil/evaluation/cli.py b/src/grass_mil/evaluation/cli.py deleted file mode 100644 index cdafff2..0000000 --- a/src/grass_mil/evaluation/cli.py +++ /dev/null @@ -1,271 +0,0 @@ -"""``grass-mil-scores``: AUROC CIs, DeLong comparisons and operating points.""" - -from __future__ import annotations - -import json -from pathlib import Path -from typing import Any, Dict, List - -import hydra -import numpy as np -import pandas as pd -from omegaconf import DictConfig - -from grass_mil._paths import setup_project_root - -setup_project_root() - -from grass_mil.evaluation.operating_points import ( # noqa: E402 - flag_top_within_fold, - likelihood_ratios, - operating_point, - operating_point_bootstrap, - ppv_npv_at_prevalence, - threshold_points, -) -from grass_mil.evaluation.stats import ( # noqa: E402 - auroc, - bootstrap_auroc_ci, - delong_test, - per_fold_auroc, - rank_within_fold, - stratified_auroc, - z_within_fold, -) -from grass_mil.utils import RankedLogger # noqa: E402 - -log = RankedLogger(__name__, rank_zero_only=True) - -__all__ = ["run_scores", "main"] - - -def _prepare_table(cfg: DictConfig) -> pd.DataFrame: - df = pd.read_csv(str(cfg.input)) - df.columns = [str(c).strip() for c in df.columns] - label = str(cfg.label_column) - scores = [str(s) for s in cfg.scores] - fold = cfg.get("fold_column") - patient = cfg.get("patient_column") - strata = [str(s) for s in (cfg.get("strata") or [])] - for col in [label] + scores + ([str(fold)] if fold else []) + strata: - if col not in df.columns: - raise ValueError(f"Column '{col}' not in {cfg.input}; columns: {list(df.columns)}") - df = df.dropna(subset=[label] + scores) - if patient and str(patient) in df.columns and bool(cfg.get("patient_level", True)): - agg: Dict[str, Any] = {s: "mean" for s in scores} - agg[label] = "max" - if fold: - agg[str(fold)] = "first" - for s in strata: - agg[s] = "first" - agg["n_units"] = ("__one", "sum") - df = df.assign(__one=1) - grouped = df.groupby(str(patient)) - out = grouped.agg({k: v for k, v in agg.items() if k != "n_units"}) - out["n_units"] = grouped["__one"].sum() - df = out.reset_index() - df[label] = (df[label].astype(float) >= 0.5).astype(int) - return df - - -def _standardise(cfg: DictConfig, df: pd.DataFrame, col: str) -> np.ndarray: - how = str(cfg.get("standardise", "z")).lower() - fold = cfg.get("fold_column") - values = df[col].to_numpy(dtype=float) - if not fold or how == "none": - return values - folds = df[str(fold)].to_numpy() - if how == "z": - return z_within_fold(values, folds) - if how == "rank": - return rank_within_fold(values, folds) - raise ValueError("standardise must be z, rank or none") - - -def run_scores(cfg: DictConfig) -> Dict[str, Any]: - out_dir = Path(str(cfg.output_dir)) - out_dir.mkdir(parents=True, exist_ok=True) - df = _prepare_table(cfg) - label = str(cfg.label_column) - y = df[label].to_numpy(dtype=int) - fold = cfg.get("fold_column") - folds = df[str(fold)].to_numpy() if fold else None - n_boot = int(cfg.n_boot) - seed = int(cfg.seed) - unit = ( - "patients" - if cfg.get("patient_column") and bool(cfg.get("patient_level", True)) - else "units" - ) - prevalence = float(y.mean()) - md: List[str] = [ - f"# Held-out score evaluation ({len(y)} {unit}, {int(y.sum())} positive, prevalence {100 * prevalence:.0f} %)", - "", - f"Scores standardised within fold ({cfg.get('standardise', 'z')}) before pooling; " - f"{n_boot} bootstrap resamples, seed {seed}.", - "", - ] - - scores = [str(s) for s in cfg.scores] - std_scores = {s: _standardise(cfg, df, s) for s in scores} - summary_rows = [] - md += [ - "## AUROC", - "", - "| score | pooled AUROC [95% CI] | per-fold mean ± sd | raw pooled |", - "|---|---|---|---|", - ] - for s in scores: - ci = bootstrap_auroc_ci(y, std_scores[s], n_boot=n_boot, seed=seed) - row = { - "score": s, - "auroc": ci.auroc, - "lower": ci.lower, - "upper": ci.upper, - "raw_pooled_auroc": auroc(y, df[s].to_numpy(float)), - } - pf_txt = "–" - if folds is not None: - pf = per_fold_auroc(y, df[s].to_numpy(float), folds) - row.update({f"perfold_{k}": v for k, v in pf.items()}) - pf_txt = f"{pf['mean']:.3f} ± {pf['sd']:.3f} ({pf['n_folds']} folds)" - summary_rows.append(row) - md.append( - f"| {s} | {ci.auroc:.3f} [{ci.lower:.3f}, {ci.upper:.3f}] | {pf_txt} | {row['raw_pooled_auroc']:.3f} |" - ) - pd.DataFrame(summary_rows).to_csv(out_dir / "scores_summary.csv", index=False) - - delong_rows = [] - if len(scores) > 1: - md += [ - "", - "## DeLong comparisons (pooled, standardised scores)", - "", - "| score 1 | score 2 | ΔAUROC | z | p |", - "|---|---|---|---|---|", - ] - for i in range(len(scores)): - for j in range(i + 1, len(scores)): - d = delong_test(y, std_scores[scores[i]], std_scores[scores[j]]) - delong_rows.append( - { - "score_1": scores[i], - "score_2": scores[j], - "auroc_1": d.auroc_1, - "auroc_2": d.auroc_2, - "delta": d.delta, - "z": d.z, - "p_value": d.p_value, - } - ) - md.append( - f"| {scores[i]} | {scores[j]} | {d.delta:+.3f} | {d.z:.2f} | {d.p_value:.3g} |" - ) - pd.DataFrame(delong_rows).to_csv(out_dir / "delong.csv", index=False) - - strata_rows = [] - for col in [str(c) for c in (cfg.get("strata") or [])]: - md += [ - "", - f"## AUROC by {col} (first score: {scores[0]})", - "", - "| level | n | positives | AUROC [95% CI] |", - "|---|---|---|---|", - ] - for level, r in stratified_auroc( - y, std_scores[scores[0]], df[col].to_numpy(), n_boot=n_boot, seed=seed - ).items(): - strata_rows.append({"stratum": col, "level": level, **r}) - md.append( - f"| {level} | {r['n']} | {r['n_pos']} | {r['auroc']:.3f} [{r['lower']:.3f}, {r['upper']:.3f}] |" - ) - if strata_rows: - pd.DataFrame(strata_rows).to_csv(out_dir / "strata.csv", index=False) - - op_rows = [] - rules = [float(q) for q in cfg.operating_rules] - fold_key = folds if folds is not None else np.zeros(len(y), dtype=int) - for s in scores: - sc = std_scores[s] - md += [ - "", - f"## Operating points: {s}", - "", - "| rule | flagged | sensitivity | specificity | PPV | NPV | LR+ | LR− | TP/FP/FN/TN |", - "|---|---|---|---|---|---|---|---|---|", - ] - points = [ - (f"top {100 * q:.0f} % within fold", flag_top_within_fold(sc, fold_key, q), False) - for q in rules - ] - thr = threshold_points(sc, y) - for name, key in ( - ("Youden (all data, optimistic)", "youden"), - ("rule-out: sens ≥ 90 % (optimistic)", "rule_out"), - ("rule-in: spec ≥ 90 % (optimistic)", "rule_in"), - ): - t = thr[key] - if t == t: - points.append((name, (sc >= t).astype(int), True)) - for name, pred, optimistic in points: - m = operating_point(pred, y) - ci = operating_point_bootstrap(pred, y, n_boot=n_boot, seed=seed) - lr = likelihood_ratios(m["sensitivity"], m["specificity"]) - op_rows.append( - { - "score": s, - "rule": name, - "optimistic": optimistic, - **m, - **{f"{k}_lower": v[0] for k, v in ci.items()}, - **{f"{k}_upper": v[1] for k, v in ci.items()}, - **lr, - } - ) - md.append( - f"| {name} | {100 * m['flagged']:.0f} % | {100 * m['sensitivity']:.0f} % [{100 * ci['sensitivity'][0]:.0f}, {100 * ci['sensitivity'][1]:.0f}] | " - f"{100 * m['specificity']:.0f} % [{100 * ci['specificity'][0]:.0f}, {100 * ci['specificity'][1]:.0f}] | " - f"{100 * m['ppv']:.0f} % [{100 * ci['ppv'][0]:.0f}, {100 * ci['ppv'][1]:.0f}] | {100 * m['npv']:.0f} % [{100 * ci['npv'][0]:.0f}, {100 * ci['npv'][1]:.0f}] | " - f"{lr['lr_plus']:.2f} | {lr['lr_minus']:.2f} | {m['tp']}/{m['fp']}/{m['fn']}/{m['tn']} |" - ) - pd.DataFrame(op_rows).to_csv(out_dir / "operating_points.csv", index=False) - - prev_rows = [] - if rules: - s = scores[0] - q = rules[len(rules) // 2] - m = operating_point(flag_top_within_fold(std_scores[s], fold_key, q), y) - md += [ - "", - f"## PPV / NPV at other prevalences ({s}, top {100 * q:.0f} % rule: sensitivity {100 * m['sensitivity']:.0f} %, specificity {100 * m['specificity']:.0f} %)", - "", - "| prevalence | PPV | NPV |", - "|---|---|---|", - ] - for p in [float(v) for v in cfg.prevalences]: - r = ppv_npv_at_prevalence(m["sensitivity"], m["specificity"], p) - prev_rows.append({"score": s, "rule_q": q, **r}) - md.append(f"| {100 * p:.0f} % | {100 * r['ppv']:.0f} % | {100 * r['npv']:.0f} % |") - pd.DataFrame(prev_rows).to_csv(out_dir / "ppv_npv_by_prevalence.csv", index=False) - - report = out_dir / str(cfg.report_filename) - report.write_text("\n".join(md) + "\n") - summary = { - "n": int(len(y)), - "n_pos": int(y.sum()), - "scores": summary_rows, - "delong": delong_rows, - "report": str(report), - } - (out_dir / "scores_summary.json").write_text(json.dumps(summary, indent=1, default=float)) - log.info("\n".join(md[:12])) - return summary - - -@hydra.main(version_base="1.3", config_path="../configs", config_name="scores.yaml") -def main(cfg: DictConfig) -> None: - run_scores(cfg) - - -if __name__ == "__main__": - main() diff --git a/src/grass_mil/evaluation/operating_points.py b/src/grass_mil/evaluation/operating_points.py deleted file mode 100644 index 8c3034b..0000000 --- a/src/grass_mil/evaluation/operating_points.py +++ /dev/null @@ -1,136 +0,0 @@ -"""Sensitivity / specificity / PPV / NPV at pre-specified operating rules. - -Thresholds are quantile rules applied *within fold* (flag the top fraction q -of units by score), so they do not depend on the fold models' score scales. -Threshold-free summaries (Youden, rule-in, rule-out) are chosen on all data -and reported as optimistic. -""" - -from __future__ import annotations - -from dataclasses import asdict, dataclass -from typing import Dict, List - -import numpy as np - -__all__ = [ - "Confusion", - "confusion_counts", - "flag_top_within_fold", - "likelihood_ratios", - "metrics_from_confusion", - "operating_point", - "operating_point_bootstrap", - "ppv_npv_at_prevalence", - "threshold_points", -] - - -@dataclass -class Confusion: - tp: int - fp: int - fn: int - tn: int - - -def confusion_counts(pred, y) -> Confusion: - pred = np.asarray(pred).astype(int).reshape(-1) - y = np.asarray(y).astype(int).reshape(-1) - return Confusion( - tp=int(((pred == 1) & (y == 1)).sum()), - fp=int(((pred == 1) & (y == 0)).sum()), - fn=int(((pred == 0) & (y == 1)).sum()), - tn=int(((pred == 0) & (y == 0)).sum()), - ) - - -def metrics_from_confusion(c: Confusion) -> Dict[str, float]: - n = c.tp + c.fp + c.fn + c.tn - return { - "sensitivity": c.tp / (c.tp + c.fn) if (c.tp + c.fn) else float("nan"), - "specificity": c.tn / (c.tn + c.fp) if (c.tn + c.fp) else float("nan"), - "ppv": c.tp / (c.tp + c.fp) if (c.tp + c.fp) else float("nan"), - "npv": c.tn / (c.tn + c.fn) if (c.tn + c.fn) else float("nan"), - "flagged": (c.tp + c.fp) / n if n else float("nan"), - "accuracy": (c.tp + c.tn) / n if n else float("nan"), - **asdict(c), - } - - -def flag_top_within_fold(score, folds, q: float) -> np.ndarray: - """Flag the top fraction ``q`` of units within each fold (threshold = fold quantile).""" - s = np.asarray(score, dtype=float).reshape(-1) - f = np.asarray(folds).reshape(-1) - pred = np.zeros(s.size, dtype=int) - for k in np.unique(f): - m = f == k - thr = np.quantile(s[m], 1.0 - float(q)) - pred[m] = (s[m] >= thr).astype(int) - return pred - - -def operating_point(pred, y) -> Dict[str, float]: - return metrics_from_confusion(confusion_counts(pred, y)) - - -def operating_point_bootstrap( - pred, y, *, n_boot: int = 2000, seed: int = 0, alpha: float = 0.05 -) -> Dict[str, List[float]]: - """Percentile bootstrap CIs over units for sensitivity, specificity, PPV, NPV.""" - pred = np.asarray(pred).astype(int).reshape(-1) - y = np.asarray(y).astype(int).reshape(-1) - rng = np.random.default_rng(seed) - keys = ("sensitivity", "specificity", "ppv", "npv") - draws = {k: [] for k in keys} - for _ in range(int(n_boot)): - idx = rng.integers(0, y.size, y.size) - m = operating_point(pred[idx], y[idx]) - for k in keys: - if m[k] == m[k]: - draws[k].append(m[k]) - return { - k: [float(v) for v in np.percentile(vals, [100 * alpha / 2, 100 * (1 - alpha / 2)])] - if vals - else [float("nan"), float("nan")] - for k, vals in draws.items() - } - - -def ppv_npv_at_prevalence( - sensitivity: float, specificity: float, prevalence: float -) -> Dict[str, float]: - p = float(prevalence) - ppv_den = sensitivity * p + (1 - specificity) * (1 - p) - npv_den = specificity * (1 - p) + (1 - sensitivity) * p - return { - "prevalence": p, - "ppv": sensitivity * p / ppv_den if ppv_den > 0 else float("nan"), - "npv": specificity * (1 - p) / npv_den if npv_den > 0 else float("nan"), - } - - -def likelihood_ratios(sensitivity: float, specificity: float) -> Dict[str, float]: - return { - "lr_plus": sensitivity / (1 - specificity) if specificity < 1 else float("inf"), - "lr_minus": (1 - sensitivity) / specificity if specificity > 0 else float("inf"), - } - - -def threshold_points( - score, y, *, rule_out_sens: float = 0.9, rule_in_spec: float = 0.9 -) -> Dict[str, float]: - """Youden-optimal, rule-out and rule-in thresholds chosen on ALL data (optimistic).""" - from sklearn.metrics import roc_curve - - s = np.asarray(score, dtype=float).reshape(-1) - yy = np.asarray(y).astype(int).reshape(-1) - fpr, tpr, thr = roc_curve(yy, s) - out: Dict[str, float] = {} - j = int(np.argmax(tpr - fpr)) - out["youden"] = float(thr[j]) - ro = np.where(tpr >= rule_out_sens)[0] - out["rule_out"] = float(thr[ro[0]]) if ro.size else float("nan") - ri = np.where(fpr <= 1 - rule_in_spec)[0] - out["rule_in"] = float(thr[ri[-1]]) if ri.size else float("nan") - return out diff --git a/src/grass_mil/evaluation/stats.py b/src/grass_mil/evaluation/stats.py deleted file mode 100644 index 8fe9782..0000000 --- a/src/grass_mil/evaluation/stats.py +++ /dev/null @@ -1,187 +0,0 @@ -"""AUROC statistics for held-out (out-of-fold) scores. - -Fold models of a cross-validation have different score scales, so pooled -AUROCs are computed on scores standardised within fold (z-score or rank); -per-fold means are reported alongside. The DeLong test uses the Sun & Xu -structural components for two correlated AUROCs on the same labels. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Dict, Tuple - -import numpy as np - -__all__ = [ - "AurocCI", - "DeLongResult", - "auroc", - "bootstrap_auroc_ci", - "delong_test", - "per_fold_auroc", - "rank_within_fold", - "stratified_auroc", - "z_within_fold", -] - - -def _arr(y, s) -> Tuple[np.ndarray, np.ndarray]: - y = np.asarray(y).astype(int).reshape(-1) - s = np.asarray(s, dtype=float).reshape(-1) - if y.shape[0] != s.shape[0]: - raise ValueError("labels and scores differ in length") - return y, s - - -def auroc(y, s) -> float: - """AUROC via the Mann-Whitney statistic (ties count one half); nan if one class.""" - y, s = _arr(y, s) - pos = s[y == 1] - neg = s[y == 0] - if pos.size == 0 or neg.size == 0: - return float("nan") - from scipy.stats import rankdata - - ranks = rankdata(np.concatenate([pos, neg])) - return float( - (ranks[: pos.size].sum() - pos.size * (pos.size + 1) / 2.0) / (pos.size * neg.size) - ) - - -@dataclass -class AurocCI: - auroc: float - lower: float - upper: float - n_boot: int - n_valid: int - - -def bootstrap_auroc_ci( - y, s, *, n_boot: int = 2000, seed: int = 0, alpha: float = 0.05, groups=None -) -> AurocCI: - """Percentile bootstrap CI over samples (or over ``groups``, e.g. patients).""" - y, s = _arr(y, s) - rng = np.random.default_rng(seed) - point = auroc(y, s) - vals = [] - if groups is not None: - groups = np.asarray(groups) - uniq = np.unique(groups) - members = {g: np.where(groups == g)[0] for g in uniq} - for _ in range(int(n_boot)): - pick = rng.integers(0, uniq.size, uniq.size) - idx = np.concatenate([members[uniq[i]] for i in pick]) - if 0 < y[idx].sum() < idx.size: - vals.append(auroc(y[idx], s[idx])) - else: - n = y.size - for _ in range(int(n_boot)): - idx = rng.integers(0, n, n) - if 0 < y[idx].sum() < n: - vals.append(auroc(y[idx], s[idx])) - if not vals: - return AurocCI(point, float("nan"), float("nan"), int(n_boot), 0) - lo, hi = np.percentile(vals, [100 * alpha / 2, 100 * (1 - alpha / 2)]) - return AurocCI(point, float(lo), float(hi), int(n_boot), len(vals)) - - -@dataclass -class DeLongResult: - auroc_1: float - auroc_2: float - delta: float - z: float - p_value: float - variance: float - - -def _structural_components(s: np.ndarray, pos: np.ndarray, neg: np.ndarray): - sp, sn = s[pos], s[neg] - v10 = np.array([(np.sum(sn < x) + 0.5 * np.sum(sn == x)) / sn.size for x in sp]) - v01 = np.array([(np.sum(sp > x) + 0.5 * np.sum(sp == x)) / sp.size for x in sn]) - return v10, v01 - - -def delong_test(y, s1, s2) -> DeLongResult: - """DeLong test for two correlated AUROCs on the same labels (Sun & Xu components).""" - from scipy import stats - - y, s1 = _arr(y, s1) - _, s2 = _arr(y, s2) - pos = np.where(y == 1)[0] - neg = np.where(y == 0)[0] - if pos.size < 2 or neg.size < 2: - raise ValueError("DeLong needs at least two positives and two negatives.") - a10, a01 = _structural_components(s1, pos, neg) - b10, b01 = _structural_components(s2, pos, neg) - auc1, auc2 = float(a10.mean()), float(b10.mean()) - s10 = np.cov(np.stack([a10, b10])) - s01 = np.cov(np.stack([a01, b01])) - S = s10 / pos.size + s01 / neg.size - var = float(S[0, 0] + S[1, 1] - 2 * S[0, 1]) - z = (auc1 - auc2) / np.sqrt(var) if var > 0 else 0.0 - p = float(2 * stats.norm.sf(abs(z))) - return DeLongResult(auc1, auc2, auc1 - auc2, float(z), p, var) - - -def z_within_fold(scores, folds) -> np.ndarray: - """Standardise scores to zero mean / unit sd within each fold.""" - s = np.asarray(scores, dtype=float).reshape(-1) - f = np.asarray(folds).reshape(-1) - out = np.zeros_like(s) - for k in np.unique(f): - m = f == k - sd = s[m].std() - out[m] = (s[m] - s[m].mean()) / (sd if sd > 0 else 1.0) - return out - - -def rank_within_fold(scores, folds) -> np.ndarray: - """Map scores to (rank - 0.5) / n within each fold (uniform on (0, 1)).""" - from scipy.stats import rankdata - - s = np.asarray(scores, dtype=float).reshape(-1) - f = np.asarray(folds).reshape(-1) - out = np.zeros_like(s) - for k in np.unique(f): - m = f == k - out[m] = (rankdata(s[m]) - 0.5) / m.sum() - return out - - -def per_fold_auroc(y, s, folds) -> Dict[str, float]: - """AUROC within each fold plus their mean and sd (the preprint's headline number).""" - y, s = _arr(y, s) - f = np.asarray(folds).reshape(-1) - vals = {} - for k in sorted(np.unique(f), key=str): - m = f == k - vals[str(k)] = auroc(y[m], s[m]) - arr = np.array([v for v in vals.values() if v == v]) - out = {f"fold_{k}": v for k, v in vals.items()} - out["mean"] = float(arr.mean()) if arr.size else float("nan") - out["sd"] = float(arr.std(ddof=1)) if arr.size > 1 else 0.0 - out["n_folds"] = int(arr.size) - return out - - -def stratified_auroc( - y, s, strata, *, n_boot: int = 2000, seed: int = 0 -) -> Dict[str, Dict[str, float]]: - """AUROC (with bootstrap CI) within each level of a stratifying variable.""" - y, s = _arr(y, s) - st = np.asarray(strata).reshape(-1) - out: Dict[str, Dict[str, float]] = {} - for level in sorted(np.unique(st), key=str): - m = st == level - ci = bootstrap_auroc_ci(y[m], s[m], n_boot=n_boot, seed=seed) - out[str(level)] = { - "n": int(m.sum()), - "n_pos": int(y[m].sum()), - "auroc": ci.auroc, - "lower": ci.lower, - "upper": ci.upper, - } - return out diff --git a/src/grass_mil/extract.py b/src/grass_mil/extract.py index 18b0811..dc36abd 100644 --- a/src/grass_mil/extract.py +++ b/src/grass_mil/extract.py @@ -32,7 +32,9 @@ def resolve_models(ecfg: DictConfig) -> list[dict[str, Any]]: for fold in summary["folds"]: ckpt = Path(fold["run_dir"]) / "checkpoints" / "last.ckpt" held_out = oof.loc[oof["outer"] == fold["outer"], "unit_id"].astype(str).tolist() - models.append({"name": f"fold{fold['outer']}", "ckpt_path": str(ckpt), "held_out": held_out}) + models.append( + {"name": f"fold{fold['outer']}", "ckpt_path": str(ckpt), "held_out": held_out} + ) for name, ckpt in dict(ecfg.get("models") or {}).items(): models.append({"name": str(name), "ckpt_path": str(ckpt), "held_out": []}) if not models: @@ -61,7 +63,9 @@ def run_extract(cfg: DictConfig) -> list[str]: ckpt = torch.load(spec["ckpt_path"], map_location="cpu", weights_only=False) model = SupervisedModule.from_checkpoint(spec["ckpt_path"]) if tuple(model.input_spec.cell_types) != tuple(datamodule.input_spec.cell_types): - raise ValueError(f"{spec['ckpt_path']} was trained on other cell types than this cohort.") + raise ValueError( + f"{spec['ckpt_path']} was trained on other cell types than this cohort." + ) model_dir = extract_cells( model, units, diff --git a/src/grass_mil/inference/cell_cache.py b/src/grass_mil/inference/cell_cache.py index 521cff0..443901c 100644 --- a/src/grass_mil/inference/cell_cache.py +++ b/src/grass_mil/inference/cell_cache.py @@ -95,7 +95,9 @@ def __iter__(self) -> Iterator[tuple[str, dict[str, np.ndarray]]]: def sampler_of_checkpoint(ckpt: dict[str, Any], batch_size: int) -> SamplerConfig: """Ego-graph settings the model was trained with, scoring every cell.""" cfg = dict(ckpt["datamodule_hyper_parameters"]["sampler"]) - cfg.update(roots_per_unit=None, batch_size=int(batch_size), merge_units=False, root_weighting=None) + cfg.update( + roots_per_unit=None, batch_size=int(batch_size), merge_units=False, root_weighting=None + ) cfg["interior_roots"] = {**dict(cfg.get("interior_roots") or {}), "enabled": False} return SamplerConfig.from_dict(cfg) @@ -136,7 +138,12 @@ def extract_cells( for i, entry in enumerate(dataset.entries): if wanted is not None and entry.unit_id not in wanted: continue - info = {"unit_id": entry.unit_id, "region_id": entry.region_id, "sample_id": entry.sample_id, "n_cells": entry.n_cells} + info = { + "unit_id": entry.unit_id, + "region_id": entry.region_id, + "sample_id": entry.sample_id, + "n_cells": entry.n_cells, + } info.update((unit_meta or {}).get(entry.unit_id, {})) index_units.append(info) path = _npz_path(model_dir, entry.unit_id) @@ -146,7 +153,9 @@ def extract_cells( n = int(data.num_nodes) arrays: dict[str, np.ndarray] = {} seen = np.zeros(n, dtype=bool) - for batch in EgoGraphLoader(_OneUnit(data, entry), sampler, mode="chunks", roots_per_unit=None): + for batch in EgoGraphLoader( + _OneUnit(data, entry), sampler, mode="chunks", roots_per_unit=None + ): roots = batch.root_cell.numpy() batch = batch.to(device) with torch.autocast(device.type, dtype=torch.float16, enabled=autocast): @@ -164,7 +173,9 @@ def extract_cells( arrays[key][roots] = value.float().cpu().numpy() seen[roots] = True if not seen.all(): - raise RuntimeError(f"Unit {entry.unit_id}: {int((~seen).sum())} cells were not scored.") + raise RuntimeError( + f"Unit {entry.unit_id}: {int((~seen).sum())} cells were not scored." + ) arrays["pos"] = data.pos.numpy().astype(np.float32) arrays["edge_index"] = data.edge_index.numpy().astype(np.int64) arrays["cell_ids"] = np.asarray(dataset.load(i)["cell_ids"], dtype=str) diff --git a/src/grass_mil/inference/predictions.py b/src/grass_mil/inference/predictions.py index a91318b..6f022a3 100644 --- a/src/grass_mil/inference/predictions.py +++ b/src/grass_mil/inference/predictions.py @@ -84,8 +84,14 @@ def frame_metrics(frame: pd.DataFrame, target_type: str) -> dict[str, float]: logits = torch.log(torch.as_tensor(probs, dtype=torch.float32).clamp_min(1e-12)) targets = torch.as_tensor(frame[["label"]].to_numpy(), dtype=torch.float32) elif target_type == "regression": - logits = torch.as_tensor(frame[[c for c in frame.columns if c.startswith("pred_")]].to_numpy(), dtype=torch.float32) - targets = torch.as_tensor(frame[[c for c in frame.columns if c.startswith("target_")]].to_numpy(), dtype=torch.float32) + logits = torch.as_tensor( + frame[[c for c in frame.columns if c.startswith("pred_")]].to_numpy(), + dtype=torch.float32, + ) + targets = torch.as_tensor( + frame[[c for c in frame.columns if c.startswith("target_")]].to_numpy(), + dtype=torch.float32, + ) else: logits = torch.as_tensor(frame[["risk"]].to_numpy(), dtype=torch.float32) targets = torch.as_tensor(frame[["time", "event"]].to_numpy(), dtype=torch.float32) diff --git a/src/grass_mil/interpretability/__init__.py b/src/grass_mil/interpretability/__init__.py deleted file mode 100644 index 1f3d041..0000000 --- a/src/grass_mil/interpretability/__init__.py +++ /dev/null @@ -1,51 +0,0 @@ -from grass_mil.interpretability.core.biomarkers import ( - niche_composition_summary, -) -from grass_mil.interpretability.core.clustering import ( - cosine_hdbscan, - fit_niche_clustering, - run_clustering, -) -from grass_mil.interpretability.core.attribution import niche_attribution_summary -from grass_mil.interpretability.core.agreement import compute_niche_agreement -from grass_mil.interpretability.niches import run_niche_discovery -from grass_mil.interpretability.core.reduction import run_reduction -from grass_mil.interpretability.core.transfer import ( - apply_niche_transfer, - fit_niche_transfer_from_report_bundle, - load_niche_transfer_bundle, - save_niche_transfer_bundle, -) -from grass_mil.interpretability.pipeline import run_interpretability_pipeline -from grass_mil.interpretability.reporting.render import render_interpretability_report -from grass_mil.interpretability.tier2.filtration import compute_filtration_curves -from grass_mil.interpretability.tier2.neighborhood import ( - run_diff_neighborhood_enrichment, - run_neighborhood_enrichment, -) -from grass_mil.interpretability.tier2.tissue_graph import ( - build_tissue_graph_figure, - prepare_tissue_graph_view, -) - -__all__ = [ - "run_reduction", - "run_clustering", - "cosine_hdbscan", - "fit_niche_clustering", - "run_niche_discovery", - "niche_attribution_summary", - "compute_niche_agreement", - "niche_composition_summary", - "run_neighborhood_enrichment", - "run_diff_neighborhood_enrichment", - "compute_filtration_curves", - "prepare_tissue_graph_view", - "build_tissue_graph_figure", - "fit_niche_transfer_from_report_bundle", - "apply_niche_transfer", - "save_niche_transfer_bundle", - "load_niche_transfer_bundle", - "run_interpretability_pipeline", - "render_interpretability_report", -] diff --git a/src/grass_mil/interpretability/core/__init__.py b/src/grass_mil/interpretability/core/__init__.py deleted file mode 100644 index 2bf512d..0000000 --- a/src/grass_mil/interpretability/core/__init__.py +++ /dev/null @@ -1,27 +0,0 @@ -from grass_mil.interpretability.core.biomarkers import ( - niche_composition_summary, -) -from grass_mil.interpretability.core.clustering import run_clustering -from grass_mil.interpretability.core.data import ( - extract_embedding_set, - load_interpretability_dataset, -) -from grass_mil.interpretability.core.reduction import run_reduction -from grass_mil.interpretability.core.transfer import ( - apply_niche_transfer, - fit_niche_transfer_from_report_bundle, - load_niche_transfer_bundle, - save_niche_transfer_bundle, -) - -__all__ = [ - "run_reduction", - "run_clustering", - "niche_composition_summary", - "load_interpretability_dataset", - "extract_embedding_set", - "fit_niche_transfer_from_report_bundle", - "apply_niche_transfer", - "save_niche_transfer_bundle", - "load_niche_transfer_bundle", -] diff --git a/src/grass_mil/interpretability/core/agreement.py b/src/grass_mil/interpretability/core/agreement.py deleted file mode 100644 index bbb030b..0000000 --- a/src/grass_mil/interpretability/core/agreement.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Agreement between partitions of the same instances. - -The suite niches several representation spaces -- the encoder embedding, the -attention head space, and a model-independent cell-type composition space. -Structure that appears in a learned space but not in the composition baseline is -attributable to the model rather than to raw abundance, so comparing partitions -is how that claim is checked. - -All indices are chance-corrected or normalised, and computed over the same -instances in the same order. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from itertools import combinations -from typing import Dict, Mapping, Sequence, Tuple - -import numpy as np -import pandas as pd - -__all__ = ["NicheAgreementResult", "niche_jaccard_matrix", "compute_niche_agreement"] - -_METRICS = ("ari", "ami", "nmi") - - -@dataclass(frozen=True) -class NicheAgreementResult: - """Pairwise agreement between labelings. - - Attributes: - pairwise_metrics: One row per space pair, columns ``ari``/``ami``/``nmi``. - jaccard: ``(space_a, space_b) -> niche x niche`` Jaccard overlap. - """ - - pairwise_metrics: pd.DataFrame - jaccard: Dict[Tuple[str, str], pd.DataFrame] - - -def niche_jaccard_matrix(labels_a: Sequence, labels_b: Sequence) -> pd.DataFrame: - """Jaccard overlap between every pair of niches from two labelings.""" - a = np.asarray(labels_a) - b = np.asarray(labels_b) - if a.shape != b.shape: - raise ValueError(f"Labelings must align; got {a.shape} and {b.shape}.") - - clusters_a = sorted(set(a.tolist())) - clusters_b = sorted(set(b.tolist())) - matrix = np.zeros((len(clusters_a), len(clusters_b)), dtype=float) - for i, ca in enumerate(clusters_a): - mask_a = a == ca - for j, cb in enumerate(clusters_b): - mask_b = b == cb - union = int(np.sum(mask_a | mask_b)) - matrix[i, j] = float(np.sum(mask_a & mask_b) / union) if union else np.nan - return pd.DataFrame(matrix, index=clusters_a, columns=clusters_b) - - -def compute_niche_agreement( - labelings: Mapping[str, Sequence], - *, - metrics: Sequence[str] = _METRICS, - include_jaccard: bool = True, -) -> NicheAgreementResult: - """Chance-corrected agreement between every pair of labelings.""" - from sklearn.metrics import ( - adjusted_mutual_info_score, - adjusted_rand_score, - normalized_mutual_info_score, - ) - - unknown = set(metrics) - set(_METRICS) - if unknown: - raise ValueError( - f"Unsupported agreement metric(s): {sorted(unknown)}. Expected a subset of {_METRICS}." - ) - if len(labelings) < 2: - raise ValueError("Niche agreement needs at least two labelings to compare.") - - arrays = {name: np.asarray(values) for name, values in labelings.items()} - sizes = {name: values.shape[0] for name, values in arrays.items()} - if len(set(sizes.values())) != 1: - raise ValueError(f"All labelings must cover the same instances; got sizes {sizes}.") - - functions = { - "ari": adjusted_rand_score, - "ami": adjusted_mutual_info_score, - "nmi": normalized_mutual_info_score, - } - - rows = [] - jaccard: Dict[Tuple[str, str], pd.DataFrame] = {} - for name_a, name_b in combinations(sorted(arrays), 2): - a, b = arrays[name_a], arrays[name_b] - row: Dict[str, object] = {"space_a": name_a, "space_b": name_b} - for metric in metrics: - row[metric] = float(functions[metric](a, b)) - rows.append(row) - if include_jaccard: - jaccard[(name_a, name_b)] = niche_jaccard_matrix(a, b) - - frame = pd.DataFrame(rows).set_index(["space_a", "space_b"]) - return NicheAgreementResult(pairwise_metrics=frame, jaccard=jaccard) diff --git a/src/grass_mil/interpretability/core/attribution.py b/src/grass_mil/interpretability/core/attribution.py deleted file mode 100644 index 7566099..0000000 --- a/src/grass_mil/interpretability/core/attribution.py +++ /dev/null @@ -1,347 +0,0 @@ -"""Exact attribution of a bag prediction to its instances. - -Because the bag logit is an attention-weighted sum of instance logits with no -added bias, - -.. math:: L_c = \\sum_{i \\in B_r} A_{i,c}\\, \\ell_{i,c}, - -each instance carries an exact additive contribution -:math:`M_{i,c} = A_{i,c} \\ell_{i,c}`. This is an identity, not an -approximation: it needs no baseline, no integration path, and no surrogate -model. :func:`niche_attribution_summary` reports ``identity_residual`` so the -assumption is checked rather than trusted. - -The reported margin removes the head's per-class bias, so it reflects the -instance-driven part of the decision rather than the head's prior. It is a -one-vs-rest contrast, which for a binary head is exactly - -.. math:: m_i = A_{i,1}(\\ell_{i,1} - \\beta_1) - A_{i,0}(\\ell_{i,0} - \\beta_0) - -and generalises to any number of classes (see :func:`instance_ovr_margins`). - -Niche-level summaries aggregate within niche x region cells and are -reported region-equal with percentile bootstrap intervals over regions, because -regions differ enormously in instance count and pooling over instances would -let one large region dominate. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Dict, List, Optional, Sequence, Tuple - -import numpy as np -import pandas as pd - -__all__ = [ - "AttributionResult", - "niche_attribution_summary", - "instance_margin_contributions", - "instance_ovr_margins", - "percentile_bootstrap_ci", -] - - -@dataclass(frozen=True) -class AttributionResult: - """Per-instance contributions and their niche-level summaries. - - Attributes: - per_instance: One row per instance with its contribution columns. - per_niche: One row per niche; each statistic carries ``_lo``/``_hi`` - percentile bootstrap bounds. - identity_residual: ``max_r |sum_i M[i,c] - L_c|``. Non-zero means the - additive identity does not hold and the attribution is invalid. - """ - - per_instance: pd.DataFrame - per_niche: pd.DataFrame - identity_residual: float - - -def percentile_bootstrap_ci( - values: np.ndarray, - *, - weights: Optional[np.ndarray] = None, - n_boot: int = 200, - q_lo: float = 2.5, - q_hi: float = 97.5, - rng: Optional[np.random.Generator] = None, -) -> Tuple[float, float, float]: - """Point estimate and percentile bootstrap interval over resampled units. - - Returns ``(point, lo, hi)``; all three are NaN for an empty input. - """ - values = np.asarray(values, dtype=float) - finite = np.isfinite(values) - if weights is not None: - weights = np.asarray(weights, dtype=float) - finite &= np.isfinite(weights) - values = values[finite] - if weights is not None: - weights = weights[finite] - if values.size == 0: - return float("nan"), float("nan"), float("nan") - - def _estimate(sample_values: np.ndarray, sample_weights: Optional[np.ndarray]) -> float: - if sample_weights is None: - return float(np.mean(sample_values)) - total = float(np.sum(np.abs(sample_weights))) - if total <= 0: - return float("nan") - return float(np.sum(sample_weights * sample_values) / total) - - point = _estimate(values, weights) - if values.size == 1 or n_boot <= 0: - return point, point, point - - rng = rng or np.random.default_rng(0) - draws = rng.integers(0, values.size, size=(int(n_boot), values.size)) - samples = np.array( - [_estimate(values[idx], None if weights is None else weights[idx]) for idx in draws], - dtype=float, - ) - samples = samples[np.isfinite(samples)] - if samples.size == 0: - return point, float("nan"), float("nan") - return point, float(np.percentile(samples, q_lo)), float(np.percentile(samples, q_hi)) - - -def instance_margin_contributions( - attention: np.ndarray, - logits: np.ndarray, - *, - logit_bias: Optional[Sequence[float]] = None, -) -> np.ndarray: - """``M[i,c] = A[i,c] * (l[i,c] - beta_c)``. - - ``logit_bias`` is the head's per-class output bias. Removing it isolates the - instance-driven component; pass ``None`` to keep the raw identity. - """ - attention = np.asarray(attention, dtype=float) - logits = np.asarray(logits, dtype=float) - if attention.ndim == 1: - attention = attention[:, None] - if logits.ndim == 1: - logits = logits[:, None] - if attention.shape[0] != logits.shape[0]: - raise ValueError( - "attention and logits must have the same number of rows; got " - f"{attention.shape[0]} and {logits.shape[0]}." - ) - if attention.shape[1] not in (1, logits.shape[1]): - raise ValueError( - f"attention has {attention.shape[1]} channel(s), which is neither 1 nor the " - f"{logits.shape[1]} class logit(s)." - ) - - centred = logits - if logit_bias is not None: - bias = np.asarray(logit_bias, dtype=float).reshape(1, -1) - if bias.shape[1] != logits.shape[1]: - raise ValueError( - f"logit_bias has {bias.shape[1]} entries but there are {logits.shape[1]} classes." - ) - centred = logits - bias - return attention * centred - - -def instance_ovr_margins(contributions: np.ndarray) -> np.ndarray: - """One-vs-rest margin per instance and class. - - For class :math:`c`, the margin contrasts that class's contribution against - the mean of the others: - - .. math:: m_{i,c} = M_{i,c} - \\frac{1}{C-1} \\sum_{c' \\neq c} M_{i,c'} - - With ``C == 2`` this reduces exactly to :math:`M_{i,1} - M_{i,0}`, the - classic binary margin, so binary results are unchanged. With ``C == 1`` - (regression, or Cox on the log-hazard scale) the contribution is already the - margin. - - Returns: - ``[n_instances, n_classes]``. - """ - contributions = np.asarray(contributions, dtype=float) - if contributions.ndim == 1: - contributions = contributions[:, None] - n_classes = contributions.shape[1] - if n_classes == 1: - return contributions - total = contributions.sum(axis=1, keepdims=True) - rest_mean = (total - contributions) / (n_classes - 1) - return contributions - rest_mean - - -def _resolve_focus_classes(n_classes: int, focus_classes) -> List[int]: - """Which classes get a niche-level summary. - - Binary tasks summarise the positive class only: the class-0 margin is its - exact negation, so reporting both is redundant. Multi-class tasks summarise - every class, since no single contrast represents the decision. - """ - if focus_classes is not None: - resolved = [int(c) for c in focus_classes] - for c in resolved: - if c < 0 or c >= n_classes: - raise ValueError(f"focus class {c} is out of range for {n_classes} class(es).") - return resolved - if n_classes <= 2: - return [n_classes - 1] - return list(range(n_classes)) - - -def niche_attribution_summary( - table: pd.DataFrame, - niche_labels: np.ndarray, - *, - attention_columns: Sequence[str], - logit_columns: Sequence[str], - bag_id_column: str = "bag_id", - logit_bias: Optional[Sequence[float]] = None, - n_bootstrap: int = 200, - random_state: int = 0, - margin_eps: float = 1e-6, - bag_logits: Optional[Dict[str, Sequence[float]]] = None, - focus_classes: Optional[Sequence[int]] = None, -) -> AttributionResult: - """Summarise each niche's contribution to the bag decisions. - - Args: - table: Instance table; must carry the attention, logit and bag columns. - niche_labels: One niche id per row of ``table``. - attention_columns: Per-class attention columns, or a single shared one. - logit_columns: Per-class instance logit columns. - bag_id_column: Column grouping instances into bags. - logit_bias: Head output bias per class, removed from the margin. - bag_logits: Optional bag logits keyed by bag id. When given, the - additive identity is checked against them. - focus_classes: Classes to summarise. Defaults to the positive class for - binary tasks (the class-0 margin is its exact negation) and to every - class for multi-class tasks. - """ - for column in (*attention_columns, *logit_columns, bag_id_column): - if column not in table.columns: - raise ValueError(f"Missing required column {column!r} for attribution.") - if len(niche_labels) != len(table): - raise ValueError(f"niche_labels has {len(niche_labels)} entries for {len(table)} rows.") - - attention = table.loc[:, list(attention_columns)].to_numpy(dtype=float) - logits = table.loc[:, list(logit_columns)].to_numpy(dtype=float) - contributions = instance_margin_contributions(attention, logits, logit_bias=logit_bias) - raw_contributions = instance_margin_contributions(attention, logits, logit_bias=None) - - per_instance = pd.DataFrame( - { - bag_id_column: table[bag_id_column].astype(str).to_numpy(), - "niche_label": np.asarray(niche_labels), - } - ) - margins = instance_ovr_margins(contributions) - n_classes = int(contributions.shape[1]) - for index in range(n_classes): - per_instance[f"contribution_c{index}"] = contributions[:, index] - per_instance[f"margin_c{index}"] = margins[:, index] - - identity_residual = _identity_residual( - per_instance[bag_id_column].to_numpy(), raw_contributions, bag_logits - ) - - summaries = [] - for class_index in _resolve_focus_classes(n_classes, focus_classes): - frame = per_instance.copy() - frame["margin"] = margins[:, class_index] - part = _summarise_niches( - frame, - bag_id_column=bag_id_column, - n_bootstrap=n_bootstrap, - random_state=random_state, - margin_eps=margin_eps, - ) - part["class_index"] = class_index - summaries.append(part.set_index("class_index", append=True)) - summary = pd.concat(summaries).sort_index() - return AttributionResult( - per_instance=per_instance, - per_niche=summary, - identity_residual=identity_residual, - ) - - -def _identity_residual( - bag_ids: np.ndarray, - contributions: np.ndarray, - bag_logits: Optional[Dict[str, Sequence[float]]], -) -> float: - """``max_r |sum_i M[i,c] - L_c|``, or NaN when bag logits are unavailable.""" - if not bag_logits: - return float("nan") - residual = 0.0 - frame = pd.DataFrame(contributions) - frame["__bag"] = bag_ids - for bag_id, group in frame.groupby("__bag"): - expected = bag_logits.get(str(bag_id)) - if expected is None: - continue - got = group.drop(columns="__bag").to_numpy(dtype=float).sum(axis=0) - expected_arr = np.asarray(expected, dtype=float).reshape(-1) - width = min(got.size, expected_arr.size) - residual = max(residual, float(np.max(np.abs(got[:width] - expected_arr[:width])))) - return residual - - -def _summarise_niches( - per_instance: pd.DataFrame, - *, - bag_id_column: str, - n_bootstrap: int, - random_state: int, - margin_eps: float, -) -> pd.DataFrame: - frame = per_instance.copy() - bag_margin = frame.groupby(bag_id_column)["margin"].sum() - bag_size = frame.groupby(bag_id_column).size() - - cell_margin = frame.groupby([bag_id_column, "niche_label"])["margin"].sum() - cell_count = frame.groupby([bag_id_column, "niche_label"]).size() - - rng = np.random.default_rng(random_state) - rows = [] - for niche in sorted(frame["niche_label"].unique()): - niche_margin = cell_margin.xs(niche, level="niche_label") - niche_count = cell_count.xs(niche, level="niche_label") - bags = niche_margin.index - - totals = bag_margin.loc[bags].to_numpy(dtype=float) - margins = niche_margin.to_numpy(dtype=float) - counts = niche_count.to_numpy(dtype=float) - sizes = bag_size.loc[bags].to_numpy(dtype=float) - - prevalence = counts / np.maximum(sizes, 1.0) - - # Shares are fractions of the bag margin and sum to 1 across a full - # partition, but an individual niche is NOT bounded to [-1, 1]: it can - # exceed 1 when another niche pushes the opposite way. Regions whose - # margin is near zero are excluded rather than clipped, since the ratio - # is meaningless there. - decisive = np.abs(totals) > margin_eps - share = np.divide(margins, totals, out=np.full_like(margins, np.nan), where=decisive) - consistency = np.where(decisive, np.sign(margins) == np.sign(totals), np.nan) - - row: Dict[str, object] = {"niche_label": niche, "n_regions": int(len(bags))} - for name, values, weights in ( - ("margin_signed_share", share, None), - ("margin_weighted", margins, np.abs(totals)), - ("margin_sign_consistency", consistency, None), - ("prevalence", prevalence, None), - ): - point, lo, hi = percentile_bootstrap_ci( - values, weights=weights, n_boot=n_bootstrap, rng=rng - ) - row[name] = point - row[f"{name}_lo"] = lo - row[f"{name}_hi"] = hi - row["n_instances"] = int(counts.sum()) - rows.append(row) - - return pd.DataFrame(rows).set_index("niche_label").sort_index() diff --git a/src/grass_mil/interpretability/core/biomarkers.py b/src/grass_mil/interpretability/core/biomarkers.py deleted file mode 100644 index 11e7a25..0000000 --- a/src/grass_mil/interpretability/core/biomarkers.py +++ /dev/null @@ -1,66 +0,0 @@ -from __future__ import annotations - -import numpy as np -import pandas as pd - -from grass_mil.contracts import NicheSummary - - -def _resolve_variance_estimator(variance_estimator: str) -> int: - key = str(variance_estimator).strip().lower() - if key in {"unbiased", "sample", "n-1"}: - return 1 - if key in {"biased", "population", "n"}: - return 0 - raise ValueError("Unsupported variance_estimator. Expected one of: " "'unbiased' or 'biased'.") - - -def _infer_composition_matrix( - table: pd.DataFrame, - *, - composition_prefix: str, -) -> pd.DataFrame: - comp_cols = [c for c in table.columns if c.startswith(composition_prefix)] - if comp_cols: - return table[comp_cols].astype(float).copy() - - raise ValueError( - f"No composition columns found with prefix {composition_prefix!r}. " - "Explicit composition features are required." - ) - - -def niche_composition_summary( - table: pd.DataFrame, - niche_labels: np.ndarray, - *, - composition_prefix: str = "comp_", - variance_estimator: str = "unbiased", -) -> NicheSummary: - """Mean composition and z-scored enrichment per niche. - - Rows come out in niche-id order, which is dendrogram order: the pipeline - renumbers niches by composition once, right after clustering, so nothing - downstream has to re-derive or thread an ordering. - """ - if len(table) != int(niche_labels.shape[0]): - raise ValueError("table rows and niche_labels length mismatch.") - - comp = _infer_composition_matrix(table, composition_prefix=composition_prefix) - labels = pd.Series(niche_labels, name="niche_label") - - ddof = _resolve_variance_estimator(variance_estimator) - comp_mean = comp.mean(axis=0) - comp_std = comp.std(axis=0, ddof=ddof).replace(0.0, 1.0) - z_comp = (comp - comp_mean) / comp_std - - composition = comp.groupby(labels).mean().sort_index() - enrichment = z_comp.groupby(labels).mean().sort_index() - counts = labels.value_counts().sort_index() - - return NicheSummary( - niche_labels=np.asarray(niche_labels), - composition=composition, - enrichment=enrichment, - niche_counts=counts, - ) diff --git a/src/grass_mil/interpretability/core/clustering.py b/src/grass_mil/interpretability/core/clustering.py deleted file mode 100644 index 5bd9c47..0000000 --- a/src/grass_mil/interpretability/core/clustering.py +++ /dev/null @@ -1,309 +0,0 @@ -from __future__ import annotations - -from typing import Any, Dict, Optional - -import numpy as np -from sklearn.cluster import ( - OPTICS, - AgglomerativeClustering, - DBSCAN, - KMeans, - SpectralClustering, -) - -from grass_mil.contracts import ClusteringResult - - -def run_clustering( - method: str, - X: np.ndarray, - *, - n_clusters: Optional[int] = None, - random_state: Optional[int] = 42, - **params: Any, -) -> ClusteringResult: - """Run a clustering backend with estimator kwargs forwarded from `params`. - - Notes: - - `n_clusters` and `random_state` are accepted as common convenience args. - - Any additional `params` are passed directly to the selected estimator constructor. - """ - method_key = str(method).strip().lower() - if method_key == "kmeans": - if n_clusters is None: - raise ValueError("kmeans requires `n_clusters`.") - model = KMeans(n_clusters=n_clusters, random_state=random_state, **params) - labels = model.fit_predict(X) - elif method_key == "agglomerative": - if n_clusters is None: - raise ValueError("agglomerative requires `n_clusters`.") - model = AgglomerativeClustering(n_clusters=n_clusters, **params) - labels = model.fit_predict(X) - elif method_key == "hdbscan": - try: - from sklearn.cluster import HDBSCAN # type: ignore - except Exception as exc: # pragma: no cover - optional dependency - raise RuntimeError( - "HDBSCAN requested but sklearn HDBSCAN is unavailable in this environment." - ) from exc - model = HDBSCAN(**params) - labels = model.fit_predict(X) - elif method_key == "spectral": - if n_clusters is None: - raise ValueError("spectral requires `n_clusters`.") - model = SpectralClustering(n_clusters=n_clusters, random_state=random_state, **params) - labels = model.fit_predict(X) - elif method_key == "optics": - model = OPTICS(**params) - labels = model.fit_predict(X) - elif method_key == "dbscan": - model = DBSCAN(**params) - labels = model.fit_predict(X) - else: - raise ValueError(f"Unsupported clustering method: {method!r}") - - all_params: Dict[str, Any] = {"n_clusters": n_clusters, "random_state": random_state} - all_params.update(params) - return ClusteringResult( - method=method_key, - labels=np.asarray(labels), - params=all_params, - fitted_object=model, - ) - - -# ---------------------------------------------------------------------------- -# Sample -> reduce -> (L2) -> cluster -> assign: the GRASS-MIL niche recipe -# ---------------------------------------------------------------------------- - - -_HDBSCAN_ONLY = ( - "min_cluster_size", - "min_samples", - "cluster_selection_method", - "cluster_selection_epsilon", - "allow_single_cluster", - "alpha", - "max_cluster_size", - "store_centers", -) - - -class NicheClusterModel: - """A fitted "reduce on a sample, cluster, assign the rest" model. - - ``predict(X)`` reduces new embeddings with the fitted reducer and assigns - them with the k-NN classifier fitted on the clustered sample (HDBSCAN has - no native predict; k-NN on the reduced coordinates is what the preprint - uses). Backends with their own ``predict`` (k-means) use it directly. - """ - - def __init__( - self, - *, - reducer, - clusterer, - assigner, - sample_idx: np.ndarray, - labels_sample: np.ndarray, - params: Dict[str, Any], - l2_normalise: bool, - reducer_method: str, - clustering_method: str, - ) -> None: - self.reducer = reducer - self.clusterer = clusterer - self.assigner = assigner - self.sample_idx = np.asarray(sample_idx) - self.labels_sample = np.asarray(labels_sample) - self.params = dict(params) - self.l2_normalise = bool(l2_normalise) - self.reducer_method = reducer_method - self.clustering_method = clustering_method - - def reduce(self, X: np.ndarray) -> np.ndarray: - X = np.asarray(X, dtype=np.float32) - return X if self.reducer is None else self.reducer.transform(X) - - def predict(self, X: np.ndarray) -> np.ndarray: - P = self.reduce(X) - if self.assigner is not None: - return np.asarray(self.assigner.predict(P)) - from sklearn.preprocessing import normalize - - Pc = normalize(P, norm="l2") if self.l2_normalise else P - return np.asarray(self.clusterer.predict(Pc)) - - @property - def n_clusters(self) -> int: - return int(len(set(self.labels_sample.tolist()) - {-1})) - - -def fit_niche_clustering( - X: np.ndarray, - *, - reduction: Optional[Dict[str, Any]] = None, - clustering: Optional[Dict[str, Any]] = None, - sample_size: Optional[int] = 100_000, - l2_normalise: bool = True, - assign_k: int = 15, - seed: int = 0, -) -> tuple[NicheClusterModel, np.ndarray]: - """Fit the niche clustering and label every row of ``X``. - - Default (preprint): PCA to 10 components fitted on a random sample of at - most ``sample_size`` rows, L2 normalisation (so euclidean HDBSCAN acts on - the cosine geometry), ``HDBSCAN(min_cluster_size=300, min_samples=1, - cluster_selection_method="eom")`` on the sample, and 15-nearest-neighbour - assignment of the remaining rows on the (un-normalised) PCs. - - ``reduction`` = ``{"method": "pca"|"none", "n_components": int, ...}``; - ``clustering`` = ``{"method": , "params": {...}, - "n_clusters": int}``. Any backend of :func:`run_clustering` is accepted. - Returns ``(model, labels_all)`` with ``-1`` for noise. - """ - from sklearn.neighbors import KNeighborsClassifier - from sklearn.preprocessing import normalize - - from grass_mil.interpretability.core.reduction import run_reduction - - X = np.asarray(X, dtype=np.float32) - n = int(X.shape[0]) - reduction = dict(reduction or {"method": "pca", "n_components": 10}) - clustering = dict( - clustering - or { - "method": "hdbscan", - "params": { - "min_cluster_size": 300, - "min_samples": 1, - "cluster_selection_method": "eom", - }, - } - ) - rng = np.random.default_rng(int(seed)) - if sample_size is None or int(sample_size) >= n: - sample_idx = np.arange(n) - else: - sample_idx = np.sort(rng.choice(n, int(sample_size), replace=False)) - - red_method = str(reduction.get("method", "pca")).strip().lower() - reducer = None - if red_method not in {"none", "identity"}: - red_params = { - k: v - for k, v in reduction.items() - if k not in {"method", "n_components", "random_state"} - } - n_components = int(reduction.get("n_components", 10)) - n_components = max(1, min(n_components, int(X.shape[1]), int(sample_idx.size))) - res = run_reduction( - red_method, - X[sample_idx], - n_components=n_components, - random_state=int(reduction.get("random_state", 42)), - **red_params, - ) - reducer = res.fitted_object - P_sample = X[sample_idx] if reducer is None else reducer.transform(X[sample_idx]) - P_cluster = normalize(P_sample, norm="l2") if l2_normalise else P_sample - - cl_method = str(clustering.get("method", "hdbscan")).strip().lower() - cl_params = dict(clustering.get("params", {}) or {}) - if cl_method == "hdbscan" and l2_normalise: - cl_params.pop("metric", None) - if cl_method != "hdbscan": - # The default params block is HDBSCAN's; switching the backend on the - # command line must not forward keys the other estimator does not take. - for key in _HDBSCAN_ONLY: - cl_params.pop(key, None) - if cl_method == "hdbscan": - cl_params.setdefault("min_cluster_size", 300) - cl_params.setdefault("min_samples", 1) - cl_params.setdefault("cluster_selection_method", "eom") - cl_params["min_cluster_size"] = int( - min(int(cl_params["min_cluster_size"]), max(2, sample_idx.size)) - ) - result = run_clustering( - cl_method, - P_cluster, - n_clusters=clustering.get("n_clusters"), - random_state=int(clustering.get("random_state", 42)), - **cl_params, - ) - labels_sample = np.asarray(result.labels).astype(int) - - # Backends with a native predict (k-means, ...) assign unsampled rows - # themselves, exactly reproducing the sample labels; density backends - # (HDBSCAN, DBSCAN, OPTICS) have none, so a k-NN classifier on the reduced - # sample coordinates assigns the rest (the preprint's 15-NN step). - assigner = None - if not hasattr(result.fitted_object, "predict"): - k = max(1, min(int(assign_k) if int(assign_k) > 0 else 15, int(sample_idx.size))) - assigner = KNeighborsClassifier(k).fit(P_sample, labels_sample) - model = NicheClusterModel( - reducer=reducer, - clusterer=result.fitted_object, - assigner=assigner, - sample_idx=sample_idx, - labels_sample=labels_sample, - params={ - "reduction": { - "method": red_method, - **( - {"n_components": int(reducer.n_components_)} - if reducer is not None and hasattr(reducer, "n_components_") - else {} - ), - }, - "clustering": { - "method": cl_method, - **cl_params, - **( - {"n_clusters": clustering.get("n_clusters")} - if clustering.get("n_clusters") - else {} - ), - }, - "sample_size": int(sample_idx.size), - "l2_normalise": bool(l2_normalise), - "assign_k": int(assign_k), - "seed": int(seed), - }, - l2_normalise=l2_normalise, - reducer_method=red_method, - clustering_method=cl_method, - ) - labels_all = model.predict(X) - labels_all[sample_idx] = labels_sample - return model, labels_all - - -def cosine_hdbscan( - X: np.ndarray, - *, - n_components: int = 10, - min_cluster_size: int = 300, - min_samples: Optional[int] = 1, - cluster_selection_method: str = "eom", - sample_size: Optional[int] = 100_000, - assign_k: int = 15, - seed: int = 0, -) -> tuple[NicheClusterModel, np.ndarray]: - """The preprint's niche clustering with its defaults (see :func:`fit_niche_clustering`).""" - return fit_niche_clustering( - X, - reduction={"method": "pca", "n_components": int(n_components)}, - clustering={ - "method": "hdbscan", - "params": { - "min_cluster_size": int(min_cluster_size), - "min_samples": min_samples, - "cluster_selection_method": str(cluster_selection_method), - }, - }, - sample_size=sample_size, - l2_normalise=True, - assign_k=assign_k, - seed=seed, - ) diff --git a/src/grass_mil/interpretability/core/data.py b/src/grass_mil/interpretability/core/data.py deleted file mode 100644 index 81cc739..0000000 --- a/src/grass_mil/interpretability/core/data.py +++ /dev/null @@ -1,73 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from typing import Iterable, Optional - -import numpy as np -import pandas as pd - -from grass_mil.contracts import EmbeddingSet, InterpretabilityDataset - - -def load_table(path: Path) -> pd.DataFrame: - path = Path(path) - if not path.exists(): - raise FileNotFoundError(path) - suffix = path.suffix.lower() - if suffix == ".csv": - return pd.read_csv(path) - if suffix == ".parquet": - return pd.read_parquet(path) - raise ValueError(f"Unsupported table format for {path}. Expected .csv or .parquet") - - -def load_interpretability_dataset( - *, - instance_table_path: Path, - bag_table_path: Optional[Path] = None, - spatial_table_path: Optional[Path] = None, - cell_table_path: Optional[Path] = None, - cell_edge_table_path: Optional[Path] = None, - id_column: str = "instance_id", - bag_id_column: str = "bag_id", - cell_type_column: Optional[str] = None, - condition_column: Optional[str] = None, -) -> InterpretabilityDataset: - instance_table = load_table(instance_table_path) - bag_table = load_table(bag_table_path) if bag_table_path else None - spatial_table = load_table(spatial_table_path) if spatial_table_path else None - cell_table = load_table(cell_table_path) if cell_table_path else None - cell_edge_table = load_table(cell_edge_table_path) if cell_edge_table_path else None - return InterpretabilityDataset( - instance_table=instance_table, - bag_table=bag_table, - spatial_table=spatial_table, - cell_table=cell_table, - cell_edge_table=cell_edge_table, - id_column=id_column, - bag_id_column=bag_id_column, - cell_type_column=cell_type_column, - condition_column=condition_column, - ) - - -def extract_embedding_set( - dataset: InterpretabilityDataset, - *, - embedding_prefixes: Iterable[str] = ("inst_emb_", "emb_", "graph_emb_"), -) -> EmbeddingSet: - columns = list(dataset.instance_table.columns) - emb_cols: list[str] = [] - for prefix in embedding_prefixes: - emb_cols = [c for c in columns if c.startswith(prefix)] - if emb_cols: - break - if not emb_cols: - raise ValueError( - "No embedding columns found in instance table. " - "Expected prefixes: " + ", ".join(embedding_prefixes) - ) - matrix = dataset.instance_table[emb_cols].to_numpy(dtype=float) - return EmbeddingSet( - matrix=np.asarray(matrix), feature_columns=emb_cols, id_column=dataset.id_column - ) diff --git a/src/grass_mil/interpretability/core/niches.py b/src/grass_mil/interpretability/core/niches.py deleted file mode 100644 index d9ed4fb..0000000 --- a/src/grass_mil/interpretability/core/niches.py +++ /dev/null @@ -1,167 +0,0 @@ -"""Niche identity: naming, the background label, and composition ordering. - -A *niche* is a recurring local tissue neighbourhood — a group of ego-graphs with -similar cellular composition. The clustering algorithm produces integer labels; -this module turns them into something a biologist reads. - -Two conventions: - -- Label ``-1`` is **Background**, not "noise". Density-based clustering assigns - it to neighbourhoods that do not fit any dense group. Those are usually - transitional or sparsely populated tissue, which is a meaningful category - rather than a failure, so it is kept and named. -- Niches are ordered by a dendrogram over their cellular composition, so - compositionally similar niches sit next to each other in every table and - heatmap. Background is pinned last: it is a catch-all, not a niche, and - letting it participate in the ordering would distort the tree. -""" - -from __future__ import annotations - -from typing import Iterable, List - -import numpy as np -import pandas as pd - -__all__ = [ - "BACKGROUND_NICHE_ID", - "BACKGROUND_NICHE_NAME", - "is_background", - "niche_display_name", - "niche_display_names", - "order_niches_by_composition", - "relabel_niches_by_composition", -] - -BACKGROUND_NICHE_ID = -1 -BACKGROUND_NICHE_NAME = "Background" - - -def is_background(niche_id: object) -> bool: - """True for the catch-all label produced by density-based clustering.""" - try: - return int(niche_id) == BACKGROUND_NICHE_ID - except (TypeError, ValueError): - return str(niche_id) == BACKGROUND_NICHE_NAME - - -def niche_display_name(niche_id: object) -> str: - """Human-readable name for one niche. - - The numeric id is retained in the name (``Niche 3``) so a reader can join a - figure back to the ``niche_label`` column without a lookup table. - """ - if is_background(niche_id): - return BACKGROUND_NICHE_NAME - return f"Niche {niche_id}" - - -def niche_display_names(niche_ids: Iterable[object]) -> List[str]: - return [niche_display_name(value) for value in niche_ids] - - -def order_niches_by_composition( - composition: pd.DataFrame, - *, - method: str = "average", - metric: str = "correlation", - background_last: bool = True, -) -> List: - """Order niches so compositionally similar ones are adjacent. - - Uses hierarchical linkage over the per-niche composition profiles and - returns the dendrogram leaf order. Correlation distance is the default - because composition vectors are proportions: two niches with the same - cellular makeup at different densities should read as similar. - - Args: - composition: Niche x feature composition, indexed by niche id. - background_last: Pin Background to the end instead of letting it join - the tree. - - Returns: - Niche ids in dendrogram order. Falls back to the existing order when - there are too few niches to build a tree. - """ - if composition is None or composition.empty: - return [] - - index = list(composition.index) - if background_last: - ordered_ids = [i for i in index if not is_background(i)] - background = [i for i in index if is_background(i)] - else: - ordered_ids, background = list(index), [] - - if len(ordered_ids) < 3: - # Fewer than three leaves: linkage adds nothing over the given order. - return ordered_ids + background - - values = composition.loc[ordered_ids].to_numpy(dtype=float) - finite = np.isfinite(values) - if not finite.all(): - values = np.where(finite, values, 0.0) - - # scipy is a hard dependency; import outside the guard so a genuinely - # missing install raises instead of silently degrading to input order. - from scipy.cluster.hierarchy import leaves_list, linkage - from scipy.spatial.distance import pdist - - distances = pdist(values, metric=metric) - if not np.isfinite(distances).all(): - # Correlation distance is undefined for a constant profile. - distances = pdist(values, metric="euclidean") - try: - order = leaves_list(linkage(distances, method=method)) - except (ValueError, RuntimeError): - return ordered_ids + background - - return [ordered_ids[int(position)] for position in order] + background - - -def relabel_niches_by_composition( - table: pd.DataFrame, - labels: np.ndarray, - *, - composition_prefix: str = "comp_", - method: str = "average", - metric: str = "correlation", - start_at: int = 1, -) -> tuple[np.ndarray, dict]: - """Renumber niches so that the id *is* the dendrogram position. - - Called once, immediately after clustering. Niche 1 is first in dendrogram - order, Niche 2 second, and so on, so every downstream table, figure and - export can simply sort by id and get composition order for free -- no - ordering state has to be threaded through the pipeline. Background keeps its - own label and takes no position in the tree. - - Returns: - ``(relabelled, mapping)`` where ``mapping`` is old id -> new id. - """ - labels = np.asarray(labels) - composition_columns = [c for c in table.columns if str(c).startswith(composition_prefix)] - if not composition_columns: - # Nothing to order by; leave the labels untouched rather than guessing. - return labels, {value: value for value in np.unique(labels)} - - composition = ( - table.loc[:, composition_columns] - .astype(float) - .groupby(pd.Series(labels, name="niche_label")) - .mean() - .sort_index() - ) - order = order_niches_by_composition(composition, method=method, metric=metric) - - mapping = {} - next_id = int(start_at) - for niche_id in order: - if is_background(niche_id): - mapping[niche_id] = BACKGROUND_NICHE_ID - continue - mapping[niche_id] = next_id - next_id += 1 - - relabelled = np.array([mapping.get(value, value) for value in labels]) - return relabelled, mapping diff --git a/src/grass_mil/interpretability/core/reduction.py b/src/grass_mil/interpretability/core/reduction.py deleted file mode 100644 index c0e889a..0000000 --- a/src/grass_mil/interpretability/core/reduction.py +++ /dev/null @@ -1,46 +0,0 @@ -from __future__ import annotations - -from typing import Any, Dict, Optional - -import numpy as np -from sklearn.decomposition import PCA -from sklearn.manifold import TSNE - -from grass_mil.contracts import ReductionResult - - -def run_reduction( - method: str, - X: np.ndarray, - *, - n_components: int = 2, - random_state: Optional[int] = 42, - **params: Any, -) -> ReductionResult: - method_key = str(method).strip().lower() - if method_key == "pca": - model = PCA(n_components=n_components, random_state=random_state, **params) - embedding = model.fit_transform(X) - elif method_key == "umap": - try: - import umap # type: ignore - except Exception as exc: # pragma: no cover - optional dependency - raise RuntimeError( - "UMAP reduction requested but package `umap-learn` is unavailable." - ) from exc - model = umap.UMAP(n_components=n_components, random_state=random_state, **params) - embedding = model.fit_transform(X) - elif method_key == "tsne": - model = TSNE(n_components=n_components, random_state=random_state, **params) - embedding = model.fit_transform(X) - else: - raise ValueError(f"Unsupported reduction method: {method!r}") - - all_params: Dict[str, Any] = {"n_components": int(n_components), "random_state": random_state} - all_params.update(params) - return ReductionResult( - method=method_key, - embedding=np.asarray(embedding), - params=all_params, - fitted_object=model, - ) diff --git a/src/grass_mil/interpretability/core/transfer.py b/src/grass_mil/interpretability/core/transfer.py deleted file mode 100644 index cdc8ebe..0000000 --- a/src/grass_mil/interpretability/core/transfer.py +++ /dev/null @@ -1,217 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from typing import Any, Iterable - -import joblib -import numpy as np -import pandas as pd -from sklearn.neighbors import KNeighborsClassifier - -from grass_mil.contracts import ( - NicheTransferBundle, - InterpretabilityDataset, - ReportBundle, -) -from grass_mil.interpretability.core.data import extract_embedding_set - -_TRANSFER_BUNDLE_SCHEMA_VERSION = 1 - - -def fit_niche_transfer_from_report_bundle( - dataset: InterpretabilityDataset, - report_bundle: ReportBundle, - *, - embedding_prefixes: Iterable[str] = ("inst_emb_", "emb_", "graph_emb_"), - n_neighbors: int = 15, - weights: str = "uniform", - metric: str = "minkowski", -) -> NicheTransferBundle: - if report_bundle.clustering is None: - raise ValueError("Cannot fit niche transfer: report bundle is missing clustering result.") - - emb_set = extract_embedding_set(dataset, embedding_prefixes=tuple(embedding_prefixes)) - labels = np.asarray(report_bundle.clustering.labels) - if labels.ndim != 1: - raise ValueError("Cannot fit niche transfer: clustering labels must be 1D.") - if emb_set.matrix.shape[0] != labels.shape[0]: - raise ValueError( - "Cannot fit niche transfer: embedding rows and clustering labels length mismatch." - ) - if emb_set.matrix.shape[0] == 0: - raise ValueError("Cannot fit niche transfer on an empty instance table.") - - n_neighbors = int(n_neighbors) - if n_neighbors < 1: - raise ValueError("Cannot fit niche transfer: n_neighbors must be >= 1.") - if n_neighbors > int(labels.shape[0]): - raise ValueError( - "Cannot fit niche transfer: n_neighbors cannot exceed number of training rows." - ) - - cluster_on_pca = bool(report_bundle.metadata.get("cluster_on_pca", False)) - pca_model: Any = None - train_features = np.asarray(emb_set.matrix) - if cluster_on_pca: - if report_bundle.niche_feature_reduction is None: - raise ValueError( - "Cannot fit niche transfer: cluster_on_pca=true but " - "report bundle is missing niche_feature_reduction." - ) - pca_result = report_bundle.niche_feature_reduction - pca_model = pca_result.fitted_object - if pca_model is None or not hasattr(pca_model, "transform"): - raise ValueError( - "Cannot fit niche transfer: niche_feature_reduction fitted object " - "is missing a transform(...) method." - ) - train_features = np.asarray(pca_result.embedding) - if train_features.ndim != 2: - raise ValueError( - "Cannot fit niche transfer: niche_feature_reduction embedding must be 2D." - ) - if train_features.shape[0] != labels.shape[0]: - raise ValueError( - "Cannot fit niche transfer: niche_feature_reduction rows and clustering " - "labels length mismatch." - ) - - knn = KNeighborsClassifier(n_neighbors=n_neighbors, weights=weights, metric=metric) - knn.fit(train_features, labels) - label_values = np.asarray(knn.classes_) - metadata = { - "schema_version": _TRANSFER_BUNDLE_SCHEMA_VERSION, - "rows": int(labels.shape[0]), - "n_neighbors": int(n_neighbors), - "weights": str(weights), - "metric": str(metric), - "cluster_on_pca": bool(cluster_on_pca), - "cluster_method": report_bundle.clustering.method, - } - return NicheTransferBundle( - embedding_columns=list(emb_set.feature_columns), - id_column=str(dataset.id_column), - cluster_on_pca=cluster_on_pca, - pca_model=pca_model, - knn_model=knn, - label_values=label_values, - metadata=metadata, - ) - - -def apply_niche_transfer( - bundle: NicheTransferBundle, - instance_table: pd.DataFrame, - *, - id_column: str | None = None, - output_label_column: str = "transferred_cluster_label", - output_confidence_column: str = "transfer_confidence", -) -> pd.DataFrame: - if len(instance_table) == 0: - raise ValueError("Cannot apply niche transfer to an empty instance table.") - - resolved_id_column = str(id_column or bundle.id_column) - if resolved_id_column not in instance_table.columns: - raise ValueError(f"Cannot apply niche transfer: missing id column {resolved_id_column!r}.") - - missing_embeddings = [ - col for col in bundle.embedding_columns if str(col) not in instance_table.columns - ] - if missing_embeddings: - preview = ", ".join(sorted(missing_embeddings)[:8]) - raise ValueError( - "Cannot apply niche transfer: instance table is missing embedding columns: " - f"{preview}." - ) - - if bundle.knn_model is None or not hasattr(bundle.knn_model, "predict"): - raise ValueError( - "Cannot apply niche transfer: transfer bundle is missing a fitted kNN model." - ) - - query_matrix = instance_table[bundle.embedding_columns].to_numpy(dtype=float) - query_features = np.asarray(query_matrix) - if bundle.pca_model is not None: - if not hasattr(bundle.pca_model, "transform"): - raise ValueError( - "Cannot apply niche transfer: transfer bundle PCA model is missing " - "transform(...)." - ) - query_features = np.asarray(bundle.pca_model.transform(query_features)) - - predicted = np.asarray(bundle.knn_model.predict(query_features)) - confidence = np.full((predicted.shape[0],), np.nan, dtype=float) - if hasattr(bundle.knn_model, "predict_proba"): - proba = np.asarray(bundle.knn_model.predict_proba(query_features)) - if proba.ndim == 2 and proba.shape[0] == predicted.shape[0] and proba.shape[1] > 0: - confidence = np.max(proba, axis=1).astype(float, copy=False) - - return pd.DataFrame( - { - resolved_id_column: instance_table[resolved_id_column].tolist(), - output_label_column: predicted.tolist(), - output_confidence_column: confidence.tolist(), - } - ) - - -def save_niche_transfer_bundle(bundle: NicheTransferBundle, path: str | Path) -> None: - output_path = Path(path) - output_path.parent.mkdir(parents=True, exist_ok=True) - payload = { - "schema_version": _TRANSFER_BUNDLE_SCHEMA_VERSION, - "embedding_columns": list(bundle.embedding_columns), - "id_column": str(bundle.id_column), - "cluster_on_pca": bool(bundle.cluster_on_pca), - "pca_model": bundle.pca_model, - "knn_model": bundle.knn_model, - "label_values": np.asarray(bundle.label_values), - "metadata": dict(bundle.metadata), - } - joblib.dump(payload, output_path) - - -def load_niche_transfer_bundle(path: str | Path) -> NicheTransferBundle: - payload = joblib.load(Path(path)) - if not isinstance(payload, dict): - raise ValueError("Invalid transfer bundle format: expected dictionary payload.") - - schema_version = payload.get("schema_version") - try: - schema_version_int = int(schema_version) - except Exception as exc: - raise ValueError( - "Invalid transfer bundle format: missing or invalid schema_version." - ) from exc - if schema_version_int != _TRANSFER_BUNDLE_SCHEMA_VERSION: - raise ValueError( - "Unsupported transfer bundle schema version: " - f"{schema_version!r}. Expected {_TRANSFER_BUNDLE_SCHEMA_VERSION}." - ) - - required_keys = ( - "embedding_columns", - "id_column", - "cluster_on_pca", - "knn_model", - "label_values", - "metadata", - ) - missing_keys = [key for key in required_keys if key not in payload] - if missing_keys: - missing_text = ", ".join(sorted(missing_keys)) - raise ValueError(f"Invalid transfer bundle format: missing keys: {missing_text}.") - - knn_model = payload.get("knn_model") - if knn_model is None or not hasattr(knn_model, "predict"): - raise ValueError("Invalid transfer bundle format: missing or invalid knn_model.") - - return NicheTransferBundle( - embedding_columns=[str(col) for col in payload["embedding_columns"]], - id_column=str(payload["id_column"]), - cluster_on_pca=bool(payload["cluster_on_pca"]), - pca_model=payload.get("pca_model"), - knn_model=knn_model, - label_values=np.asarray(payload["label_values"]), - metadata=dict(payload["metadata"]), - ) diff --git a/src/grass_mil/interpretability/metaclusters/__init__.py b/src/grass_mil/interpretability/metaclusters/__init__.py deleted file mode 100644 index a23c872..0000000 --- a/src/grass_mil/interpretability/metaclusters/__init__.py +++ /dev/null @@ -1,22 +0,0 @@ -"""Cross-fold niche replication: per-fold clustering, matching into metaclusters, -whole-dataset transfer and a shared label-free basis.""" - -from .basis import run_shared_basis -from .match import ( - constrained_average_linkage, - hellinger_matrix, - match_metaclusters, - scan_threshold, -) -from .records import cluster_fold_models -from .transfer import run_transfer - -__all__ = [ - "cluster_fold_models", - "constrained_average_linkage", - "hellinger_matrix", - "match_metaclusters", - "run_shared_basis", - "run_transfer", - "scan_threshold", -] diff --git a/src/grass_mil/interpretability/metaclusters/basis.py b/src/grass_mil/interpretability/metaclusters/basis.py deleted file mode 100644 index 11a7ca4..0000000 --- a/src/grass_mil/interpretability/metaclusters/basis.py +++ /dev/null @@ -1,252 +0,0 @@ -"""Stage 4: a shared, label-free k-means basis with per-model attribution. - -Fit ``MiniBatchKMeans(k)`` once on standardised embeddings of ONE space (a -chosen model, or a pretrained encoder cache), label every cell, then report -per niche the relattn margin of every fold model split into held-out / train / -all (model, region) pairs with bootstrap intervals, next to the empirical -association. Comparable across models because the niches are the same. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Dict, Optional, Sequence - -import numpy as np -import pandas as pd - -from grass_mil.inference.cell_cache import CellCache -from grass_mil.interpretability.niches.margins import relattn_margins -from grass_mil.interpretability.niches.pipeline import load_cells -from grass_mil.interpretability.niches.tables import bootstrap_mean - -__all__ = ["BasisResult", "run_shared_basis", "write_basis_outputs"] - - -@dataclass -class BasisResult: - labels: np.ndarray # (N_all,) niche id - region_ids: np.ndarray - pairs: pd.DataFrame # (model, region, niche, set, relattn, lift, n) - table: pd.DataFrame - basis_model: Any - scaler: Any - settings: Dict[str, Any] - - -def run_shared_basis( - caches: Dict[str, CellCache], - *, - basis_cache: Optional[CellCache] = None, - labels_of: Dict[str, int], - k: int = 20, - sample_size: int = 300_000, - seed: int = 0, - n_init: int = 10, - batch_size: int = 4096, - regions: Optional[Sequence[str]] = None, - positive_class: int = 1, - n_boot: int = 2000, - log=None, -) -> BasisResult: - from sklearn.cluster import MiniBatchKMeans - from sklearn.preprocessing import StandardScaler - - models = list(caches) - basis = basis_cache if basis_cache is not None else caches[models[0]] - regions = sorted(str(r) for r in (regions if regions is not None else basis.regions)) - cells_b = load_cells(basis, regions) - RID = cells_b.region_ids - E = cells_b.E - rng = np.random.default_rng(int(seed)) - sub = rng.choice(len(E), min(int(sample_size), len(E)), replace=False) - scaler = StandardScaler().fit(E[sub]) - km = MiniBatchKMeans( - int(k), random_state=int(seed), n_init=int(n_init), batch_size=int(batch_size) - ).fit(scaler.transform(E[sub])) - labels = km.predict(scaler.transform(E)).astype(int) - fold_of_region = {str(r): name for name, c in caches.items() for r in c.held_out} - rows = [] - for name, cache in caches.items(): - cells = load_cells(cache, regions) - if not np.array_equal(cells.region_ids, RID): - raise ValueError(f"Cache {name} does not cover the same cells as the basis cache.") - rel = relattn_margins(cells.LG, cells.AR, cells.bias, RID, positive_class=positive_class) - a_pos = rel.attention[:, min(positive_class, rel.attention.shape[1] - 1)] - for r in regions: - m = RID == r - lab = labels[m] - mm = rel.margin[m] - aa = a_pos[m] - which = "heldout" if fold_of_region.get(r) == name else "train" - for cl in np.unique(lab): - msk = lab == cl - rows.append( - dict( - model=name, - region=r, - niche=int(cl), - set=which, - relattn=float(mm[msk].mean()), - lift=float(aa[msk].sum() / msk.mean()), - n=int(msk.sum()), - ) - ) - if log is not None: - log.info(f"[{name}] basis attribution over {len(regions)} regions") - pairs = pd.DataFrame(rows) - names = cells_b.cell_type_names - K = len(names) - ct = cells_b.ct - y_cell = np.array([int(labels_of[r]) for r in RID]) - cohort = np.bincount(ct, minlength=K) / max(1, ct.size) - out = [] - for cl in sorted(np.unique(labels)): - msk = labels == cl - comp = np.bincount(ct[msk], minlength=K) / max(1, msk.sum()) - enr = comp / np.maximum(cohort, 1e-9) - rec = dict( - niche=f"N{int(cl)}", - niche_id=int(cl), - share=float(msk.mean()), - regions=int(pairs[pairs.niche == cl].region.nunique()), - emp_frac_pos=float(y_cell[msk].mean()), - top_types=", ".join(f"{names[j]} {100 * comp[j]:.0f}%" for j in np.argsort(-comp)[:3]), - top_enriched=", ".join( - f"{names[j]} x{enr[j]:.1f}" for j in np.argsort(-enr)[:3] if enr[j] > 1.5 - ), - ) - for s in ("heldout", "train"): - v = pairs[(pairs.niche == cl) & (pairs.set == s)] - rec[f"relattn_{s}"], rec[f"relattn_{s}_lo"], rec[f"relattn_{s}_hi"] = bootstrap_mean( - v.relattn, n_boot=n_boot, seed=seed - ) - rec[f"lift_{s}"] = float(v.lift.mean()) if len(v) else float("nan") - v = pairs[pairs.niche == cl] - rec["relattn_all"], rec["relattn_all_lo"], rec["relattn_all_hi"] = bootstrap_mean( - v.relattn, n_boot=n_boot, seed=seed - ) - rec["lift_all"] = float(v.lift.mean()) if len(v) else float("nan") - rec.update({f"comp_{names[j]}": float(comp[j]) for j in range(K)}) - out.append(rec) - table = pd.DataFrame(out) - return BasisResult( - labels=labels, - region_ids=RID, - pairs=pairs, - table=table, - basis_model=km, - scaler=scaler, - settings={ - "k": int(k), - "sample_size": int(sub.size), - "seed": int(seed), - "basis": basis.index.get("model_name"), - "models": models, - "n_boot": int(n_boot), - }, - ) - - -def write_basis_outputs( - result: BasisResult, - out_dir: str | Path, - *, - positive_name: str = "positive", - figures: bool = True, -) -> Dict[str, str]: - import joblib - - out = Path(out_dir) - out.mkdir(parents=True, exist_ok=True) - np.save(out / "basis_labels.npy", result.labels) - joblib.dump( - {"kmeans": result.basis_model, "scaler": result.scaler, "settings": result.settings}, - out / "basis_model.pkl", - ) - result.pairs.to_csv(out / "basis_attribution_pairs.csv", index=False) - result.table.to_csv(out / "basis_attribution.csv", index=False) - T = result.table - md = [ - f"# Shared k-means basis (k = {result.settings['k']}, fitted on {result.settings['basis']}) with per-model attribution", - "", - f"Margins = relattn, region-equal mean, 95% bootstrap over (model, region) pairs; held-out = region scored by the model that held it out; {len(result.settings['models'])} models.", - "", - "| Niche | Share | Regions | Top types | Top enriched | Emp. frac. pos | Margin HELD-OUT | Margin TRAIN | Margin ALL | Lift held-out / train |", - "|---|---|---|---|---|---|---|---|---|---|", - ] - for _, t in T.sort_values("emp_frac_pos", ascending=False).iterrows(): - md.append( - f"| {t.niche} | {100 * t.share:.1f}% | {t.regions} | {t.top_types} | {t.top_enriched} | {t.emp_frac_pos:.2f} | {t.relattn_heldout:+.2f} [{t.relattn_heldout_lo:+.2f}, {t.relattn_heldout_hi:+.2f}] | {t.relattn_train:+.2f} [{t.relattn_train_lo:+.2f}, {t.relattn_train_hi:+.2f}] | {t.relattn_all:+.2f} [{t.relattn_all_lo:+.2f}, {t.relattn_all_hi:+.2f}] | {t.lift_heldout:.2f} / {t.lift_train:.2f} |" - ) - if len(T) >= 3: - from scipy.stats import spearmanr - - for s in ("heldout", "train", "all"): - md.append( - f"\nSpearman(emp_frac_pos, margin_{s}) over {len(T)} niches: rho={spearmanr(T.emp_frac_pos, T[f'relattn_{s}'], nan_policy='omit')[0]:+.2f}" - ) - (out / "report_basis.md").write_text("\n".join(md) + "\n") - paths = {"report": str(out / "report_basis.md")} - if figures: - try: - import matplotlib - - matplotlib.use("Agg") - import matplotlib.pyplot as plt - - T2 = T.sort_values("emp_frac_pos", ascending=False).reset_index(drop=True) - base = float(np.average(T2.emp_frac_pos, weights=T2.share)) - yy = np.arange(len(T2))[::-1] - fig, ax = plt.subplots( - 1, 2, figsize=(10, 0.45 * len(T2) + 1.8), gridspec_kw=dict(width_ratios=[1, 1.6]) - ) - ax[0].barh( - yy, - T2.emp_frac_pos, - color=["#b2182b" if v > base else "#2166ac" for v in T2.emp_frac_pos], - height=0.6, - ) - ax[0].axvline(base, color="k", ls=":", lw=1) - ax[0].set_yticks(yy) - ax[0].set_yticklabels( - [ - f"{n} ({100 * s:.1f}%): {t.split(',')[0]}" - for n, s, t in zip(T2.niche, T2.share, T2.top_enriched) - ], - fontsize=7, - ) - ax[0].set_xlabel(f"fraction of niche cells from {positive_name} regions") - for s, col, off in ( - ("heldout", "#1b9e77", 0.22), - ("train", "#d95f02", 0.0), - ("all", "#7570b3", -0.22), - ): - m_, lo, hi = T2[f"relattn_{s}"], T2[f"relattn_{s}_lo"], T2[f"relattn_{s}_hi"] - ax[1].errorbar( - m_, - yy + off, - xerr=[m_ - lo, hi - m_], - fmt="o", - ms=4, - color=col, - capsize=2, - lw=1, - label=s, - ) - ax[1].axvline(0, color="k", lw=0.8) - ax[1].set_yticks(yy) - ax[1].set_yticklabels([]) - ax[1].set_xlabel(f"attention-weighted margin (relattn), + = {positive_name}") - ax[1].legend(fontsize=7, loc="lower right") - for a in ax: - a.spines["top"].set_visible(False) - a.spines["right"].set_visible(False) - fig.tight_layout() - fig.savefig(out / "basis_attribution.png", dpi=160) - plt.close(fig) - paths["figure"] = str(out / "basis_attribution.png") - except Exception: # pragma: no cover - pass - return paths diff --git a/src/grass_mil/interpretability/metaclusters/match.py b/src/grass_mil/interpretability/metaclusters/match.py deleted file mode 100644 index 4c995ff..0000000 --- a/src/grass_mil/interpretability/metaclusters/match.py +++ /dev/null @@ -1,420 +0,0 @@ -"""Stage 2: match fold-clusters across folds into metaclusters. - -Distance between two fold-clusters = Hellinger distance of their cell-type -compositions (or Aitchison / CLR euclidean). Constrained average linkage never -merges two clusters of the same fold (cannot-link). The threshold is either -fixed or chosen by a scan that maximises the number of metaclusters replicated -in at least ``min_folds`` folds (ties -> fewer metaclusters). -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Dict, List, Optional, Sequence, Tuple - -import numpy as np -import pandas as pd - -__all__ = [ - "MatchResult", - "aitchison_matrix", - "constrained_average_linkage", - "hellinger_matrix", - "match_metaclusters", - "metacluster_table", - "patient_shares", - "patient_tests", - "scan_threshold", -] - - -def hellinger_matrix(comp: np.ndarray) -> np.ndarray: - """Pairwise Hellinger distance between composition rows (in [0, 1]).""" - H = np.sqrt(np.clip(np.asarray(comp, dtype=float), 0.0, None)) - return np.sqrt(0.5 * ((H[:, None, :] - H[None, :, :]) ** 2).sum(-1)) - - -def aitchison_matrix(comp: np.ndarray, pseudo: float = 1e-3) -> np.ndarray: - L = np.log(np.asarray(comp, dtype=float) + pseudo) - L = L - L.mean(axis=1, keepdims=True) - return np.sqrt(((L[:, None, :] - L[None, :, :]) ** 2).sum(-1)) - - -def constrained_average_linkage( - D: np.ndarray, fold: Sequence[Any], threshold: float -) -> List[List[int]]: - """Agglomerate greedily by average distance below ``threshold``; groups sharing a fold never merge.""" - fold = np.asarray(fold) - n = D.shape[0] - groups: List[List[int]] = [[i] for i in range(n)] - while True: - best = None - for a in range(len(groups)): - fa = set(fold[groups[a]].tolist()) - for b in range(a + 1, len(groups)): - if fa & set(fold[groups[b]].tolist()): - continue - d = float(np.mean([D[i, j] for i in groups[a] for j in groups[b]])) - if d < threshold and (best is None or d < best[0]): - best = (d, a, b) - if best is None: - break - _, a, b = best - groups[a] = groups[a] + groups[b] - del groups[b] - return groups - - -def scan_threshold( - D: np.ndarray, fold: Sequence[Any], *, grid: Sequence[float], min_folds: int = 3 -) -> Tuple[float, List[List[int]], pd.DataFrame]: - """Pick the threshold maximising replicated metaclusters, ties -> fewer metaclusters.""" - rows = [] - best = None - for thr in grid: - groups = constrained_average_linkage(D, fold, float(thr)) - rep = sum(len(g) >= int(min_folds) for g in groups) - key = (rep, -len(groups)) - rows.append({"threshold": float(thr), "n_metaclusters": len(groups), "n_replicated": rep}) - if best is None or key > best[0]: - best = (key, float(thr), groups) - return best[1], best[2], pd.DataFrame(rows) - - -@dataclass -class MatchResult: - threshold: float - distance: str - records: pd.DataFrame # fold-cluster records with a `meta` column - table: pd.DataFrame # one row per metacluster - D: np.ndarray - scan: Optional[pd.DataFrame] = None - patient_shares: Optional[pd.DataFrame] = None - patient_tests: Optional[pd.DataFrame] = None - region_shares: Optional[pd.DataFrame] = None - settings: Dict[str, Any] = field(default_factory=dict) - - -def metacluster_table( - records: pd.DataFrame, D: np.ndarray, cohort: np.ndarray, names: Sequence[str] -) -> pd.DataFrame: - comp_cols = [f"comp_{n}" for n in names] - rows = [] - for k in sorted(records["meta"].unique()): - g = records[records["meta"] == k] - comp = np.average( - g[comp_cols].to_numpy(dtype=float), axis=0, weights=g["n_cells"].to_numpy(dtype=float) - ) - enr = comp / np.maximum(np.asarray(cohort, dtype=float), 1e-9) - idx = g.index.to_list() - within = ( - float(np.mean([D[i, j] for i in idx for j in idx if i < j])) if len(idx) > 1 else 0.0 - ) - rows.append( - dict( - meta=int(k), - n_folds=int(g["fold"].nunique()), - n_members=int(len(g)), - folds=",".join(str(f) for f in sorted(g["fold"].unique())), - members=", ".join(g["id"]), - within_dist=within, - share_mean=float(g["share"].mean()), - cores=int(g["cores"].sum()), - pos_cores=int(g["pos_cores"].sum()), - neg_cores=int(g["neg_cores"].sum()), - emp_mean=float(g["emp"].mean()), - emp_rel_mean=float(g["emp_rel"].mean()), - emp_rel_sd=float(g["emp_rel"].std(ddof=1)) if len(g) > 1 else float("nan"), - emp_per_fold=";".join(f"{f}:{e:+.2f}" for f, e in zip(g["fold"], g["emp_rel"])), - margin_mean=float(g["margin"].mean()), - margin_sd=float(g["margin"].std(ddof=1)) if len(g) > 1 else float("nan"), - margin_per_fold=";".join(f"{f}:{m:+.2f}" for f, m in zip(g["fold"], g["margin"])), - sign_agree=float(np.mean(np.sign(g["margin"]) == np.sign(g["margin"].mean()))), - ci_excl0=int(((g["margin_lo"] > 0) | (g["margin_hi"] < 0)).sum()), - emp_margin_agree=float(np.mean(np.sign(g["emp_rel"]) == np.sign(g["margin"]))), - lift_mean=float(g["lift"].mean()), - top_types=", ".join( - f"{names[j]} {100 * comp[j]:.0f}%" for j in np.argsort(-comp)[:4] - ), - top_enriched=", ".join( - f"{names[j]} x{enr[j]:.1f}" for j in np.argsort(-enr)[:4] if enr[j] > 1.5 - ), - **{c: float(v) for c, v in zip(comp_cols, comp)}, - ) - ) - return ( - pd.DataFrame(rows) - .sort_values(["n_folds", "emp_rel_mean"], ascending=[False, False]) - .reset_index(drop=True) - ) - - -def patient_shares( - region_labels: Dict[str, np.ndarray], - region_fold: Dict[str, str], - id_to_meta: Dict[str, int], - patient_of: Dict[str, str], - metas: Sequence[int], -) -> Tuple[pd.DataFrame, pd.DataFrame]: - """Per-region and per-patient metacluster shares (each region labelled once).""" - cols = [f"M{k}" for k in metas] - regions = sorted(region_labels) - SH = pd.DataFrame(0.0, index=regions, columns=cols) - for r in regions: - lab = region_labels[r] - f = region_fold[r] - for c in set(lab.tolist()) - {-1}: - k = id_to_meta.get(f"{f}_C{c}") - if k is not None: - SH.loc[r, f"M{k}"] += float(np.mean(lab == c)) - SH.index.name = "region_id" - pat = SH.copy() - pat["patient_id"] = [patient_of.get(r, r) for r in regions] - return SH, pat.groupby("patient_id").mean() - - -def patient_tests( - shares: pd.DataFrame, y: pd.Series, *, present_threshold: float = 0.0 -) -> pd.DataFrame: - """Mann-Whitney (BH) and AUROC of each metacluster share between classes.""" - from scipy import stats - from sklearn.metrics import roc_auc_score - - yy = y.reindex(shares.index) - ok = yy.notna() - shares = shares.loc[ok] - yy = yy.loc[ok].astype(int).to_numpy() - rows = [] - for col in shares.columns: - v = shares[col].to_numpy(dtype=float) - a, b = v[yy == 1], v[yy == 0] - both = a.size > 2 and b.size > 2 - rows.append( - dict( - meta=int(col[1:]), - share_pos=float(a.mean()) if a.size else float("nan"), - share_neg=float(b.mean()) if b.size else float("nan"), - present_pos=float((a > present_threshold).mean()) if a.size else float("nan"), - present_neg=float((b > present_threshold).mean()) if b.size else float("nan"), - auroc=float(roc_auc_score(yy, v)) if both else float("nan"), - p=float(stats.mannwhitneyu(a, b, alternative="two-sided").pvalue) - if both - else float("nan"), - ) - ) - T = pd.DataFrame(rows) - T["q"] = _bh(T["p"].to_numpy()) - return T.sort_values("p").reset_index(drop=True) - - -def _bh(p: np.ndarray) -> np.ndarray: - p = np.asarray(p, dtype=float) - q = np.full(p.size, np.nan) - ok = ~np.isnan(p) - if ok.sum(): - from scipy.stats import false_discovery_control - - q[ok] = false_discovery_control(p[ok], method="bh") - return q - - -def match_metaclusters( - records: pd.DataFrame, - *, - cell_type_names: Sequence[str], - cohort_frequency: np.ndarray, - distance: str = "hellinger", - threshold: Optional[float] = None, - scan_grid: Sequence[float] = tuple(np.linspace(0.15, 0.6, 46)), - min_folds: int = 3, - region_labels: Optional[Dict[str, np.ndarray]] = None, - region_fold: Optional[Dict[str, str]] = None, - patient_of: Optional[Dict[str, str]] = None, - labels_of: Optional[Dict[str, int]] = None, -) -> MatchResult: - records = records.reset_index(drop=True).copy() - comp_cols = [f"comp_{n}" for n in cell_type_names] - comp = records[comp_cols].to_numpy(dtype=float) - if distance == "hellinger": - D = hellinger_matrix(comp) - elif distance == "aitchison": - D = aitchison_matrix(comp) - else: - raise ValueError("distance must be hellinger or aitchison") - fold = records["fold"].to_numpy() - scan = None - if threshold is None: - threshold, groups, scan = scan_threshold(D, fold, grid=scan_grid, min_folds=min_folds) - else: - groups = constrained_average_linkage(D, fold, float(threshold)) - meta = np.zeros(len(records), dtype=int) - for k, g in enumerate(sorted(groups, key=lambda x: -len(x)), 1): - meta[g] = k - records["meta"] = meta - table = metacluster_table(records, D, cohort_frequency, cell_type_names) - result = MatchResult( - threshold=float(threshold), - distance=distance, - records=records, - table=table, - D=D, - scan=scan, - settings={ - "min_folds": min_folds, - "n_metaclusters": int(len(groups)), - "n_replicated": int(sum(len(g) >= min_folds for g in groups)), - }, - ) - if region_labels is not None and region_fold is not None: - id_to_meta = dict(zip(records["id"], records["meta"])) - pat_of = patient_of or {r: r for r in region_labels} - region_sh, pat_sh = patient_shares( - region_labels, region_fold, id_to_meta, pat_of, sorted(records["meta"].unique()) - ) - result.region_shares = region_sh - result.patient_shares = pat_sh - if labels_of is not None: - y_pat = pd.Series( - {pat_of.get(r, r): int(v) for r, v in labels_of.items() if r in region_labels} - ) - result.patient_tests = patient_tests(pat_sh, y_pat) - return result - - -def write_match_outputs( - result: MatchResult, - out_dir: str | Path, - *, - cell_type_names: Sequence[str], - figures: bool = True, - positive_name: str = "positive", -) -> Dict[str, str]: - out = Path(out_dir) - out.mkdir(parents=True, exist_ok=True) - paths: Dict[str, str] = {} - result.records.to_csv(out / "fold_clusters_matched.csv", index=False) - result.table.to_csv(out / "metaclusters.csv", index=False) - paths["metaclusters"] = str(out / "metaclusters.csv") - np.save(out / "fold_cluster_distances.npy", result.D) - if result.scan is not None: - result.scan.to_csv(out / "threshold_scan.csv", index=False) - if result.patient_shares is not None: - result.patient_shares.to_csv(out / "metacluster_patient_shares.csv") - result.region_shares.to_csv(out / "metacluster_region_shares.csv") - if result.patient_tests is not None: - result.patient_tests.to_csv(out / "metacluster_patient_tests.csv", index=False) - MT = result.table - md = [ - f"# Cross-fold niche replication: {result.distance} distance on composition, cannot-link within fold, threshold {result.threshold:.3f} -> {result.settings['n_metaclusters']} metaclusters, {result.settings['n_replicated']} replicated in >= {result.settings['min_folds']} folds ({len(result.records)} fold-clusters)", - "", - "| Meta | folds | share | regions (pos/neg) | composition | enriched | within-dist | emp. minus base, per fold | mean (sd) | margin per fold | mean (sd) | sign agree | CI excl 0 | emp/margin agree |", - "|---|---|---|---|---|---|---|---|---|---|---|---|---|---|", - ] - for _, r in MT.iterrows(): - md.append( - f"| M{r.meta} | {r.n_folds} ({r.folds}) | {100 * r.share_mean:.1f}% | {r.cores} ({r.pos_cores}/{r.neg_cores}) | {r.top_types} | {r.top_enriched} | {r.within_dist:.2f} | {r.emp_per_fold} | {r.emp_rel_mean:+.2f} ({r.emp_rel_sd:.2f}) | {r.margin_per_fold} | {r.margin_mean:+.2f} ({r.margin_sd:.2f}) | {r.sign_agree:.2f} | {r.ci_excl0}/{r.n_members} | {r.emp_margin_agree:.2f} |" - ) - if result.patient_tests is not None: - nf = dict(zip(MT.meta, MT.n_folds)) - md += [ - "", - f"## Patient-level metacluster shares (each region labelled once, by the model that held it out); {positive_name} vs rest", - "", - "| Meta | folds | share pos | share neg | present pos / neg | AUROC | p | q |", - "|---|---|---|---|---|---|---|---|", - ] - for _, r in result.patient_tests.iterrows(): - md.append( - f"| M{r.meta} | {nf.get(r.meta, '')} | {100 * r.share_pos:.2f}% | {100 * r.share_neg:.2f}% | {100 * r.present_pos:.0f}% / {100 * r.present_neg:.0f}% | {r.auroc:.3f} | {r.p:.3g} | {r.q:.2g} |" - ) - (out / "report.md").write_text("\n".join(md) + "\n") - paths["report"] = str(out / "report.md") - if figures: - try: - _overview_figure( - result, - out / "metaclusters_overview.png", - cell_type_names=cell_type_names, - positive_name=positive_name, - ) - paths["overview"] = str(out / "metaclusters_overview.png") - except Exception: # pragma: no cover - matplotlib missing - pass - return paths - - -def _overview_figure( - result: MatchResult, - path: Path, - *, - cell_type_names: Sequence[str], - positive_name: str, - max_types: int = 20, -) -> None: - import matplotlib - - matplotlib.use("Agg") - import matplotlib.pyplot as plt - - MT = result.table - if MT.empty: - return - comp_cols = [f"comp_{n}" for n in cell_type_names] - cohort = MT[comp_cols].to_numpy(dtype=float).mean(0) - typ = np.argsort(-cohort)[:max_types] - H = MT[[comp_cols[j] for j in typ]].to_numpy(dtype=float) * 100 - m = len(MT) - folds = sorted(result.records["fold"].unique()) - cmap = plt.get_cmap("tab10") - fig, ax = plt.subplots( - 1, 3, figsize=(17, 0.5 * m + 2.4), gridspec_kw=dict(width_ratios=[2.6, 1.1, 1.3]) - ) - im = ax[0].imshow(H, aspect="auto", cmap="viridis", vmin=0, vmax=min(100, H.max())) - ax[0].set_xticks(range(len(typ))) - ax[0].set_xticklabels([cell_type_names[j] for j in typ], rotation=90, fontsize=7) - ax[0].set_yticks(range(m)) - ax[0].set_yticklabels( - [ - f"M{a} ({b} folds, {100 * s:.1f}%)" - for a, b, s in zip(MT.meta, MT.n_folds, MT.share_mean) - ], - fontsize=7, - ) - plt.colorbar(im, ax=ax[0], label="% of metacluster cells", fraction=0.03) - yy = np.arange(m)[::-1] - for k, (_, r) in enumerate(MT.iterrows()): - g = result.records[result.records["meta"] == r.meta] - for f, e, mg, lo, hi in zip( - g["fold"], g["emp_rel"], g["margin"], g["margin_lo"], g["margin_hi"] - ): - fi = folds.index(f) - off = (fi - len(folds) / 2) * 0.12 - ax[1].scatter(e, yy[k] + off, s=20, color=cmap(fi % 10), zorder=3) - ax[2].errorbar( - mg, - yy[k] + off, - xerr=[[mg - lo], [hi - mg]], - fmt="o", - ms=3.5, - color=cmap(fi % 10), - capsize=1.5, - lw=0.8, - ) - for a in ax[1:]: - a.axvline(0, color="k", lw=0.8) - a.set_yticks(yy) - a.set_yticklabels([]) - a.set_ylim(-0.7, m - 0.3) - a.spines["top"].set_visible(False) - a.spines["right"].set_visible(False) - ax[1].set_xlabel(f"empirical {positive_name} fraction minus fold base rate") - ax[2].set_xlabel(f"attention-weighted margin (+ = {positive_name})") - ax[1].set_title("association, one dot per fold", fontsize=9) - ax[2].set_title("attribution, one point per fold", fontsize=9) - for fi, f in enumerate(folds): - ax[2].scatter([], [], color=cmap(fi % 10), label=str(f)) - ax[2].legend(fontsize=7, loc="lower right") - fig.tight_layout() - fig.savefig(path, dpi=150) - plt.close(fig) diff --git a/src/grass_mil/interpretability/metaclusters/records.py b/src/grass_mil/interpretability/metaclusters/records.py deleted file mode 100644 index 91abca7..0000000 --- a/src/grass_mil/interpretability/metaclusters/records.py +++ /dev/null @@ -1,183 +0,0 @@ -"""Stage 1: cluster every fold model's held-out cells in its own space.""" - -from __future__ import annotations - -import json -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Dict, List, Optional, Sequence - -import numpy as np -import pandas as pd - -from grass_mil.inference.cell_cache import CellCache -from grass_mil.interpretability.niches.pipeline import NicheResult, run_niche_discovery - -__all__ = ["FoldClustering", "cluster_fold_models", "load_fold_clustering", "fold_cluster_records"] - - -@dataclass -class FoldClustering: - """Per-fold outputs needed by the later stages.""" - - records: pd.DataFrame # one row per fold-cluster (+ comp_* columns) - region_labels: Dict[str, np.ndarray] = field(default_factory=dict) # region -> cluster ids - region_fold: Dict[str, str] = field(default_factory=dict) # region -> fold model name - cell_type_names: List[str] = field(default_factory=list) - cohort_frequency: Optional[np.ndarray] = None - settings: Dict[str, Any] = field(default_factory=dict) - - -def fold_cluster_records(result: NicheResult, fold: str, base_rate: float) -> pd.DataFrame: - """Records of one fold's clusters from its niche table (noise excluded).""" - T = result.table[result.table["cluster"] != "Noise"].copy() - comp_cols = [c for c in T.columns if c.startswith("comp_")] - rows = [] - for _, t in T.iterrows(): - rows.append( - { - "fold": fold, - "cluster": int(t["cluster_id"]), - "id": f"{fold}_C{int(t['cluster_id'])}", - "share": float(t["share"]), - "n_cells": int(t["n_cells"]), - "cores": int(t["cores_present"]), - "pos_cores": int(t["pos_cores_present"]), - "neg_cores": int(t["neg_cores_present"]), - "emp": float(t["emp_frac_pos"]), - "emp_rel": float(t["emp_frac_pos"]) - float(base_rate), - "margin": float(t["relattn"]), - "margin_lo": float(t["relattn_lo"]), - "margin_hi": float(t["relattn_hi"]), - "lift": float(t["lift"]), - "share_auroc": float(t["share_auroc"]), - "top_types": t["top_types"], - "top_enriched": t["top_enriched"], - **{c: float(t[c]) for c in comp_cols}, - } - ) - return pd.DataFrame(rows) - - -def cluster_fold_models( - caches: Dict[str, CellCache], - *, - labels_of: Optional[Dict[str, int]] = None, - out_dir: Optional[str | Path] = None, - reduction: Optional[Dict[str, Any]] = None, - clustering: Optional[Dict[str, Any]] = None, - sample_size: Optional[int] = 100_000, - l2_normalise: bool = True, - assign_k: int = 15, - seed: int = 0, - n_boot: int = 2000, - positive_class: int = 1, - regions_of: Optional[Dict[str, Sequence[str]]] = None, - log=None, -) -> FoldClustering: - """Cluster each fold model's held-out regions (or ``regions_of[fold]``) with one - fixed setting; save ``labels_.npy``, ``cluster_model_.pkl`` and - the fold-cluster records.""" - import joblib - - out = Path(out_dir) if out_dir is not None else None - if out is not None: - out.mkdir(parents=True, exist_ok=True) - frames: List[pd.DataFrame] = [] - region_labels: Dict[str, np.ndarray] = {} - region_fold: Dict[str, str] = {} - names: List[str] = [] - all_ct: List[np.ndarray] = [] - per_fold_summary = {} - for fold, cache in caches.items(): - regions = ( - list(regions_of[fold]) - if regions_of and fold in regions_of - else (cache.held_out or cache.regions) - ) - res = run_niche_discovery( - cache, - regions=regions, - labels_of=labels_of, - space="gnn", - reduction=reduction, - clustering=clustering, - sample_size=sample_size, - l2_normalise=l2_normalise, - assign_k=assign_k, - seed=seed, - positive_class=positive_class, - n_boot=n_boot, - ) - base = float(res.y_cell.mean()) - frames.append(fold_cluster_records(res, fold, base)) - for r in res.cells.regions: - m = res.cells.region_ids == r - region_labels[r] = res.labels[m] - region_fold[r] = fold - names = res.cells.cell_type_names - all_ct.append(res.cells.ct) - per_fold_summary[fold] = { - "regions": len(res.cells.regions), - "cells": int(res.labels.size), - "clusters": res.n_clusters, - "noise": res.noise_fraction, - "base_rate": base, - **res.concordance, - } - if out is not None: - np.save(out / f"labels_{fold}.npy", res.labels) - joblib.dump(res.model, out / f"cluster_model_{fold}.pkl") - (out / f"regions_{fold}.json").write_text(json.dumps(list(res.cells.regions))) - if log is not None: - log.info( - f"[{fold}] {len(res.cells.regions)} regions, {res.labels.size} cells, {res.n_clusters} clusters, noise {res.noise_fraction:.2f}" - ) - records = pd.concat(frames, ignore_index=True) if frames else pd.DataFrame() - ct = np.concatenate(all_ct) if all_ct else np.zeros(0, dtype=int) - cohort = np.bincount(ct, minlength=len(names)) / max(1, ct.size) if names else None - fc = FoldClustering( - records=records, - region_labels=region_labels, - region_fold=region_fold, - cell_type_names=list(names), - cohort_frequency=cohort, - settings={ - "per_fold": per_fold_summary, - "reduction": reduction, - "clustering": clustering, - "sample_size": sample_size, - "l2_normalise": l2_normalise, - "assign_k": assign_k, - "seed": seed, - }, - ) - if out is not None: - records.to_csv(out / "fold_clusters.csv", index=False) - joblib.dump( - { - "region_labels": region_labels, - "region_fold": region_fold, - "cell_type_names": list(names), - "cohort_frequency": cohort, - "settings": fc.settings, - }, - out / "fold_clustering.pkl", - ) - return fc - - -def load_fold_clustering(out_dir: str | Path) -> FoldClustering: - import joblib - - out = Path(out_dir) - payload = joblib.load(out / "fold_clustering.pkl") - records = pd.read_csv(out / "fold_clusters.csv") - return FoldClustering( - records=records, - cell_type_names=payload["cell_type_names"], - cohort_frequency=payload["cohort_frequency"], - settings=payload["settings"], - region_labels=payload["region_labels"], - region_fold=payload["region_fold"], - ) diff --git a/src/grass_mil/interpretability/metaclusters/transfer.py b/src/grass_mil/interpretability/metaclusters/transfer.py deleted file mode 100644 index ca5d673..0000000 --- a/src/grass_mil/interpretability/metaclusters/transfer.py +++ /dev/null @@ -1,231 +0,0 @@ -"""Stage 3: transfer every fold model's clustering to ALL regions in its own space. - -Each fold's saved cluster model (reducer + k-NN assigner fitted on its held-out -sample) labels every region's cells under that model's embedding; fold-cluster -ids map to metaclusters. Reports, per metacluster and model, share / -association / margin on TRAIN vs HELD-OUT regions, the pairwise ARI of the -models' metacluster labels on cells both assign, and consensus patient shares. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path -from typing import Dict, List, Optional, Sequence - -import numpy as np -import pandas as pd - -from grass_mil.inference.cell_cache import CellCache -from grass_mil.interpretability.niches.margins import relattn_margins -from grass_mil.interpretability.niches.pipeline import load_cells - -from .match import patient_tests - -__all__ = ["TransferResult", "run_transfer", "write_transfer_outputs"] - - -@dataclass -class TransferResult: - meta_labels: np.ndarray # (n_models, N_all) metacluster id or -1 - region_ids: np.ndarray # (N_all,) - models: List[str] - by_model: pd.DataFrame - ari: pd.DataFrame - consensus_region_shares: pd.DataFrame - consensus_patient_shares: Optional[pd.DataFrame] = None - consensus_patient_tests: Optional[pd.DataFrame] = None - refit_agreement: Dict[str, float] = field(default_factory=dict) - - -def run_transfer( - caches: Dict[str, CellCache], - *, - cluster_dir: str | Path, - matched_records: pd.DataFrame, - labels_of: Dict[str, int], - patient_of: Optional[Dict[str, str]] = None, - regions: Optional[Sequence[str]] = None, - positive_class: int = 1, - min_cells: int = 100, - present_threshold: float = 0.02, - log=None, -) -> TransferResult: - import joblib - from sklearn.metrics import adjusted_rand_score - - cluster_dir = Path(cluster_dir) - models = list(caches) - first = caches[models[0]] - regions = sorted(str(r) for r in (regions if regions is not None else first.regions)) - cells0 = load_cells(first, regions) - RID = cells0.region_ids - N = RID.size - y_cell = np.array([int(labels_of[r]) for r in RID]) - fold_of_region = {} - for name, cache in caches.items(): - for r in cache.held_out: - fold_of_region[str(r)] = name - split_model = np.array([fold_of_region.get(r, "") for r in RID], dtype=object) - meta_of = {(str(r.fold), int(r.cluster)): int(r.meta) for _, r in matched_records.iterrows()} - ML = np.full((len(models), N), -1, dtype=int) - MG = np.zeros((len(models), N), dtype=np.float32) - refit_agreement: Dict[str, float] = {} - for mi, name in enumerate(models): - cache = caches[name] - cells = load_cells(cache, regions) - if not np.array_equal(cells.region_ids, RID): - raise ValueError(f"Cache {name} does not cover the same cells as {models[0]}.") - model = joblib.load(cluster_dir / f"cluster_model_{name}.pkl") - labels_all = model.predict(cells.E) - ref_path = cluster_dir / f"labels_{name}.npy" - if ref_path.exists(): - held = split_model == name - ref = np.load(ref_path) - if ref.size == int(held.sum()): - refit_agreement[name] = float((labels_all[held] == ref).mean()) - labels_all[held] = ref # the held-out cells keep their stage-1 labels exactly - ML[mi] = np.array([meta_of.get((name, int(c)), -1) if c >= 0 else -1 for c in labels_all]) - rel = relattn_margins(cells.LG, cells.AR, cells.bias, RID, positive_class=positive_class) - MG[mi] = rel.margin - if log is not None: - log.info( - f"[{name}] transferred to {len(regions)} regions; refit agreement {refit_agreement.get(name, float('nan')):.3f}" - ) - metas = sorted(set(matched_records["meta"].astype(int))) - rows = [] - for k in metas: - for mi, name in enumerate(models): - mk = ML[mi] == k - if mk.sum() < int(min_cells): - continue - for split_name, sm in ( - ("train", split_model != name), - ("heldout", split_model == name), - ): - ms = mk & sm - if not sm.any(): - continue - base = float(y_cell[sm].mean()) - present = [ - r - for r in regions - if ((split_model[RID == r][0] == name) == (split_name == "heldout")) - and ms[RID == r].mean() >= present_threshold - ] - vals = [ - float(MG[mi][ms & (RID == r)].mean()) - for r in present - if (ms & (RID == r)).any() - ] - rows.append( - dict( - meta=k, - model=name, - split=split_name, - share=float(ms.sum() / sm.sum()), - cores=len(present), - pos_cores=int(sum(labels_of[r] for r in present)), - emp=float(y_cell[ms].mean()) if ms.any() else float("nan"), - emp_rel=(float(y_cell[ms].mean()) - base) if ms.any() else float("nan"), - margin=float(np.mean(vals)) if vals else float("nan"), - ) - ) - by_model = pd.DataFrame(rows) - ari = pd.DataFrame(np.nan, index=models, columns=models) - for i, a in enumerate(models): - for j, b in enumerate(models): - both = (ML[i] >= 0) & (ML[j] >= 0) - if both.sum() > 10: - ari.loc[a, b] = adjusted_rand_score(ML[i][both], ML[j][both]) - region_shares = pd.DataFrame( - { - f"M{k}": [ - float(np.mean([np.mean(ML[i][RID == r] == k) for i in range(len(models))])) - for r in regions - ] - for k in metas - }, - index=regions, - ) - region_shares.index.name = "region_id" - result = TransferResult( - meta_labels=ML, - region_ids=RID, - models=models, - by_model=by_model, - ari=ari, - consensus_region_shares=region_shares, - refit_agreement=refit_agreement, - ) - pat_of = patient_of or {r: r for r in regions} - pat = region_shares.copy() - pat["patient_id"] = [pat_of.get(r, r) for r in regions] - result.consensus_patient_shares = pat.groupby("patient_id").mean() - y_pat = pd.Series({pat_of.get(r, r): int(labels_of[r]) for r in regions}) - result.consensus_patient_tests = patient_tests( - result.consensus_patient_shares, y_pat, present_threshold=0.005 - ) - return result - - -def write_transfer_outputs( - result: TransferResult, - out_dir: str | Path, - *, - metaclusters: pd.DataFrame, - positive_name: str = "positive", -) -> Dict[str, str]: - out = Path(out_dir) - out.mkdir(parents=True, exist_ok=True) - np.save(out / "transfer_metalabels.npy", result.meta_labels) - pd.DataFrame({"region_id": result.region_ids}).to_csv(out / "transfer_cells.csv", index=False) - result.by_model.to_csv(out / "transfer_meta_by_model.csv", index=False) - result.ari.to_csv(out / "transfer_model_ari.csv") - result.consensus_region_shares.to_csv(out / "transfer_region_shares.csv") - if result.consensus_patient_shares is not None: - result.consensus_patient_shares.to_csv(out / "transfer_patient_shares.csv") - if result.consensus_patient_tests is not None: - result.consensus_patient_tests.to_csv(out / "transfer_patient_tests.csv", index=False) - MT = metaclusters.set_index("meta") - models = result.models - pairs = [(a, b) for i, a in enumerate(models) for b in models[i + 1 :]] - md = [ - "# Metaclusters transferred to all regions by each fold model", - "", - "Stage-1 refit agreement (held-out cells relabelled by the saved model): " - + "; ".join(f"{k} {v:.3f}" for k, v in result.refit_agreement.items()), - "", - "Pairwise ARI of metacluster labels between models on cells both assign: " - + "; ".join(f"{a}-{b} {result.ari.loc[a, b]:.2f}" for a, b in pairs), - "", - "| Meta | folds | composition | model | share train / held-out | regions train (pos) / held-out (pos) | emp-base train / held-out | margin train / held-out |", - "|---|---|---|---|---|---|---|---|", - ] - for k in MT.index: - g = result.by_model[result.by_model["meta"] == k] - first = True - for m in models: - a = g[(g.model == m) & (g.split == "train")] - h = g[(g.model == m) & (g.split == "heldout")] - if a.empty or h.empty: - continue - a, h = a.iloc[0], h.iloc[0] - md.append( - f"| {'M' + str(k) if first else ''} | {MT.loc[k, 'n_folds'] if first else ''} | {MT.loc[k, 'top_types'] if first else ''} | {m} | {100 * a.share:.2f}% / {100 * h.share:.2f}% | {a.cores} ({a.pos_cores}) / {h.cores} ({h.pos_cores}) | {a.emp_rel:+.2f} / {h.emp_rel:+.2f} | {a.margin:+.2f} / {h.margin:+.2f} |" - ) - first = False - if result.consensus_patient_tests is not None: - md += [ - "", - "## Consensus patient-level shares (mean over the models' assignments of every region)", - "", - "| Meta | folds | composition | share pos | share neg | present pos / neg | AUROC | p | q |", - "|---|---|---|---|---|---|---|---|---|", - ] - for _, r in result.consensus_patient_tests.iterrows(): - md.append( - f"| M{r.meta} | {MT.loc[r.meta, 'n_folds'] if r.meta in MT.index else ''} | {MT.loc[r.meta, 'top_types'] if r.meta in MT.index else ''} | {100 * r.share_pos:.2f}% | {100 * r.share_neg:.2f}% | {100 * r.present_pos:.0f}% / {100 * r.present_neg:.0f}% | {r.auroc:.3f} | {r.p:.3g} | {r.q:.2g} |" - ) - (out / "report_transfer.md").write_text("\n".join(md) + "\n") - return {"report": str(out / "report_transfer.md")} diff --git a/src/grass_mil/interpretability/metaclusters_cli.py b/src/grass_mil/interpretability/metaclusters_cli.py deleted file mode 100644 index 7e7aed1..0000000 --- a/src/grass_mil/interpretability/metaclusters_cli.py +++ /dev/null @@ -1,187 +0,0 @@ -"""``grass-mil-metaclusters``: cross-fold niche replication (cluster / match / transfer / basis).""" - -from __future__ import annotations - -import json -from pathlib import Path -from typing import Any, Dict, List, Optional - -import hydra -import numpy as np -import pandas as pd -from omegaconf import DictConfig, OmegaConf - -from grass_mil._paths import setup_project_root - -setup_project_root() - -from grass_mil.inference.cell_cache import CellCache # noqa: E402 -from grass_mil.interpretability.metaclusters.basis import run_shared_basis, write_basis_outputs # noqa: E402 -from grass_mil.interpretability.metaclusters.match import match_metaclusters, write_match_outputs # noqa: E402 -from grass_mil.interpretability.metaclusters.records import ( - cluster_fold_models, - load_fold_clustering, -) # noqa: E402 -from grass_mil.interpretability.metaclusters.transfer import run_transfer, write_transfer_outputs # noqa: E402 -from grass_mil.utils import RankedLogger # noqa: E402 - -log = RankedLogger(__name__, rank_zero_only=True) - -__all__ = ["run_metaclusters", "main", "discover_caches", "load_labels_and_patients"] - - -def discover_caches( - cells_dir: str | Path, models: Optional[List[str]] = None -) -> Dict[str, CellCache]: - root = Path(cells_dir) - if (root / "index.json").exists(): - c = CellCache.open(root) - return {c.index.get("model_name", root.name): c} - found: Dict[str, CellCache] = {} - for sub in sorted(p for p in root.iterdir() if p.is_dir() and (p / "index.json").exists()): - c = CellCache.open(sub) - name = c.index.get("model_name", sub.name) - if models and name not in models: - continue - found[name] = c - if not found: - raise FileNotFoundError(f"No cell caches under {root}") - return found - - -def load_labels_and_patients(cfg_labels): - df = pd.read_csv(str(cfg_labels.label_file)) - df.columns = [str(c).strip() for c in df.columns] - idc = str(cfg_labels.id_column) - lc = str(cfg_labels.label_column) - df = df.dropna(subset=[lc]) - labels_of = {str(r): int(float(v) >= 0.5) for r, v in zip(df[idc], df[lc])} - pc = cfg_labels.get("patient_column") - patient_of = {str(r): str(p) for r, p in zip(df[idc], df[pc])} if pc else None - return labels_of, patient_of - - -def run_metaclusters(cfg: DictConfig) -> Dict[str, Any]: - out = Path(str(cfg.output_dir)) - out.mkdir(parents=True, exist_ok=True) - caches = discover_caches(str(cfg.cells_dir), list(cfg.models) if cfg.get("models") else None) - held = {k: v for k, v in caches.items() if v.held_out} - if not held: - raise ValueError( - "No cache records held-out regions; extract from a nested-CV summary (or +extract.held_out)." - ) - caches = held - labels_of, patient_of = load_labels_and_patients(cfg.labels) - stage = str(cfg.get("stage", "all")).lower() - stages = ["cluster", "match", "transfer", "basis"] if stage == "all" else [stage] - cluster_dir = out / "cluster" - match_dir = out / "match" - summary: Dict[str, Any] = {"models": list(caches)} - names = None - cohort = None - fc = None - if "cluster" in stages: - fc = cluster_fold_models( - caches, - labels_of=labels_of, - out_dir=cluster_dir, - reduction=OmegaConf.to_container(cfg.reduction, resolve=True), - clustering=OmegaConf.to_container(cfg.clustering, resolve=True), - sample_size=cfg.get("sample_size"), - l2_normalise=bool(cfg.get("l2_normalise", True)), - assign_k=int(cfg.get("assign_k", 15)), - seed=int(cfg.get("seed", 0)), - n_boot=int(cfg.get("n_boot", 2000)), - positive_class=int(cfg.get("positive_class", 1)), - log=log, - ) - summary["cluster"] = fc.settings["per_fold"] - if "match" in stages or "transfer" in stages or "basis" in stages: - if fc is None: - fc = load_fold_clustering(cluster_dir) - names = fc.cell_type_names - cohort = fc.cohort_frequency - matched = None - if "match" in stages: - m = cfg.match - grid = np.linspace(float(m.scan.min), float(m.scan.max), int(m.scan.n)) - res = match_metaclusters( - fc.records, - cell_type_names=names, - cohort_frequency=cohort, - distance=str(m.distance), - threshold=m.get("threshold"), - scan_grid=grid, - min_folds=int(m.min_folds), - region_labels=fc.region_labels, - region_fold=fc.region_fold, - patient_of=patient_of, - labels_of=labels_of, - ) - paths = write_match_outputs( - res, - match_dir, - cell_type_names=names, - figures=bool(cfg.get("figures", True)), - positive_name=str(cfg.positive_name), - ) - matched = res.records - summary["match"] = {"threshold": res.threshold, **res.settings, "report": paths["report"]} - log.info( - f"match: threshold {res.threshold:.3f}, {res.settings['n_metaclusters']} metaclusters, {res.settings['n_replicated']} replicated" - ) - if "transfer" in stages: - if matched is None: - matched = pd.read_csv(match_dir / "fold_clusters_matched.csv") - mt = pd.read_csv(match_dir / "metaclusters.csv") - res_t = run_transfer( - caches, - cluster_dir=cluster_dir, - matched_records=matched, - labels_of=labels_of, - patient_of=patient_of, - positive_class=int(cfg.get("positive_class", 1)), - min_cells=int(cfg.transfer.min_cells), - present_threshold=float(cfg.transfer.present_threshold), - log=log, - ) - paths = write_transfer_outputs( - res_t, out / "transfer", metaclusters=mt, positive_name=str(cfg.positive_name) - ) - summary["transfer"] = {"refit_agreement": res_t.refit_agreement, "report": paths["report"]} - if "basis" in stages and bool(cfg.basis.get("enabled", True)): - b = cfg.basis - basis_cache = None - sm = b.get("space_model") - if sm: - basis_cache = caches.get(str(sm)) or CellCache.open(str(sm)) - res_b = run_shared_basis( - caches, - basis_cache=basis_cache, - labels_of=labels_of, - k=int(b.k), - sample_size=int(b.sample_size), - seed=int(b.seed), - n_init=int(b.n_init), - positive_class=int(cfg.get("positive_class", 1)), - n_boot=int(cfg.get("n_boot", 2000)), - log=log, - ) - paths = write_basis_outputs( - res_b, - out / "basis", - positive_name=str(cfg.positive_name), - figures=bool(cfg.get("figures", True)), - ) - summary["basis"] = {"k": int(b.k), "report": paths["report"]} - (out / "metaclusters_summary.json").write_text(json.dumps(summary, indent=1, default=str)) - return summary - - -@hydra.main(version_base="1.3", config_path="../configs", config_name="metaclusters.yaml") -def main(cfg: DictConfig) -> None: - run_metaclusters(cfg) - - -if __name__ == "__main__": - main() diff --git a/src/grass_mil/interpretability/niches/__init__.py b/src/grass_mil/interpretability/niches/__init__.py deleted file mode 100644 index 654aa8a..0000000 --- a/src/grass_mil/interpretability/niches/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -"""Niche discovery on a trained GRASS-MIL model's per-cell cache.""" - -from .colocalisation import cluster_colocalisation -from .margins import attention_space_from_checkpoint, relattn_margins -from .pipeline import NicheResult, run_niche_discovery -from .tables import niche_table, spearman_concordance - -__all__ = [ - "NicheResult", - "attention_space_from_checkpoint", - "cluster_colocalisation", - "niche_table", - "relattn_margins", - "run_niche_discovery", - "spearman_concordance", -] diff --git a/src/grass_mil/interpretability/niches/colocalisation.py b/src/grass_mil/interpretability/niches/colocalisation.py deleted file mode 100644 index 0534e25..0000000 --- a/src/grass_mil/interpretability/niches/colocalisation.py +++ /dev/null @@ -1,34 +0,0 @@ -"""Niche-niche co-localisation on the cell graph: log2(observed / expected) edge counts.""" - -from __future__ import annotations - -from typing import List, Tuple - -import numpy as np - -__all__ = ["cluster_colocalisation"] - - -def cluster_colocalisation( - labels: np.ndarray, edge_index: np.ndarray -) -> Tuple[np.ndarray, List[int]]: - """``log2((observed + 1) / (E + 1))`` with ``E_ab = d_a d_b / sum(d)`` from the - symmetrised cluster-cluster edge count matrix ``observed`` (noise cells ignored).""" - labels = np.asarray(labels).astype(int) - ei = np.asarray(edge_index).astype(int) - la, lb = labels[ei[0]], labels[ei[1]] - keep = (la >= 0) & (lb >= 0) - la, lb = la[keep], lb[keep] - clusters = sorted(set(labels.tolist()) - {-1}) - idx = {c: i for i, c in enumerate(clusters)} - n = len(clusters) - observed = np.zeros((n, n)) - if la.size: - ia = np.vectorize(idx.get)(la) - ib = np.vectorize(idx.get)(lb) - np.add.at(observed, (ia, ib), 1) - observed = observed + observed.T - deg = observed.sum(axis=1) - total = deg.sum() - expected = np.outer(deg, deg) / total if total > 0 else np.zeros_like(observed) - return np.log2((observed + 1.0) / (expected + 1.0)), clusters diff --git a/src/grass_mil/interpretability/niches/figures.py b/src/grass_mil/interpretability/niches/figures.py deleted file mode 100644 index d86f9a1..0000000 --- a/src/grass_mil/interpretability/niches/figures.py +++ /dev/null @@ -1,180 +0,0 @@ -"""Figures of the niche report (matplotlib; UMAP optional).""" - -from __future__ import annotations - -from pathlib import Path -from typing import Optional, Sequence - -import numpy as np -import pandas as pd - -__all__ = ["embedding_triptych", "clusters_overview", "colocalisation_heatmap", "umap_coordinates"] - - -def _plt(): - import matplotlib - - matplotlib.use("Agg") - import matplotlib.pyplot as plt - - return plt - - -def umap_coordinates( - P: np.ndarray, - *, - n_neighbors: int = 30, - min_dist: float = 0.1, - metric: str = "cosine", - seed: int = 0, -) -> Optional[np.ndarray]: - try: - import umap # type: ignore - except Exception: - return None - return umap.UMAP( - n_neighbors=int(n_neighbors), - min_dist=float(min_dist), - metric=str(metric), - random_state=int(seed), - ).fit_transform(P) - - -def embedding_triptych( - coords: np.ndarray, - labels: np.ndarray, - margin: np.ndarray, - y_cell: np.ndarray, - path: Path, - *, - title: str, - axis_labels: Sequence[str] = ("dim 1", "dim 2"), - positive_name: str = "positive", - margin_clip: float = 3.0, -) -> Path: - """Clusters | attention-weighted margin | empirical association, side by side.""" - plt = _plt() - clusters = sorted(set(labels.tolist()) - {-1}) - cmap = plt.get_cmap("tab20") - fig, ax = plt.subplots(1, 3, figsize=(15, 4.6)) - noise = labels < 0 - ax[0].scatter(coords[noise, 0], coords[noise, 1], s=1, c="#cccccc", lw=0, label="noise") - for i, c in enumerate(clusters): - mm = labels == c - ax[0].scatter(coords[mm, 0], coords[mm, 1], s=1.5, color=cmap(i % 20), lw=0, label=f"C{c}") - ax[0].legend(markerscale=8, fontsize=6, ncol=2, loc="best") - ax[0].set_title(f"{title}: clusters ({len(clusters)}, noise {noise.mean():.0%})", fontsize=9) - mv = np.clip(margin, -margin_clip, margin_clip) - sc = ax[1].scatter( - coords[:, 0], - coords[:, 1], - s=1.5, - c=mv, - cmap="RdBu_r", - vmin=-margin_clip, - vmax=margin_clip, - lw=0, - ) - plt.colorbar(sc, ax=ax[1], label=f"relattn margin (+ = {positive_name})") - ax[1].set_title("attention-weighted margin", fontsize=9) - hb = ax[2].hexbin( - coords[:, 0], coords[:, 1], C=y_cell, gridsize=45, cmap="RdBu_r", vmin=0, vmax=1, mincnt=5 - ) - plt.colorbar(hb, ax=ax[2], label=f"fraction of cells from {positive_name} regions") - ax[2].set_title("empirical association", fontsize=9) - for a in ax: - a.set_xlabel(axis_labels[0]) - a.set_ylabel(axis_labels[1]) - a.set_xticks([]) - a.set_yticks([]) - fig.tight_layout() - fig.savefig(path, dpi=150) - plt.close(fig) - return path - - -def clusters_overview( - table: pd.DataFrame, - path: Path, - *, - cell_type_names: Sequence[str], - cohort_frequency: np.ndarray, - base_rate: float, - positive_name: str = "positive", - max_types: int = 20, -) -> Path: - """Composition heatmap | empirical association | margin with CI, per cluster.""" - plt = _plt() - tt = table[table["cluster"] != "Noise"].sort_values("emp_frac_pos", ascending=False) - if tt.empty: - return path - typ = np.argsort(-np.asarray(cohort_frequency))[:max_types] - comp_cols = [f"comp_{cell_type_names[j]}" for j in typ] - H = tt[comp_cols].to_numpy(dtype=float) * 100 - fig, ax = plt.subplots( - 1, 3, figsize=(16, 0.42 * len(tt) + 2.2), gridspec_kw=dict(width_ratios=[2.6, 1, 1.3]) - ) - im = ax[0].imshow( - H, aspect="auto", cmap="viridis", vmin=0, vmax=min(100, H.max() if H.size else 100) - ) - ax[0].set_xticks(range(len(typ))) - ax[0].set_xticklabels([cell_type_names[j] for j in typ], rotation=90, fontsize=7) - ax[0].set_yticks(range(len(tt))) - ax[0].set_yticklabels( - [ - f"{c} ({100 * s:.1f}%, {n} regions)" - for c, s, n in zip(tt.cluster, tt.share, tt.cores_present) - ], - fontsize=7, - ) - plt.colorbar(im, ax=ax[0], label="% of cluster cells", fraction=0.03) - yy = np.arange(len(tt))[::-1] - ax[1].barh( - yy, - tt.emp_frac_pos, - color=["#b2182b" if v > base_rate else "#2166ac" for v in tt.emp_frac_pos], - height=0.6, - ) - ax[1].axvline(base_rate, color="k", ls=":") - ax[1].set_yticks([]) - ax[1].set_xlabel(f"fraction of cells from {positive_name} regions") - ax[1].set_xlim(0, 1) - ax[2].errorbar( - tt.relattn, - yy, - xerr=[tt.relattn - tt.relattn_lo, tt.relattn_hi - tt.relattn], - fmt="o", - ms=4, - color="#1b9e77", - capsize=2, - ) - ax[2].axvline(0, color="k", lw=0.8) - ax[2].set_yticks([]) - ax[2].set_xlabel(f"attention-weighted margin (+ = {positive_name})") - for a in ax[1:]: - a.spines["top"].set_visible(False) - a.spines["right"].set_visible(False) - fig.tight_layout() - fig.savefig(path, dpi=150) - plt.close(fig) - return path - - -def colocalisation_heatmap(CO: np.ndarray, clusters: Sequence[int], path: Path) -> Path: - plt = _plt() - n = len(clusters) - if n == 0: - return path - fig, ax = plt.subplots(figsize=(0.45 * n + 2, 0.45 * n + 1.5)) - v = float(np.nanmax(np.abs(CO))) if np.isfinite(CO).any() else 1.0 - im = ax.imshow(CO, cmap="RdBu_r", vmin=-v, vmax=v) - ax.set_xticks(range(n)) - ax.set_xticklabels([f"C{c}" for c in clusters], rotation=90, fontsize=7) - ax.set_yticks(range(n)) - ax.set_yticklabels([f"C{c}" for c in clusters], fontsize=7) - plt.colorbar(im, ax=ax, label="log2 obs/exp cell-cell edges") - ax.set_title("cluster co-localisation", fontsize=9) - fig.tight_layout() - fig.savefig(path, dpi=150) - plt.close(fig) - return path diff --git a/src/grass_mil/interpretability/niches/margins.py b/src/grass_mil/interpretability/niches/margins.py deleted file mode 100644 index ffc48d6..0000000 --- a/src/grass_mil/interpretability/niches/margins.py +++ /dev/null @@ -1,108 +0,0 @@ -"""Per-cell attention-weighted margins and the attention feature space.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Sequence - -import numpy as np - -__all__ = ["RelattnResult", "relattn_margins", "attention_space_from_checkpoint"] - - -@dataclass -class RelattnResult: - margin: np.ndarray # (N,) N_r * [A1 (l1 - b1) - A0 (l0 - b0)] (binary) - attention: np.ndarray # (N, C) attention normalised within region per class - contributions: np.ndarray # (N, C) A_c * (l_c - b_c) - - -def relattn_margins( - instance_logits: np.ndarray, - attention_logits: np.ndarray, - bias: Sequence[float], - region_ids: np.ndarray, - *, - positive_class: int = 1, -) -> RelattnResult: - """Attention-weighted, region-size-scaled cell margins (``relattn``). - - Attention is a softmax over the region's cells per class, so a cell's - contribution ``A[i,c] (l[i,c] - b_c)`` sums over the region to the bag - logit minus bias. Multiplying by the region's cell count ``N_r`` makes the - quantity comparable across regions of different size (a uniformly attended - cell gets weight 1). For a binary head the margin contrasts the positive - class with class 0; for more classes it is one-vs-rest. - """ - LG = np.asarray(instance_logits, dtype=np.float64) - AR = np.asarray(attention_logits, dtype=np.float64) - b = np.asarray(bias, dtype=np.float64).reshape(-1) - rid = np.asarray(region_ids) - if LG.shape != AR.shape: - raise ValueError(f"instance_logits {LG.shape} and attention_logits {AR.shape} differ") - n, c = LG.shape - if b.size != c: - raise ValueError(f"bias has {b.size} entries for {c} classes") - A = np.zeros_like(AR) - scale = np.zeros(n, dtype=np.float64) - for r in np.unique(rid): - m = rid == r - ar = AR[m] - a = np.exp(ar - ar.max(axis=0, keepdims=True)) - a /= a.sum(axis=0, keepdims=True) - A[m] = a - scale[m] = float(m.sum()) - contrib = A * (LG - b[None, :]) - pos = int(positive_class) - if c == 1: - margin = contrib[:, 0] - elif c == 2: - margin = contrib[:, pos] - contrib[:, 1 - pos] - else: - rest = (contrib.sum(axis=1) - contrib[:, pos]) / (c - 1) - margin = contrib[:, pos] - rest - return RelattnResult( - margin=(margin * scale).astype(np.float32), - attention=A.astype(np.float32), - contributions=contrib.astype(np.float32), - ) - - -def attention_space_from_checkpoint( - embeddings: np.ndarray, ckpt_path: str, *, batch_size: int = 65536 -) -> np.ndarray: - """The attention head's pre-tanh activation for every cell, from a checkpoint. - - Rebuilds the projected gated attention head from the ``attention.*`` weights - (dimensions inferred from the tensors) and applies ``attention_space``. - """ - import torch - - from grass_mil.inference.model_loading import read_checkpoint_state - from grass_mil.models.components.attention import AttnNetGatedProjected - - state = read_checkpoint_state(ckpt_path) - att = {k[len("attention.") :]: v for k, v in state.items() if k.startswith("attention.")} - if "projection.0.weight" not in att: - raise ValueError( - f"{ckpt_path} has no projected gated attention (attention.projection.0.weight)." - ) - proj_dim, in_dim = att["projection.0.weight"].shape - hidden = att["attention_a.0.weight"].shape[0] - n_classes = att["attention_c.weight"].shape[0] - module = AttnNetGatedProjected( - input_dim=int(in_dim), - projection_dim=int(proj_dim), - hidden_dim=int(hidden), - n_classes=int(n_classes), - ) - module.load_state_dict(att, strict=True) - module.eval() - out = [] - E = np.asarray(embeddings, dtype=np.float32) - with torch.no_grad(): - for i in range(0, E.shape[0], int(batch_size)): - out.append( - module.attention_space(torch.from_numpy(E[i : i + int(batch_size)])).numpy() - ) - return np.concatenate(out) if out else np.zeros((0, int(hidden)), dtype=np.float32) diff --git a/src/grass_mil/interpretability/niches/pipeline.py b/src/grass_mil/interpretability/niches/pipeline.py deleted file mode 100644 index ef7eb77..0000000 --- a/src/grass_mil/interpretability/niches/pipeline.py +++ /dev/null @@ -1,372 +0,0 @@ -"""Niche discovery on one model's cell cache. - -Steps (preprint): take the cells of the chosen regions (held-out by default), -compute the relattn margin of every cell, cluster the chosen feature space -(encoder embedding or the attention head's 8-d space) with the user's -reduction + clustering settings (default cosine HDBSCAN: PCA-10 on a 100k -sample, L2, HDBSCAN mcs 300 / min_samples 1 / eom, 15-NN assignment), then -describe every cluster (composition, enrichment, empirical association, -margin with bootstrap CI over regions, share AUROC, co-localisation) and -draw the figures. -""" - -from __future__ import annotations - -import json -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Dict, List, Optional, Sequence - -import numpy as np -import pandas as pd - -from grass_mil.inference.cell_cache import CellCache -from grass_mil.interpretability.core.clustering import NicheClusterModel, fit_niche_clustering - -from .colocalisation import cluster_colocalisation -from .margins import attention_space_from_checkpoint, relattn_margins -from .tables import niche_table, spearman_concordance - -__all__ = ["NicheResult", "load_cells", "run_niche_discovery", "write_niche_outputs"] - -DEFAULT_REDUCTION = {"method": "pca", "n_components": 10} -DEFAULT_CLUSTERING = { - "method": "hdbscan", - "params": {"min_cluster_size": 300, "min_samples": 1, "cluster_selection_method": "eom"}, -} - - -@dataclass -class CellSet: - regions: List[str] - region_ids: np.ndarray # (N,) - E: np.ndarray # (N, d) float32 - LG: np.ndarray - AR: Optional[np.ndarray] - ct: np.ndarray - pos: np.ndarray - edge_index: np.ndarray # (2, M) global indices - bias: np.ndarray - cell_type_names: List[str] - - -def load_cells(cache: CellCache, regions: Sequence[str]) -> CellSet: - """Concatenate the cached arrays of ``regions`` (in the given order).""" - E, LG, AR, CT, POS, EDGES, RID = [], [], [], [], [], [], [] - off = 0 - for r in regions: - d = cache.load(r) - n = int(d["E"].shape[0]) - E.append(d["E"].astype(np.float32)) - LG.append(d["LG"]) - AR.append(d["AR"] if "AR" in d else None) - CT.append(d["ct"].astype(int) if "ct" in d else np.zeros(n, dtype=int)) - POS.append(d["pos"]) - EDGES.append(d["edge_index"].astype(np.int64) + off) - RID.append(np.full(n, str(r), dtype=object)) - off += n - if not E: - raise ValueError("No regions to load.") - names = cache.cell_type_names - ct = np.concatenate(CT) - if names is None: - names = [str(i) for i in range(int(ct.max()) + 1 if ct.size else 0)] - return CellSet( - regions=list(regions), - region_ids=np.concatenate(RID), - E=np.concatenate(E), - LG=np.concatenate(LG), - AR=np.concatenate(AR) if all(a is not None for a in AR) else None, - ct=ct, - pos=np.concatenate(POS), - edge_index=np.concatenate(EDGES, axis=1), - bias=cache.bias, - cell_type_names=list(names), - ) - - -@dataclass -class NicheResult: - labels: np.ndarray - model: NicheClusterModel - table: pd.DataFrame - margin: np.ndarray - attention: np.ndarray - colocalisation: np.ndarray - clusters: List[int] - concordance: Dict[str, float] - cells: CellSet - features: np.ndarray - settings: Dict[str, Any] = field(default_factory=dict) - y_cell: Optional[np.ndarray] = None - - @property - def n_clusters(self) -> int: - return len(self.clusters) - - @property - def noise_fraction(self) -> float: - return float((self.labels == -1).mean()) if self.labels.size else float("nan") - - -def _region_labels( - cache: CellCache, regions: Sequence[str], labels_of: Optional[Dict[str, int]] -) -> Dict[str, int]: - out: Dict[str, int] = {} - for r in regions: - if labels_of is not None and r in labels_of: - out[r] = int(labels_of[r]) - else: - v = cache.label_of(r) - if v is None: - raise ValueError(f"No label for region {r}: pass labels_of or a labelled cache.") - out[r] = int(round(float(v))) - return out - - -def run_niche_discovery( - cache: CellCache, - *, - regions: Optional[Sequence[str]] = None, - labels_of: Optional[Dict[str, int]] = None, - space: str = "gnn", - attention_ckpt: Optional[str] = None, - reduction: Optional[Dict[str, Any]] = None, - clustering: Optional[Dict[str, Any]] = None, - sample_size: Optional[int] = 100_000, - l2_normalise: bool = True, - assign_k: int = 15, - seed: int = 0, - positive_class: int = 1, - present_threshold: float = 0.02, - n_boot: int = 2000, -) -> NicheResult: - if regions is None: - regions = cache.held_out or cache.regions - regions = sorted(str(r) for r in regions) - cells = load_cells(cache, regions) - region_labels = _region_labels(cache, regions, labels_of) - if cells.AR is None: - raise ValueError( - "The cache has no attention logits (AR); niche margins need an attention model." - ) - rel = relattn_margins( - cells.LG, cells.AR, cells.bias, cells.region_ids, positive_class=positive_class - ) - space = str(space).lower() - if space in {"gnn", "encoder", "embedding"}: - X = cells.E - elif space in {"attention", "attn"}: - if not attention_ckpt: - raise ValueError("space='attention' needs attention_ckpt (the model checkpoint).") - X = attention_space_from_checkpoint(cells.E, attention_ckpt) - else: - raise ValueError("space must be 'gnn' or 'attention'") - model, labels = fit_niche_clustering( - X, - reduction=reduction or DEFAULT_REDUCTION, - clustering=clustering or DEFAULT_CLUSTERING, - sample_size=sample_size, - l2_normalise=l2_normalise, - assign_k=assign_k, - seed=seed, - ) - table = niche_table( - labels, - cell_types=cells.ct, - region_ids=cells.region_ids, - region_labels=region_labels, - margin=rel.margin, - attention_pos=rel.attention[:, min(positive_class, rel.attention.shape[1] - 1)], - cell_type_names=cells.cell_type_names, - present_threshold=present_threshold, - n_boot=n_boot, - seed=seed, - ) - CO, clusters = cluster_colocalisation(labels, cells.edge_index) - y_cell = np.array([region_labels[r] for r in cells.region_ids]) - return NicheResult( - labels=labels, - model=model, - table=table, - margin=rel.margin, - attention=rel.attention, - colocalisation=CO, - clusters=clusters, - concordance=spearman_concordance(table), - cells=cells, - features=X, - settings={ - "space": space, - "regions": regions, - "n_cells": int(labels.size), - "positive_class": int(positive_class), - **model.params, - }, - y_cell=y_cell, - ) - - -def write_niche_outputs( - result: NicheResult, - out_dir: str | Path, - *, - title: str = "niche discovery", - positive_name: str = "positive", - figures: bool = True, - umap: Optional[Dict[str, Any]] = None, - held_out_auroc: Optional[float] = None, -) -> Dict[str, str]: - """Write labels, cluster model, tables, report and figures; returns the paths.""" - import joblib - - out = Path(out_dir) - out.mkdir(parents=True, exist_ok=True) - paths: Dict[str, str] = {} - np.save(out / "labels.npy", result.labels) - paths["labels"] = str(out / "labels.npy") - joblib.dump(result.model, out / "cluster_model.pkl") - paths["cluster_model"] = str(out / "cluster_model.pkl") - cells = pd.DataFrame( - { - "region_id": result.cells.region_ids, - "cell_index": np.concatenate( - [np.arange((result.cells.region_ids == r).sum()) for r in result.cells.regions] - ), - "cluster": result.labels, - "cell_type": [result.cells.cell_type_names[c] for c in result.cells.ct], - "relattn": result.margin, - "x": result.cells.pos[:, 0], - "y": result.cells.pos[:, 1], - } - ) - cells.to_csv(out / "cells.csv", index=False) - paths["cells"] = str(out / "cells.csv") - result.table.to_csv(out / "clusters.csv", index=False) - paths["clusters"] = str(out / "clusters.csv") - np.save(out / "coloc_log2.npy", result.colocalisation) - (out / "settings.json").write_text(json.dumps(result.settings, indent=1, default=str)) - - T = result.table - base = float(result.y_cell.mean()) if result.y_cell is not None else float("nan") - n_pos = sum( - 1 - for r in result.cells.regions - if result.y_cell is not None and result.y_cell[result.cells.region_ids == r][0] == 1 - ) - md = [ - f"# {title}", - "", - f"{len(result.cells.regions)} regions / {result.labels.size} cells; space = {result.settings['space']}; " - f"clustering {json.dumps(result.model.params['clustering'])} on {json.dumps(result.model.params['reduction'])} " - f"(sample {result.model.params['sample_size']}, L2 {result.model.params['l2_normalise']}, {result.model.params['assign_k']}-NN assignment).", - f"Base rate (cells from {positive_name} regions) {base:.3f}; {positive_name} regions {n_pos}/{len(result.cells.regions)}" - + (f"; held-out AUROC {held_out_auroc:.3f}" if held_out_auroc is not None else "") - + f"; {result.n_clusters} clusters, noise {result.noise_fraction:.1%}.", - "", - f"| Cluster | Share | Regions (pos/neg, >{100 * 0.02:.0f}%) | Composition | Enriched (>1.5x) | Emp. frac. pos | Share AUROC | Margin relattn [95% CI] | Signed share | Sign consist. | Lift |", - "|---|---|---|---|---|---|---|---|---|---|---|", - ] - for _, t in T.sort_values("emp_frac_pos", ascending=False).iterrows(): - md.append( - f"| {t.cluster} | {100 * t.share:.1f}% | {t.cores_present} ({t.pos_cores_present}/{t.neg_cores_present}) | {t.top_types} | {t.top_enriched} | " - f"{t.emp_frac_pos:.2f} | {t.share_auroc:.2f} | {t.relattn:+.2f} [{t.relattn_lo:+.2f}, {t.relattn_hi:+.2f}] | {t.signed_share:+.2f} | {t.sign_consistency:.2f} | {t.lift:.2f} |" - ) - conc = result.concordance - if "spearman_emp_relattn" in conc: - md.append( - f"\nSpearman(emp_frac_pos, relattn) over {conc['n_clusters']} clusters: {conc['spearman_emp_relattn']:+.2f}" - + ( - f"; Spearman(share AUROC, relattn): {conc['spearman_share_auroc_relattn']:+.2f}" - if "spearman_share_auroc_relattn" in conc - else "" - ) - ) - CO, CL = result.colocalisation, result.clusters - n = len(CL) - if n >= 2: - pairs = [(i, j) for i in range(n) for j in range(i + 1, n)] - md.append( - "\nCo-localisation log2(obs/exp) of cell-cell edges, top pairs: " - + "; ".join( - f"C{CL[i]}-C{CL[j]} {CO[i, j]:+.1f}" - for i, j in sorted(pairs, key=lambda p: -CO[p])[:8] - ) - ) - md.append( - "Most avoided pairs: " - + "; ".join( - f"C{CL[i]}-C{CL[j]} {CO[i, j]:+.1f}" - for i, j in sorted(pairs, key=lambda p: CO[p])[:5] - ) - ) - (out / "report.md").write_text("\n".join(md) + "\n") - paths["report"] = str(out / "report.md") - - if figures: - try: - from .figures import ( - clusters_overview, - colocalisation_heatmap, - embedding_triptych, - umap_coordinates, - ) - except Exception: # pragma: no cover - matplotlib missing - return paths - samp = result.model.sample_idx - P = result.model.reduce(result.features[samp]) - y_s = result.y_cell[samp] if result.y_cell is not None else np.zeros(samp.size) - if P.shape[1] >= 2: - ev = getattr(result.model.reducer, "explained_variance_ratio_", None) - axis = ( - (f"PC1 ({100 * ev[0]:.1f}% var)", f"PC2 ({100 * ev[1]:.1f}% var)") - if ev is not None - else ("dim 1", "dim 2") - ) - embedding_triptych( - P[:, :2], - result.labels[samp], - result.margin[samp], - y_s, - out / "pca_triptych.png", - title=f"{result.settings['space']} (PCA)", - axis_labels=axis, - positive_name=positive_name, - ) - paths["pca_triptych"] = str(out / "pca_triptych.png") - ucfg = dict(umap or {}) - if ucfg.get("enabled", True): - U = umap_coordinates( - P, - n_neighbors=int(ucfg.get("n_neighbors", 30)), - min_dist=float(ucfg.get("min_dist", 0.1)), - metric=str(ucfg.get("metric", "cosine")), - seed=int(ucfg.get("seed", 0)), - ) - if U is not None: - embedding_triptych( - U, - result.labels[samp], - result.margin[samp], - y_s, - out / "umap_triptych.png", - title=f"{result.settings['space']} (UMAP)", - axis_labels=("UMAP 1", "UMAP 2"), - positive_name=positive_name, - ) - paths["umap_triptych"] = str(out / "umap_triptych.png") - cohort = np.bincount(result.cells.ct, minlength=len(result.cells.cell_type_names)) / max( - 1, result.cells.ct.size - ) - clusters_overview( - T, - out / "clusters_overview.png", - cell_type_names=result.cells.cell_type_names, - cohort_frequency=cohort, - base_rate=base, - positive_name=positive_name, - ) - paths["clusters_overview"] = str(out / "clusters_overview.png") - colocalisation_heatmap(CO, CL, out / "colocalisation.png") - paths["colocalisation"] = str(out / "colocalisation.png") - return paths diff --git a/src/grass_mil/interpretability/niches/tables.py b/src/grass_mil/interpretability/niches/tables.py deleted file mode 100644 index 6326508..0000000 --- a/src/grass_mil/interpretability/niches/tables.py +++ /dev/null @@ -1,131 +0,0 @@ -"""Per-niche summary table (the preprint's cluster table).""" - -from __future__ import annotations - -from typing import Dict, List, Optional, Sequence - -import numpy as np -import pandas as pd - -__all__ = ["niche_table", "spearman_concordance", "bootstrap_mean"] - - -def bootstrap_mean(values, *, n_boot: int = 2000, seed: int = 0): - v = np.asarray(values, dtype=float) - v = v[np.isfinite(v)] - if v.size == 0: - return float("nan"), float("nan"), float("nan") - g = np.random.default_rng(seed) - bs = [g.choice(v, v.size, replace=True).mean() for _ in range(int(n_boot))] - return float(v.mean()), float(np.percentile(bs, 2.5)), float(np.percentile(bs, 97.5)) - - -def niche_table( - labels: np.ndarray, - *, - cell_types: np.ndarray, - region_ids: np.ndarray, - region_labels: Dict[str, int], - margin: np.ndarray, - attention_pos: np.ndarray, - cell_type_names: Optional[Sequence[str]] = None, - present_threshold: float = 0.02, - enrichment_threshold: float = 1.5, - n_boot: int = 2000, - seed: int = 0, - include_noise: bool = True, -) -> pd.DataFrame: - """One row per niche: share, cores present, composition, enrichment, empirical - fraction of cells from positive regions, share-AUROC, relattn margin with a - bootstrap CI over regions, signed share, sign consistency and attention lift. - """ - from sklearn.metrics import roc_auc_score - - labels = np.asarray(labels).astype(int) - ct = np.asarray(cell_types).astype(int) - rid = np.asarray(region_ids) - M = np.asarray(margin, dtype=float) - A1 = np.asarray(attention_pos, dtype=float) - K = int(ct.max()) + 1 if ct.size else 0 - if cell_type_names is not None: - K = max(K, len(cell_type_names)) - names = list(cell_type_names) if cell_type_names is not None else [str(i) for i in range(K)] - regions = sorted(set(rid.tolist())) - y_region = np.array([int(region_labels[r]) for r in regions]) - y_cell = np.array([int(region_labels[r]) for r in rid]) - cohort = np.bincount(ct, minlength=K) / max(1, ct.size) - region_mask = {r: rid == r for r in regions} - region_size = {r: int(m.sum()) for r, m in region_mask.items()} - region_margin = {r: float(M[m].sum()) for r, m in region_mask.items()} - clusters = sorted(set(labels.tolist()) - {-1}) - if include_noise and (labels == -1).any(): - clusters = clusters + [-1] - rows: List[dict] = [] - for cl in clusters: - msk = labels == cl - comp = np.bincount(ct[msk], minlength=K) / max(1, msk.sum()) - enr = comp / np.maximum(cohort, 1e-9) - vals, shares, signs, lifts = [], [], [], [] - per_region_share = np.zeros(len(regions)) - for j, r in enumerate(regions): - mm = msk & region_mask[r] - s = mm.sum() / max(1, region_size[r]) - per_region_share[j] = s - if mm.any(): - vals.append(M[mm].mean()) - mcr = float(M[mm].sum()) - Mr = region_margin[r] - if abs(Mr) > 1e-6: - shares.append(mcr / Mr) - signs.append(float(np.sign(mcr) == np.sign(Mr))) - lifts.append(A1[mm].sum() / s) - present = per_region_share > present_threshold - m_, lo, hi = bootstrap_mean(vals, n_boot=n_boot, seed=seed) - ss = bootstrap_mean(shares, n_boot=n_boot, seed=seed)[0] if shares else float("nan") - both = 0 < y_region.sum() < y_region.size - rows.append( - dict( - cluster="Noise" if cl == -1 else f"C{cl}", - cluster_id=int(cl), - share=float(msk.mean()), - n_cells=int(msk.sum()), - cores_present=int(present.sum()), - pos_cores_present=int((present & (y_region == 1)).sum()), - neg_cores_present=int((present & (y_region == 0)).sum()), - emp_frac_pos=float(y_cell[msk].mean()) if msk.any() else float("nan"), - share_auroc=float(roc_auc_score(y_region, per_region_share)) - if both - else float("nan"), - relattn=m_, - relattn_lo=lo, - relattn_hi=hi, - signed_share=ss, - sign_consistency=float(np.mean(signs)) if signs else float("nan"), - lift=float(np.mean(lifts)) if lifts else float("nan"), - top_types=", ".join( - f"{names[j]} {100 * comp[j]:.0f}%" for j in np.argsort(-comp)[:4] - ), - top_enriched=", ".join( - f"{names[j]} x{enr[j]:.1f}" - for j in np.argsort(-enr)[:4] - if enr[j] > enrichment_threshold - ), - **{f"comp_{names[j]}": float(comp[j]) for j in range(K)}, - ) - ) - return pd.DataFrame(rows) - - -def spearman_concordance(table: pd.DataFrame) -> Dict[str, float]: - """Spearman correlations between the empirical association and the margin.""" - from scipy.stats import spearmanr - - tt = table[table["cluster"] != "Noise"] - out = {"n_clusters": int(len(tt))} - if len(tt) >= 3: - out["spearman_emp_relattn"] = float(spearmanr(tt["emp_frac_pos"], tt["relattn"])[0]) - if tt["share_auroc"].notna().sum() >= 3: - out["spearman_share_auroc_relattn"] = float( - spearmanr(tt["share_auroc"], tt["relattn"], nan_policy="omit")[0] - ) - return out diff --git a/src/grass_mil/interpretability/niches_cli.py b/src/grass_mil/interpretability/niches_cli.py deleted file mode 100644 index 09dc887..0000000 --- a/src/grass_mil/interpretability/niches_cli.py +++ /dev/null @@ -1,105 +0,0 @@ -"""``grass-mil-niches``: niche discovery on one model's cell cache.""" - -from __future__ import annotations - -import json -from pathlib import Path -from typing import Any, Dict, Optional - -import hydra -import pandas as pd -from omegaconf import DictConfig, OmegaConf - -from grass_mil._paths import setup_project_root - -setup_project_root() - -from grass_mil.inference.cell_cache import CellCache # noqa: E402 -from grass_mil.interpretability.niches.pipeline import run_niche_discovery, write_niche_outputs # noqa: E402 -from grass_mil.utils import RankedLogger # noqa: E402 - -log = RankedLogger(__name__, rank_zero_only=True) - -__all__ = ["run_niches", "main", "resolve_regions", "load_region_labels"] - - -def resolve_regions(value, cache: CellCache): - if value is None or (isinstance(value, str) and value.lower() == "held_out"): - return cache.held_out or cache.regions - if isinstance(value, str) and value.lower() == "all": - return cache.regions - if isinstance(value, str): - p = Path(value) - if p.suffix.lower() == ".json": - return [str(v) for v in json.loads(p.read_text())] - return [line.strip() for line in p.read_text().splitlines() if line.strip()] - return [str(v) for v in value] - - -def load_region_labels(cfg_labels) -> Optional[Dict[str, int]]: - if not cfg_labels or not cfg_labels.get("label_file"): - return None - df = pd.read_csv(str(cfg_labels.label_file)) - df.columns = [str(c).strip() for c in df.columns] - idc = str(cfg_labels.get("id_column", "REGION_ID")) - lc = cfg_labels.get("label_column") - if not lc: - raise ValueError("labels.label_column is required with labels.label_file") - df = df.dropna(subset=[str(lc)]) - return {str(r): int(float(v) >= 0.5) for r, v in zip(df[idc], df[str(lc)])} - - -def run_niches(cfg: DictConfig) -> Dict[str, Any]: - cache = CellCache.open(str(cfg.cells_dir)) - regions = resolve_regions(cfg.get("regions"), cache) - clustering = OmegaConf.to_container(cfg.clustering, resolve=True) - reduction = OmegaConf.to_container(cfg.reduction, resolve=True) - result = run_niche_discovery( - cache, - regions=regions, - labels_of=load_region_labels(cfg.get("labels")), - space=str(cfg.get("space", "gnn")), - attention_ckpt=cfg.get("attention_ckpt") or cache.index.get("ckpt_path"), - reduction=reduction, - clustering=clustering, - sample_size=cfg.get("sample_size"), - l2_normalise=bool(cfg.get("l2_normalise", True)), - assign_k=int(cfg.get("assign_k", 15)), - seed=int(cfg.get("seed", 0)), - positive_class=int(cfg.get("positive_class", 1)), - present_threshold=float(cfg.get("present_threshold", 0.02)), - n_boot=int(cfg.get("n_boot", 2000)), - ) - title = ( - cfg.get("title") - or f"Niche discovery: {cache.index.get('model_name', Path(str(cfg.cells_dir)).name)}" - ) - paths = write_niche_outputs( - result, - str(cfg.output_dir), - title=str(title), - positive_name=str(cfg.get("positive_name", "positive")), - figures=bool(cfg.get("figures", True)), - umap=OmegaConf.to_container(cfg.umap, resolve=True) - if cfg.get("umap") is not None - else None, - held_out_auroc=cfg.get("held_out_auroc"), - ) - log.info( - f"{result.n_clusters} clusters, noise {result.noise_fraction:.1%}; report: {paths['report']}" - ) - return { - "n_clusters": result.n_clusters, - "noise_fraction": result.noise_fraction, - "concordance": result.concordance, - "paths": paths, - } - - -@hydra.main(version_base="1.3", config_path="../configs", config_name="niches.yaml") -def main(cfg: DictConfig) -> None: - run_niches(cfg) - - -if __name__ == "__main__": - main() diff --git a/src/grass_mil/interpretability/pipeline.py b/src/grass_mil/interpretability/pipeline.py deleted file mode 100644 index 2fb1cde..0000000 --- a/src/grass_mil/interpretability/pipeline.py +++ /dev/null @@ -1,302 +0,0 @@ -from __future__ import annotations - -import logging - -from dataclasses import replace - -import json -from pathlib import Path -from typing import Any, Dict, Optional - -import numpy as np -import pandas as pd - -log = logging.getLogger(__name__) - -from grass_mil.contracts import ( - NicheSummary, - InterpretabilityDataset, - PluginResult, - ReportBundle, -) -from grass_mil.interpretability.core.biomarkers import niche_composition_summary -from grass_mil.interpretability.core.niches import relabel_niches_by_composition -from grass_mil.contracts import ClusteringResult -from grass_mil.interpretability.core.clustering import run_clustering -from grass_mil.interpretability.core.data import extract_embedding_set -from grass_mil.interpretability.core.reduction import run_reduction -from grass_mil.interpretability.plugins.base import PluginContext -from grass_mil.interpretability.plugins.builtin import create_builtin_registry -from grass_mil.interpretability.plugins.registry import PluginRegistry - - -def _is_plugin_input_available( - dataset: InterpretabilityDataset, - context: PluginContext, - input_name: str, -) -> bool: - key = str(input_name).strip() - if not key: - return False - if key in context.state: - return context.state[key] is not None - if hasattr(dataset, key): - return getattr(dataset, key) is not None - return False - - -def _validate_plugin_required_inputs( - plugin_name: str, - required_inputs: list[str], - dataset: InterpretabilityDataset, - context: PluginContext, -) -> None: - required = [str(name).strip() for name in required_inputs if str(name).strip()] - missing = [ - name - for name in required - if not _is_plugin_input_available(dataset=dataset, context=context, input_name=name) - ] - if missing: - missing_text = ", ".join(sorted(set(missing))) - raise ValueError( - f"Plugin {plugin_name!r} is missing required inputs: {missing_text}. " - "Provide these inputs in dataset fields or plugin context state." - ) - - -def _write_table(df: pd.DataFrame, path: Path) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - if path.suffix.lower() == ".parquet": - df.to_parquet(path, index=False) - else: - df.to_csv(path, index=False) - - -def _write_json(payload: Dict[str, Any], path: Path) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(payload, indent=2, default=str)) - - -def _resolve_adaptive_min_cluster_size( - params: Dict[str, Any], *, n_instances: int, method: str -) -> Dict[str, Any]: - """Scale `min_cluster_size` to the cohort unless it was set explicitly. - - A fixed value does not transfer between runs: 100 is a reasonable floor for - ~20k instances and swallows every niche at 2k. `auto` (the default) uses - 0.5% of the instances being clustered, which keeps the smallest admissible - niche a constant *fraction* of the sample. - """ - if method.strip().lower() not in {"hdbscan", "optics"}: - return params - value = params.get("min_cluster_size", "auto") - if value not in (None, "auto"): - return params - resolved = max(2, int(round(0.005 * float(n_instances)))) - out = dict(params) - out["min_cluster_size"] = resolved - log.info( - "clustering.params.min_cluster_size=auto -> %d (0.5%% of %d instances)", - resolved, - n_instances, - ) - return out - - -def run_interpretability_pipeline( - dataset: InterpretabilityDataset, - config: Dict[str, Any], - *, - artifacts_dir: Path, - plugin_registry: Optional[PluginRegistry] = None, -) -> ReportBundle: - artifacts_dir = Path(artifacts_dir) - artifacts_dir.mkdir(parents=True, exist_ok=True) - resolved_registry = ( - plugin_registry if plugin_registry is not None else create_builtin_registry() - ) - - emb_set = extract_embedding_set( - dataset, - embedding_prefixes=tuple( - config.get("embedding_prefixes", ("inst_emb_", "emb_", "graph_emb_")) - ), - ) - - clustering_cfg = dict(config.get("clustering", {})) - clustering_enabled = bool(clustering_cfg.get("enabled", True)) - cluster_on_pca = bool(clustering_cfg.get("cluster_on_pca", False)) - cluster_pca_components_raw = clustering_cfg.get("pca_components", None) - cluster_pca_components = ( - int(cluster_pca_components_raw) if cluster_pca_components_raw is not None else None - ) - effective_cluster_pca_components = cluster_pca_components - if clustering_enabled and cluster_on_pca: - if cluster_pca_components is None: - raise ValueError("clustering.cluster_on_pca=true requires clustering.pca_components.") - max_pca_components = int(min(emb_set.matrix.shape[0], emb_set.matrix.shape[1])) - if max_pca_components < 1: - raise ValueError("Cannot run PCA-based clustering with empty embedding matrix.") - effective_cluster_pca_components = min(cluster_pca_components, max_pca_components) - - reduction_cfg = dict(config.get("reduction", {})) - reduction_enabled = bool(reduction_cfg.get("enabled", True)) - reduction_result = None - combo_pca_result = None - if reduction_enabled: - method = str(reduction_cfg.get("method", "pca")) - params = dict(reduction_cfg.get("params", {})) - reduction_input = emb_set.matrix - if method.strip().lower() == "umap" and cluster_on_pca: - combo_pca_result = run_reduction( - "pca", - emb_set.matrix, - n_components=effective_cluster_pca_components, - random_state=None, - ) - reduction_input = combo_pca_result.embedding - reduction_result = run_reduction(method, reduction_input, **params) - - clustering_result = None - labels: Optional[np.ndarray] = None - if clustering_enabled: - method = str(clustering_cfg.get("method", "kmeans")) - params = dict(clustering_cfg.get("params", {})) - # Contract: when cluster_on_pca is disabled, niche on raw embeddings. - # Reduction remains available for visualization/reporting, but does not - # change the clustering feature space. - clustering_input = emb_set.matrix - if cluster_on_pca: - if combo_pca_result is None: - combo_pca_result = run_reduction( - "pca", - emb_set.matrix, - n_components=effective_cluster_pca_components, - random_state=None, - ) - clustering_input = combo_pca_result.embedding - params = _resolve_adaptive_min_cluster_size( - params, n_instances=int(clustering_input.shape[0]), method=method - ) - sample_size = clustering_cfg.get("sample_size") - l2_normalise = bool(clustering_cfg.get("l2_normalise", False)) - if sample_size is not None or l2_normalise: - # Preprint recipe: cluster a sample on L2-normalised coordinates and - # assign the remaining instances by k-NN (see core.clustering). - from grass_mil.interpretability.core.clustering import fit_niche_clustering - - cl_params = { - k: v for k, v in params.items() if k not in {"n_clusters", "random_state"} - } - niche_model, labels = fit_niche_clustering( - clustering_input, - reduction={"method": "none"}, - clustering={ - "method": method, - "params": cl_params, - "n_clusters": params.get("n_clusters"), - "random_state": params.get("random_state", 42), - }, - sample_size=int(sample_size) if sample_size is not None else None, - l2_normalise=l2_normalise, - assign_k=int(clustering_cfg.get("assign_k", 15)), - seed=int(clustering_cfg.get("seed", 0)), - ) - clustering_result = ClusteringResult( - method=str(method).lower(), - labels=np.asarray(labels), - params={**params, **niche_model.params}, - fitted_object=niche_model, - ) - else: - clustering_result = run_clustering(method, clustering_input, **params) - labels = clustering_result.labels - - niche_relabel_map: Dict[Any, Any] = {} - if labels is not None and bool(config.get("relabel_niches_by_composition", True)): - # Renumber once so the niche id carries the dendrogram order. Everything - # downstream then gets composition order by sorting on the id. - labels, niche_relabel_map = relabel_niches_by_composition( - dataset.instance_table, - labels, - composition_prefix=str(config.get("composition_prefix", "comp_")), - ) - if clustering_result is not None: - clustering_result = replace(clustering_result, labels=labels) - - niche_summary: Optional[NicheSummary] = None - if labels is not None: - niche_summary = niche_composition_summary( - dataset.instance_table, - labels, - composition_prefix=str(config.get("composition_prefix", "comp_")), - variance_estimator=str(config.get("variance_estimator", "unbiased")), - ) - - context = PluginContext( - state={ - "reduction": reduction_result, - "clustering": clustering_result, - "niche_labels": labels, - "niche_summary": niche_summary, - } - ) - plugin_results: Dict[str, PluginResult] = {} - plugin_cfg = dict(config.get("plugins", {})) - enabled_plugins = list(plugin_cfg.get("enabled", [])) - plugin_params = dict(plugin_cfg.get("params", {})) - - for plugin_name in enabled_plugins: - plugin = resolved_registry.get(str(plugin_name)) - _validate_plugin_required_inputs( - plugin_name=str(plugin_name), - required_inputs=list(plugin.required_inputs()), - dataset=dataset, - context=context, - ) - result = plugin.run(dataset, context, **dict(plugin_params.get(plugin_name, {}))) - plugin_results[str(plugin_name)] = result - - # Persist core artifacts - base_instance = dataset.instance_table.copy() - if labels is not None: - base_instance["niche_label"] = labels - _write_table(base_instance, artifacts_dir / "instance_with_niches.csv") - if reduction_result is not None: - red_df = pd.DataFrame(reduction_result.embedding) - red_df.columns = [f"reduced_{i}" for i in range(red_df.shape[1])] - red_df.insert( - 0, dataset.id_column, dataset.instance_table[dataset.id_column].astype(str).tolist() - ) - _write_table(red_df, artifacts_dir / "reduction.csv") - if clustering_result is not None: - niche_df = pd.DataFrame( - { - dataset.id_column: dataset.instance_table[dataset.id_column].astype(str).tolist(), - "niche_label": clustering_result.labels, - } - ) - _write_table(niche_df, artifacts_dir / "niche_labels.csv") - - summary_payload: Dict[str, Any] = { - "rows": int(len(dataset.instance_table)), - "reduction_method": reduction_result.method if reduction_result else None, - "clustering_method": clustering_result.method if clustering_result else None, - "cluster_on_pca": cluster_on_pca, - "cluster_pca_components": effective_cluster_pca_components, - "enabled_plugins": enabled_plugins, - "niche_relabel_map": {str(k): int(v) for k, v in niche_relabel_map.items()}, - } - _write_json(summary_payload, artifacts_dir / "pipeline_summary.json") - - return ReportBundle( - dataset=dataset, - reduction=reduction_result, - niche_feature_reduction=combo_pca_result if cluster_on_pca else None, - clustering=clustering_result, - niche_summary=niche_summary, - plugin_results=plugin_results, - artifacts_dir=artifacts_dir, - metadata=summary_payload, - ) diff --git a/src/grass_mil/interpretability/plugins/__init__.py b/src/grass_mil/interpretability/plugins/__init__.py deleted file mode 100644 index 30bab8c..0000000 --- a/src/grass_mil/interpretability/plugins/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -from grass_mil.interpretability.plugins.base import InterpretabilityPlugin, PluginContext -from grass_mil.interpretability.plugins.builtin import ( - create_builtin_registry, - register_builtin_plugins, -) -from grass_mil.interpretability.plugins.registry import PluginRegistry, create_plugin_registry - -__all__ = [ - "InterpretabilityPlugin", - "PluginContext", - "PluginRegistry", - "create_plugin_registry", - "create_builtin_registry", - "register_builtin_plugins", -] diff --git a/src/grass_mil/interpretability/plugins/base.py b/src/grass_mil/interpretability/plugins/base.py deleted file mode 100644 index d778118..0000000 --- a/src/grass_mil/interpretability/plugins/base.py +++ /dev/null @@ -1,21 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any, Dict, List, Protocol - -from grass_mil.contracts import InterpretabilityDataset, PluginResult - - -@dataclass -class PluginContext: - state: Dict[str, Any] = field(default_factory=dict) - - -class InterpretabilityPlugin(Protocol): - name: str - - def required_inputs(self) -> List[str]: ... - - def run( - self, dataset: InterpretabilityDataset, context: PluginContext, **params: Any - ) -> PluginResult: ... diff --git a/src/grass_mil/interpretability/plugins/builtin.py b/src/grass_mil/interpretability/plugins/builtin.py deleted file mode 100644 index 337cf18..0000000 --- a/src/grass_mil/interpretability/plugins/builtin.py +++ /dev/null @@ -1,689 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Dict, List, Optional - -import numpy as np -import pandas as pd - -from grass_mil.contracts import ( - NicheSummary, - PluginResult, - ReportSection, -) -from grass_mil.interpretability.core.agreement import compute_niche_agreement -from grass_mil.interpretability.core.attribution import niche_attribution_summary -from grass_mil.interpretability.core.biomarkers import niche_composition_summary -from grass_mil.interpretability.plugins.base import InterpretabilityPlugin, PluginContext -from grass_mil.interpretability.plugins.registry import PluginRegistry, create_plugin_registry -from grass_mil.interpretability.tier2.cell_level import ( - attach_niche_labels_to_cells, - differential_cell_type_enrichment, - cell_filtration_curves, - niche_label_moran, - per_niche_cell_type_enrichment, - per_niche_cell_type_moran, - per_niche_cell_type_ripley, -) -from grass_mil.interpretability.tier2.neighborhood import ( - run_diff_neighborhood_enrichment, -) -from grass_mil.interpretability.tier2.tissue_graph import prepare_tissue_graph_view - - -def _resolve_filtration_thresholds(params: Dict[str, Any]) -> np.ndarray: - explicit = params.get("thresholds") - if explicit is not None: - return np.asarray(explicit, dtype=float) - - # 500 bins from 0 to 55 micrometers. - start = float(params.get("threshold_start", 0.0)) - stop = float(params.get("threshold_stop", 55.0)) - count = int(params.get("threshold_count", 500)) - if count < 2: - raise ValueError("filtration threshold_count must be >= 2.") - if stop < start: - raise ValueError("filtration threshold_stop must be >= threshold_start.") - return np.linspace(start, stop, count, dtype=float) - - -def _context_cluster_summary( - context: PluginContext, - *, - labels: np.ndarray, -) -> NicheSummary | None: - summary = context.state.get("niche_summary") - if not isinstance(summary, NicheSummary): - return None - if summary.niche_labels.shape != labels.shape: - return None - if not np.array_equal(summary.niche_labels, labels): - return None - return summary - - -@dataclass -class NicheProfilesPlugin(InterpretabilityPlugin): - name: str = "niche_profiles" - - def required_inputs(self) -> List[str]: - return ["niche_labels"] - - def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: - labels = np.asarray(context.state["niche_labels"]) - summary = _context_cluster_summary(context, labels=labels) - if summary is None: - summary = niche_composition_summary( - dataset.instance_table, - labels, - composition_prefix=str(params.get("composition_prefix", "comp_")), - variance_estimator=str(params.get("variance_estimator", "unbiased")), - ) - payload = { - "composition": summary.composition, - "enrichment": summary.enrichment, - "counts": summary.niche_counts, - } - sections = [ - ReportSection( - title="Niche Profiles", - description="Niche-level composition and enrichment summary.", - tables={ - "composition": summary.composition, - "enrichment": summary.enrichment, - }, - ) - ] - return PluginResult(name=self.name, payload=payload, sections=sections) - - -@dataclass -class TissueGraphPlugin(InterpretabilityPlugin): - name: str = "tissue_graph" - - def required_inputs(self) -> List[str]: - return ["spatial_table", "niche_labels"] - - def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: - if dataset.spatial_table is None: - raise ValueError("Spatial table is required for tissue graph plugin.") - sample_value = params.get("sample_value") - if sample_value is None or str(sample_value).strip() == "": - raise ValueError("tissue_graph plugin requires a non-empty sample_value parameter.") - - table = dataset.instance_table.copy() - table["niche_label"] = np.asarray(context.state["niche_labels"]) - sample_column = str(params.get("sample_column", "sample_id")) - id_column = str(params.get("id_column", dataset.id_column)) - x_column = str(params.get("x_column", "center_x")) - y_column = str(params.get("y_column", "center_y")) - label_column = str(params.get("label_column", "niche_label")) - - view = prepare_tissue_graph_view( - table, - dataset.spatial_table, - sample_column=sample_column, - sample_value=sample_value, - id_column=id_column, - x_column=x_column, - y_column=y_column, - label_column=label_column, - coerce_ids_to_str=True, - include_edge_distances=True, - ) - meta = dict(view.metadata) - niche_counts = meta.get("niche_counts", {}) - section_summary = pd.DataFrame( - [ - { - "sample_column": meta.get("sample_column"), - "sample_value": meta.get("sample_value"), - "node_count": meta.get("node_count"), - "edge_count": meta.get("edge_count"), - "niche_count": len(niche_counts) if isinstance(niche_counts, dict) else 0, - } - ] - ) - style = { - "show_edges": bool(params.get("show_edges", True)), - "node_size": float(params.get("node_size", 5.0)), - "edge_width": float(params.get("edge_width", 0.5)), - "edge_opacity": float(params.get("edge_opacity", 0.25)), - "colorscale": str(params.get("colorscale", "Viridis")), - "reverse_y": bool(params.get("reverse_y", True)), - "title_prefix": str(params.get("title_prefix", "Tissue Graph")), - } - payload = { - "tissue_graph_view": view, - "tissue_graph_meta": meta, - "tissue_graph_style": style, - } - sections = [ - ReportSection( - title=f"Tissue Graph ({sample_column}={sample_value})", - description="Single tissue/sample graph with niche-colored cells.", - tables={"summary": section_summary}, - metadata=meta, - ) - ] - return PluginResult(name=self.name, payload=payload, sections=sections) - - -def _resolve_columns(table: pd.DataFrame, prefix: str, explicit) -> List[str]: - if explicit: - return [str(c) for c in explicit] - return [c for c in table.columns if str(c).startswith(prefix)] - - -def _resolve_logit_bias(params: Dict[str, Any]) -> Optional[List[float]]: - """Head output bias, read from a checkpoint unless given explicitly. - - Subtracting it makes the margin reflect the instance-driven part of the - decision rather than the head's prior, so it should normally come straight - from the trained model rather than being transcribed by hand. - """ - explicit = params.get("logit_bias") - if explicit is not None: - return [float(v) for v in explicit] - - checkpoint = params.get("logit_bias_checkpoint") - if not checkpoint: - return None - - from grass_mil.models.training.checkpoint_init import read_graph_head_bias - - bias = read_graph_head_bias( - str(checkpoint), head_prefix=str(params.get("head_prefix", "graph_head")) - ) - if bias is None: - raise ValueError( - f"No graph head bias found in {checkpoint!r}. Check head_prefix, or set " - "logit_bias explicitly if the checkpoint holds only an encoder." - ) - return bias - - -def _bag_logits_from_table(bag_table, *, bag_id_column: str, logit_prefix: str): - """``{bag_id: [logit_0, logit_1, ...]}`` from a predictions table, or None. - - With the bag logits available the attribution reports a real - ``identity_residual`` (max |sum_i A[i,c] l[i,c] - L_c| over bags) instead - of NaN, so the additive identity is checked rather than assumed. - """ - if bag_table is None or bag_id_column not in bag_table.columns: - return None - import re - - pattern = re.compile(rf"^{re.escape(logit_prefix)}(\d+)$") - cols = sorted( - (c for c in bag_table.columns if pattern.match(str(c))), - key=lambda c: int(pattern.match(str(c)).group(1)), - ) - if not cols: - return None - return { - str(row[bag_id_column]): [float(row[c]) for c in cols] - for _, row in bag_table[[bag_id_column, *cols]].iterrows() - } - - -@dataclass -class MarginAttributionPlugin(InterpretabilityPlugin): - """Exact additive attribution of bag decisions to instance niches. - - Requires per-class attention columns, which only exist when the head emits - one attention channel per class. - """ - - name: str = "margin_attribution" - - def required_inputs(self) -> List[str]: - return ["niche_labels"] - - def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: - table = dataset.instance_table - labels = np.asarray(context.state["niche_labels"]) - - attention_columns = _resolve_columns( - table, - str(params.get("attention_prefix", "attention_c")), - params.get("attention_columns"), - ) - if not attention_columns: - attention_columns = [str(params.get("attention_column", "attention"))] - logit_columns = _resolve_columns( - table, str(params.get("logit_prefix", "logit_")), params.get("logit_columns") - ) - if not logit_columns: - raise ValueError( - "Margin attribution needs instance logit columns (default prefix 'logit_'). " - "Enable the interpretability export in the predict step." - ) - - bag_id_column = str(params.get("bag_id_column", dataset.bag_id_column)) - result = niche_attribution_summary( - table, - labels, - attention_columns=attention_columns, - logit_columns=logit_columns, - bag_id_column=bag_id_column, - logit_bias=_resolve_logit_bias(params), - n_bootstrap=int(params.get("n_bootstrap", 200)), - random_state=int(params.get("random_state", 0)), - margin_eps=float(params.get("margin_eps", 1e-6)), - bag_logits=_bag_logits_from_table( - getattr(dataset, "bag_table", None), - bag_id_column=str(params.get("bag_table_id_column", "bag_id")), - logit_prefix=str(params.get("bag_logit_prefix", "logit_")), - ), - focus_classes=params.get("focus_classes"), - ) - payload = { - "per_niche": result.per_niche, - "identity_residual": result.identity_residual, - } - sections = [ - ReportSection( - title="Margin Attribution", - description=( - "Exact additive contributions M[i,c] = A[i,c] * l[i,c], summarised per " - "niche with percentile bootstrap intervals over regions." - ), - tables={"per_niche": result.per_niche}, - ) - ] - return PluginResult(name=self.name, payload=payload, sections=sections) - - -@dataclass -class NicheAgreementPlugin(InterpretabilityPlugin): - """Agreement between the active partition and previously stored ones. - - Extra labelings are read from columns of the instance table, so a user - joins a prior run's niche_labels.csv rather than re-running clustering. - """ - - name: str = "niche_agreement" - - def required_inputs(self) -> List[str]: - return ["niche_labels"] - - def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: - table = dataset.instance_table - labelings: Dict[str, Any] = { - str(params.get("active_name", "active")): np.asarray(context.state["niche_labels"]) - } - for column in params.get("label_columns", []) or []: - if column not in table.columns: - raise ValueError( - f"niche_agreement label column {column!r} is not in the instance table." - ) - labelings[str(column)] = table[column].to_numpy() - - if len(labelings) < 2: - raise ValueError( - "niche_agreement needs at least one additional labeling; set " - "params.niche_agreement.label_columns." - ) - - result = compute_niche_agreement( - labelings, - metrics=tuple(params.get("metrics", ("ari", "ami", "nmi"))), - include_jaccard=bool(params.get("include_jaccard", True)), - ) - payload = { - "pairwise_metrics": result.pairwise_metrics, - "jaccard": result.jaccard, - } - sections = [ - ReportSection( - title="Cross-Space Niche Agreement", - description=( - "Chance-corrected agreement between partitions of the same instances." - ), - tables={"pairwise_metrics": result.pairwise_metrics.reset_index()}, - ) - ] - return PluginResult(name=self.name, payload=payload, sections=sections) - - -def register_builtin_plugins(registry: PluginRegistry) -> None: - registry.register(NicheProfilesPlugin()) - registry.register(TissueGraphPlugin()) - registry.register(MarginAttributionPlugin()) - registry.register(NicheAgreementPlugin()) - registry.register(CellTypeEnrichmentPerNichePlugin()) - registry.register(CellTypeMoranPerNichePlugin()) - registry.register(NicheLabelMoranPlugin()) - registry.register(CellFiltrationCurvesPlugin()) - registry.register(CellTypeRipleyPerNichePlugin()) - registry.register(CellTypeDiffEnrichmentByConditionPlugin()) - - -def create_builtin_registry() -> PluginRegistry: - registry = create_plugin_registry() - register_builtin_plugins(registry) - return registry - - -def _cells_with_niches(dataset, context: PluginContext, niche_column: str): - """Cell table with each cell carrying its ego-graph's niche.""" - if dataset.cell_table is None or dataset.cell_edge_table is None: - raise ValueError( - "Cell-level analyses require cell_table and cell_edge_table. Export them " - "with `grass-mil-predict interpretability.cells.enabled=true`." - ) - return attach_niche_labels_to_cells( - dataset.cell_table, - list(np.asarray(context.state["niche_labels"])), - list(dataset.instance_table[dataset.id_column]), - instance_id_column=dataset.id_column, - niche_column=niche_column, - ) - - -@dataclass -class CellTypeEnrichmentPerNichePlugin(InterpretabilityPlugin): - """Cell-type x cell-type neighbourhood enrichment, within each niche. - - The instance-level `neighborhood_enrichment` treats a whole ego-graph as one - node, so it measures how neighbourhood *labels* co-occur. This measures how - the *cells* inside a niche are wired to each other. - """ - - name: str = "cell_type_enrichment_per_niche" - - def required_inputs(self) -> List[str]: - return ["cell_table", "cell_edge_table", "niche_labels"] - - def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: - niche_column = str(params.get("niche_column", "niche_label")) - cells = _cells_with_niches(dataset, context, niche_column) - result = per_niche_cell_type_enrichment( - cells, - dataset.cell_edge_table, - cell_type_column=str(params.get("cell_type_column", "cell_type")), - niche_column=niche_column, - cell_id_column=str(params.get("cell_id_column", dataset.cell_id_column)), - n_perms=int(params.get("n_perms", 0)), - random_state=int(params.get("random_state", 42)), - undirected=bool(params.get("undirected", False)), - enrichment_mode=str(params.get("enrichment_mode", "zscore")), - min_cells=int(params.get("min_cells", 10)), - skip_background=bool(params.get("skip_background", False)), - ) - tables = { - (f"niche_{k}" if k != -1 else "Background"): v for k, v in result.enrichment.items() - } - payload: Dict[str, Any] = { - "enrichment": result.enrichment, - "pvalues": result.pvalues, - "n_cells": result.n_cells, - "n_edges": result.n_edges, - } - if bool(params.get("differential_vs_background", True)): - differential = differential_cell_type_enrichment( - result, reference=int(params.get("reference_niche", -1)) - ) - payload["differential"] = differential - for k, v in differential.items(): - tables[f"niche_{k}_vs_Background"] = v - - sections = [ - ReportSection( - title="Cell-type neighbourhood enrichment, per niche", - description=( - "Cell-type x cell-type enrichment computed over the cells inside " - "each niche, using only edges whose both endpoints lie in that " - "niche. Unlike the instance-level analysis, this measures how " - "cells are wired to one another rather than how neighbourhood " - "labels co-occur." - ), - tables=tables, - ) - ] - return PluginResult(name=self.name, payload=payload, sections=sections) - - -@dataclass -class CellTypeMoranPerNichePlugin(InterpretabilityPlugin): - """Moran's I of each cell type, within each niche.""" - - name: str = "cell_type_moran_per_niche" - - def required_inputs(self) -> List[str]: - return ["cell_table", "cell_edge_table", "niche_labels"] - - def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: - niche_column = str(params.get("niche_column", "niche_label")) - cells = _cells_with_niches(dataset, context, niche_column) - result = per_niche_cell_type_moran( - cells, - dataset.cell_edge_table, - cell_type_column=str(params.get("cell_type_column", "cell_type")), - niche_column=niche_column, - cell_id_column=str(params.get("cell_id_column", dataset.cell_id_column)), - n_perms=int(params.get("n_perms", 100)), - random_state=int(params.get("random_state", 0)), - min_nodes=int(params.get("min_nodes", 3)), - ) - sections = [ - ReportSection( - title="Cell-type spatial autocorrelation, per niche", - description=( - "Moran's I of each cell-type indicator over the cell graph of " - "each niche: whether cells of that type cluster together " - "spatially inside the niche. Rows are niches, columns cell types." - ), - tables={"morans_i": result.statistic, "qvalue": result.qvalue}, - ) - ] - return PluginResult( - name=self.name, - payload={ - "statistic": result.statistic, - "pvalue": result.pvalue, - "qvalue": result.qvalue, - "n_nodes": result.n_nodes, - }, - sections=sections, - ) - - -@dataclass -class NicheLabelMoranPlugin(InterpretabilityPlugin): - """Moran's I of the niche assignment itself, over the instance graph.""" - - name: str = "niche_label_moran" - - def required_inputs(self) -> List[str]: - return ["spatial_table", "niche_labels"] - - def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: - if dataset.spatial_table is None: - raise ValueError("Niche-label Moran's I requires a spatial table.") - niche_column = str(params.get("niche_column", "niche_label")) - table = dataset.instance_table.copy() - table[niche_column] = np.asarray(context.state["niche_labels"]) - result = niche_label_moran( - table, - dataset.spatial_table, - niche_column=niche_column, - id_column=str(params.get("id_column", dataset.id_column)), - n_perms=int(params.get("n_perms", 100)), - random_state=int(params.get("random_state", 0)), - ) - sections = [ - ReportSection( - title="Niche spatial autocorrelation (global)", - description=( - "Moran's I of each niche indicator across the whole instance " - "graph: whether niches form contiguous territories rather than " - "interleaving. The global counterpart to the per-niche " - "cell-type analyses." - ), - tables={"morans_i": result.statistic, "qvalue": result.qvalue}, - ) - ] - return PluginResult( - name=self.name, - payload={ - "statistic": result.statistic, - "pvalue": result.pvalue, - "qvalue": result.qvalue, - }, - sections=sections, - ) - - -@dataclass -class CellFiltrationCurvesPlugin(InterpretabilityPlugin): - """Filtration curves over cell-cell edges, counting both endpoints.""" - - name: str = "cell_filtration_curves" - - def required_inputs(self) -> List[str]: - return ["cell_table", "cell_edge_table", "niche_labels"] - - def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: - niche_column = str(params.get("niche_column", "niche_label")) - cells = _cells_with_niches(dataset, context, niche_column) - thresholds = np.linspace( - float(params.get("threshold_start", 0.0)), - float(params.get("threshold_stop", 55.0)), - int(params.get("threshold_count", 500)), - ) - result = cell_filtration_curves( - cells, - dataset.cell_edge_table, - thresholds=thresholds, - cell_type_column=str(params.get("cell_type_column", "cell_type")), - niche_column=niche_column, - cell_id_column=str(params.get("cell_id_column", dataset.cell_id_column)), - scale_within_niche=bool(params.get("scale_within_niche", True)), - ) - sections = [ - ReportSection( - title="Cell-type filtration curves, per niche", - description=( - "Cells of each type reached as the edge-distance threshold " - "grows, computed over the edges inside the ego-graphs. Both " - "endpoints of every edge are counted, so the curves describe " - "cells rather than neighbourhoods. Distances are micrometres." - ), - metadata={"thresholds": result.thresholds.tolist()}, - ) - ] - return PluginResult( - name=self.name, - payload={"curves": result.curves, "thresholds": result.thresholds}, - sections=sections, - ) - - -@dataclass -class CellTypeRipleyPerNichePlugin(InterpretabilityPlugin): - """Ripley cross-L between cell types, with each niche as a region.""" - - name: str = "cell_type_ripley_per_niche" - - def required_inputs(self) -> List[str]: - return ["cell_table", "niche_labels"] - - def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: - niche_column = str(params.get("niche_column", "niche_label")) - cells = _cells_with_niches(dataset, context, niche_column) - result = per_niche_cell_type_ripley( - cells, - cell_type_column=str(params.get("cell_type_column", "cell_type")), - niche_column=niche_column, - x_column=str(params.get("x_column", "x")), - y_column=str(params.get("y_column", "y")), - n_radii=int(params.get("n_radii", 50)), - max_fraction=float(params.get("max_fraction", 0.25)), - min_count=int(params.get("min_count", 5)), - radius_source=str(params.get("radius_source", "median")), - ) - sections = [ - ReportSection( - title="Cell-type Ripley cross-L, per niche", - description=( - "Centred cross-L curves between cell types, using the cells' " - "own coordinates and treating each niche as a region. Positive " - "values indicate co-aggregation at that radius, negative " - "values mutual avoidance." - ), - tables={"ripley": result.to_frame()} if hasattr(result, "to_frame") else {}, - ) - ] - return PluginResult( - name=self.name, - payload={ - "curves": result.curves, - "radii": result.radii, - "pair_counts": result.pair_counts, - }, - sections=sections, - ) - - -@dataclass -class CellTypeDiffEnrichmentByConditionPlugin(InterpretabilityPlugin): - """Cell-type contact enrichment differenced between conditions. - - The cell-level counterpart of the removed instance-level differential: it - asks how the wiring between cell types changes between, say, responders and - non-responders, rather than how neighbourhood labels co-occur. - """ - - name: str = "cell_type_diff_enrichment_by_condition" - - def required_inputs(self) -> List[str]: - return ["cell_table", "cell_edge_table", "niche_labels"] - - def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: - condition_column = str( - params.get("condition_column", dataset.condition_column or "condition") - ) - niche_column = str(params.get("niche_column", "niche_label")) - cells = _cells_with_niches(dataset, context, niche_column) - - # The condition lives on instances (it is a bag-level property), so it - # is joined onto cells the same way the niche label is. - if condition_column not in cells.columns: - instance_conditions = dataset.instance_table[ - [dataset.id_column, condition_column] - ].drop_duplicates(subset=[dataset.id_column]) - cells = cells.merge(instance_conditions, on=dataset.id_column, how="left") - - results = run_diff_neighborhood_enrichment( - cells, - dataset.cell_edge_table, - label_column=str(params.get("cell_type_column", "cell_type")), - condition_column=condition_column, - permutation_group_column=str(params.get("permutation_group_column", "sample_id")), - id_column=str(params.get("cell_id_column", dataset.cell_id_column)), - n_perms=int(params.get("n_perms", 0)), - random_state=int(params.get("random_state", 42)), - undirected=bool(params.get("undirected", False)), - warn_analytical=False, - enrichment_mode=str(params.get("enrichment_mode", "zscore")), - ) - tables = {key: result.enrichment for key, result in results.items()} - sections = [ - ReportSection( - title="Cell-type enrichment differences between conditions", - description=( - "Cell-type x cell-type contact enrichment computed per " - "condition and differenced, over the cells inside the " - "ego-graphs. Positive entries mean the pair is in closer " - "contact in the first condition." - ), - tables=tables, - ) - ] - return PluginResult( - name=self.name, - payload={key: result.enrichment for key, result in results.items()}, - sections=sections, - ) diff --git a/src/grass_mil/interpretability/plugins/registry.py b/src/grass_mil/interpretability/plugins/registry.py deleted file mode 100644 index 0020a58..0000000 --- a/src/grass_mil/interpretability/plugins/registry.py +++ /dev/null @@ -1,40 +0,0 @@ -from __future__ import annotations - -from typing import Dict, Iterable, List - -from grass_mil.interpretability.plugins.base import InterpretabilityPlugin - - -class PluginRegistry: - def __init__(self, plugins: Iterable[InterpretabilityPlugin] | None = None) -> None: - self._plugins: Dict[str, InterpretabilityPlugin] = {} - if plugins is not None: - for plugin in plugins: - self.register(plugin) - - def register(self, plugin: InterpretabilityPlugin, *, replace: bool = False) -> None: - name = str(plugin.name).strip() - if not name: - raise ValueError("Plugin name must be non-empty.") - existing = self._plugins.get(name) - if existing is not None and not replace: - raise ValueError( - f"Plugin {name!r} is already registered. " - "Use replace=True to intentionally overwrite it." - ) - self._plugins[name] = plugin - - def get(self, name: str) -> InterpretabilityPlugin: - key = str(name).strip() - if key not in self._plugins: - raise KeyError(f"Plugin {name!r} is not registered.") - return self._plugins[key] - - def list(self) -> List[str]: - return sorted(self._plugins.keys()) - - -def create_plugin_registry( - plugins: Iterable[InterpretabilityPlugin] | None = None, -) -> PluginRegistry: - return PluginRegistry(plugins=plugins) diff --git a/src/grass_mil/interpretability/report_cli.py b/src/grass_mil/interpretability/report_cli.py deleted file mode 100644 index b5221d0..0000000 --- a/src/grass_mil/interpretability/report_cli.py +++ /dev/null @@ -1,81 +0,0 @@ -from __future__ import annotations - -import json -import logging -from pathlib import Path -from typing import Any, Dict, Tuple - -import hydra -from omegaconf import DictConfig, OmegaConf - -from grass_mil._paths import setup_project_root - -setup_project_root() - -from grass_mil.interpretability.core.data import load_interpretability_dataset # noqa: E402 -from grass_mil.interpretability.pipeline import run_interpretability_pipeline # noqa: E402 -from grass_mil.interpretability.reporting.render import render_interpretability_report # noqa: E402 - -log = logging.getLogger(__name__) - - -def _cfg_to_dict(cfg: Any) -> Dict[str, Any]: - return OmegaConf.to_container(cfg, resolve=True) if not isinstance(cfg, dict) else cfg - - -def run_report(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: - dataset = load_interpretability_dataset( - instance_table_path=Path(str(cfg.data.instance_table)), - bag_table_path=Path(str(cfg.data.bag_table)) if cfg.data.get("bag_table") else None, - spatial_table_path=Path(str(cfg.data.spatial_table)) - if cfg.data.get("spatial_table") - else None, - cell_table_path=Path(str(cfg.data.cell_table)) if cfg.data.get("cell_table") else None, - cell_edge_table_path=Path(str(cfg.data.cell_edge_table)) - if cfg.data.get("cell_edge_table") - else None, - id_column=str(cfg.data.get("id_column", "instance_id")), - bag_id_column=str(cfg.data.get("bag_id_column", "bag_id")), - cell_type_column=cfg.data.get("cell_type_column"), - condition_column=cfg.data.get("condition_column"), - ) - output_dir = Path(str(cfg.output_dir)) - artifacts_dir = output_dir / "artifacts" - pipeline_cfg = { - "embedding_prefixes": list(cfg.embedding_prefixes), - "composition_prefix": str(cfg.get("composition_prefix", "comp_")), - "variance_estimator": str(cfg.get("variance_estimator", "unbiased")), - "reduction": _cfg_to_dict(cfg.reduction), - "clustering": _cfg_to_dict(cfg.clustering), - "plugins": _cfg_to_dict(cfg.plugins), - } - bundle = run_interpretability_pipeline(dataset, pipeline_cfg, artifacts_dir=artifacts_dir) - report_out = render_interpretability_report( - bundle, - output_dir=output_dir, - html_enabled=bool(cfg.report.html.enabled), - pdf_enabled=bool(cfg.report.pdf.enabled), - snapshot_dpi_scale=float(cfg.report.pdf.snapshot_dpi_scale), - ) - summary = { - "output_dir": str(output_dir), - "artifacts_dir": str(artifacts_dir), - "report": report_out, - "pipeline": bundle.metadata, - } - (output_dir / "report_summary.json").write_text(json.dumps(summary, indent=2)) - object_dict: Dict[str, Any] = {"cfg": cfg, "bundle": bundle} - return summary, object_dict - - -@hydra.main( - version_base="1.3", - config_path="../configs", - config_name="interpretability/report.yaml", -) -def main(cfg: DictConfig) -> None: - run_report(cfg) - - -if __name__ == "__main__": - main() diff --git a/src/grass_mil/interpretability/reporting/__init__.py b/src/grass_mil/interpretability/reporting/__init__.py deleted file mode 100644 index 14e9546..0000000 --- a/src/grass_mil/interpretability/reporting/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from grass_mil.interpretability.reporting.render import render_interpretability_report - -__all__ = ["render_interpretability_report"] diff --git a/src/grass_mil/interpretability/reporting/html.py b/src/grass_mil/interpretability/reporting/html.py deleted file mode 100644 index 4613a91..0000000 --- a/src/grass_mil/interpretability/reporting/html.py +++ /dev/null @@ -1,74 +0,0 @@ -from __future__ import annotations - -import html -from pathlib import Path -from typing import Any, Dict, List - -import pandas as pd - -from grass_mil.contracts import ReportSection - - -def _render_table(title: str, table: pd.DataFrame) -> str: - return f"

{html.escape(title)}

" + table.to_html( - classes="table table-striped", border=0, index=True - ) - - -def render_html_report( - sections: List[ReportSection], - figures: Dict[str, Any], - output_path: Path, -) -> Path: - output_path = Path(output_path) - output_path.parent.mkdir(parents=True, exist_ok=True) - - figure_blocks: List[str] = [] - for name, figure in figures.items(): - figure_blocks.append( - f"

{html.escape(name)}

{figure.to_html(full_html=False, include_plotlyjs='cdn')}
" - ) - - section_blocks: List[str] = [] - for section in sections: - meta_html = "" - if section.metadata: - meta_df = pd.DataFrame([section.metadata]) - meta_html = _render_table("Metadata", meta_df) - tables_html = "".join(_render_table(k, v) for k, v in section.tables.items()) - section_blocks.append( - "
" - f"

{html.escape(section.title)}

" - f"

{html.escape(section.description)}

" - f"{meta_html}{tables_html}" - "
" - ) - - html_doc = f""" - - - - - Interpretability Report - - - -

Interpretability Report

- {''.join(section_blocks)} -
-

Interactive Plots

- {''.join(figure_blocks)} - -""" - output_path.write_text(html_doc, encoding="utf-8") - return output_path diff --git a/src/grass_mil/interpretability/reporting/pdf.py b/src/grass_mil/interpretability/reporting/pdf.py deleted file mode 100644 index f0a5f47..0000000 --- a/src/grass_mil/interpretability/reporting/pdf.py +++ /dev/null @@ -1,33 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - - -def render_pdf_via_playwright( - *, - html_path: Path, - pdf_path: Path, - timeout_ms: int = 60_000, -) -> Path: - html_path = Path(html_path).resolve() - pdf_path = Path(pdf_path).resolve() - pdf_path.parent.mkdir(parents=True, exist_ok=True) - try: - from playwright.sync_api import sync_playwright - except Exception as exc: # pragma: no cover - optional dependency - raise RuntimeError( - "Playwright is unavailable. Install playwright and browser binaries to enable PDF export." - ) from exc - - with sync_playwright() as p: # pragma: no cover - depends on browser runtime - browser = p.chromium.launch() - page = browser.new_page() - page.goto(f"file://{html_path}", wait_until="networkidle", timeout=timeout_ms) - page.pdf( - path=str(pdf_path), - format="A4", - print_background=True, - margin={"top": "15mm", "bottom": "15mm", "left": "12mm", "right": "12mm"}, - ) - browser.close() - return pdf_path diff --git a/src/grass_mil/interpretability/reporting/plotly_builders.py b/src/grass_mil/interpretability/reporting/plotly_builders.py deleted file mode 100644 index 4079e3a..0000000 --- a/src/grass_mil/interpretability/reporting/plotly_builders.py +++ /dev/null @@ -1,492 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, Dict, List - -import numpy as np -import pandas as pd - -from grass_mil.contracts import ReportBundle -from grass_mil.interpretability.core.niches import ( - BACKGROUND_NICHE_ID, - niche_display_name, - niche_display_names, -) -from grass_mil.interpretability.tier2.tissue_graph import ( - TissueGraphView, - build_tissue_graph_figure, -) - -if TYPE_CHECKING: - import plotly.graph_objects as go - - -def _require_plotly(): - try: - import plotly.graph_objects as go - except ModuleNotFoundError as exc: - raise ModuleNotFoundError( - "plotly is required for interpretability plotting. " - "Install it with: pip install plotly" - ) from exc - return go - - -def _safe_slug(value: object) -> str: - text = str(value).strip().lower() - if not text: - return "unknown" - out = [] - for char in text: - out.append(char if char.isalnum() else "_") - slug = "".join(out).strip("_") - return slug or "unknown" - - -def build_reduction_scatter(bundle: ReportBundle) -> "go.Figure": - go = _require_plotly() - if bundle.reduction is None: - raise ValueError("Reduction result is missing.") - emb = np.asarray(bundle.reduction.embedding) - if emb.shape[1] < 2: - raise ValueError("Reduction embedding must have at least 2 dimensions for scatter.") - labels = ( - bundle.clustering.labels - if bundle.clustering is not None - else np.zeros((emb.shape[0],), dtype=int) - ) - fig = go.Figure( - data=[ - go.Scattergl( - x=emb[:, 0], - y=emb[:, 1], - mode="markers", - marker={"size": 6, "color": labels, "colorscale": "Viridis"}, - name="instances", - ) - ] - ) - fig.update_layout( - title="Reduced Embedding Scatter", - xaxis_title="Dim 1", - yaxis_title="Dim 2", - template="plotly_white", - ) - return fig - - -def build_niche_enrichment_heatmap(bundle: ReportBundle) -> "go.Figure": - go = _require_plotly() - if bundle.niche_summary is None: - raise ValueError("Niche summary is missing.") - enr = bundle.niche_summary.enrichment - fig = go.Figure( - data=[ - go.Heatmap( - z=enr.to_numpy(), - x=[str(c) for c in enr.columns], - # Rows are already in dendrogram order from the niche summary. - y=niche_display_names(enr.index), - colorscale="RdBu", - zmid=0.0, - ) - ] - ) - fig.update_layout( - title="Niche Enrichment Heatmap", - xaxis_title="Features", - yaxis_title="Niche", - template="plotly_white", - ) - return fig - - -def build_plugin_figures(bundle: ReportBundle) -> Dict[str, "go.Figure"]: - go = _require_plotly() - figures: Dict[str, go.Figure] = {} - for name, result in bundle.plugin_results.items(): - if "curves" in result.payload and "thresholds" in result.payload: - thresholds = np.asarray(result.payload["thresholds"]) - curves = result.payload["curves"] - fig = go.Figure() - for niche, ct_map in curves.items(): - for cell_type, values in ct_map.items(): - fig.add_trace( - go.Scatter( - x=thresholds, - y=np.asarray(values), - mode="lines", - name=f"{niche}:{cell_type}", - ) - ) - fig.update_layout( - title=f"{name} Curves", - xaxis_title="Threshold", - yaxis_title="Value", - template="plotly_white", - ) - figures[name] = fig - elif "enrichment" in result.payload and isinstance( - result.payload["enrichment"], pd.DataFrame - ): - enr = result.payload["enrichment"] - fig = go.Figure( - data=[ - go.Heatmap( - z=enr.to_numpy(), - x=[str(c) for c in enr.columns], - y=[str(i) for i in enr.index], - colorscale="RdBu", - zmid=0.0, - ) - ] - ) - fig.update_layout(title=f"{name} Enrichment", template="plotly_white") - figures[name] = fig - elif "enrichment_by_pair" in result.payload and isinstance( - result.payload["enrichment_by_pair"], dict - ): - for pair, enr in result.payload["enrichment_by_pair"].items(): - if not isinstance(enr, pd.DataFrame): - continue - fig = go.Figure( - data=[ - go.Heatmap( - z=enr.to_numpy(), - x=[str(c) for c in enr.columns], - y=[str(i) for i in enr.index], - colorscale="RdBu", - zmid=0.0, - ) - ] - ) - fig.update_layout( - title=f"{name} Enrichment ({pair})", - template="plotly_white", - ) - figures[f"{name}_{pair}"] = fig - elif "enrichment_by_niche" in result.payload and isinstance( - result.payload["enrichment_by_niche"], dict - ): - for niche_id, enr in result.payload["enrichment_by_niche"].items(): - if not isinstance(enr, pd.DataFrame): - continue - fig = go.Figure( - data=[ - go.Heatmap( - z=enr.to_numpy(), - x=[str(c) for c in enr.columns], - y=[str(i) for i in enr.index], - colorscale="RdBu", - zmid=0.0, - ) - ] - ) - fig.update_layout( - title=f"{name} — {niche_display_name(niche_id)}", - template="plotly_white", - ) - figures[f"{name}_niche_{_safe_slug(niche_id)}"] = fig - elif "tissue_graph_view" in result.payload and isinstance( - result.payload["tissue_graph_view"], TissueGraphView - ): - view = result.payload["tissue_graph_view"] - style = result.payload.get("tissue_graph_style", {}) - fig = build_tissue_graph_figure( - view, - show_edges=bool(style.get("show_edges", True)), - node_size=float(style.get("node_size", 5.0)), - edge_width=float(style.get("edge_width", 0.5)), - edge_opacity=float(style.get("edge_opacity", 0.25)), - colorscale=str(style.get("colorscale", "Viridis")), - reverse_y=bool(style.get("reverse_y", True)), - title_prefix=str(style.get("title_prefix", "Tissue Graph")), - ) - meta = result.payload.get("tissue_graph_meta", {}) - sample_column = _safe_slug(meta.get("sample_column", "sample")) - sample_value = _safe_slug(meta.get("sample_value", "unknown")) - figures[f"tissue_graph_{sample_column}_{sample_value}"] = fig - return figures - - -def _resolve_color_values(series: "pd.Series") -> tuple[np.ndarray, Dict[str, object]]: - """Map an instance-table column onto numeric marker colors. - - Numeric columns are used as-is. Non-numeric columns (``condition`` is - routinely a string label) are factorized into integer codes, with the - original labels retained as colorbar ticks, so categorical attributes stay - plottable instead of raising. - - Returns: - (values, colorbar_extra) where ``colorbar_extra`` is merged into the - marker colorbar spec. - """ - numeric = pd.to_numeric(series, errors="coerce") - if bool(numeric.notna().any()): - return numeric.to_numpy(dtype=float), {} - - codes, uniques = pd.factorize(series, sort=True) - values = codes.astype(float) - values[codes < 0] = np.nan - colorbar_extra: Dict[str, object] = { - "tickmode": "array", - "tickvals": list(range(len(uniques))), - "ticktext": [str(value) for value in uniques], - } - return values, colorbar_extra - - -def build_multi_attribute_scatters( - bundle: ReportBundle, - extra_columns: Dict[str, str] | None = None, -) -> List[tuple[str, "go.Figure"]]: - """Build scatter plots from reduction embeddings colored by multiple attributes. - - Creates scatters for both PCA (from niche_feature_reduction) and the main - reduction (e.g. UMAP), each colored by niche labels and any extra columns - found in the instance_table. - - Args: - bundle: ReportBundle with reduction and optionally niche_feature_reduction. - extra_columns: mapping of {column_name: colorscale}. Defaults to common - attributes: condition, score, attention. - - Returns: - List of (name, figure) tuples. - """ - go = _require_plotly() - figures: List[tuple[str, go.Figure]] = [] - table = bundle.dataset.instance_table - - if extra_columns is None: - extra_columns = {} - for col, cscale in [ - ("condition", "Viridis"), - ("score", "Magma"), - ("attention", "Viridis"), - ]: - if col in table.columns: - extra_columns[col] = cscale - - labels = ( - bundle.clustering.labels - if bundle.clustering is not None - else np.zeros((len(table),), dtype=int) - ) - - # True for real niches; Background (-1) is excluded from the filtered views. - background_mask = labels != BACKGROUND_NICHE_ID if BACKGROUND_NICHE_ID in labels else None - - reductions: List[tuple[str, np.ndarray, str, str]] = [] - if bundle.niche_feature_reduction is not None: - emb = np.asarray(bundle.niche_feature_reduction.embedding) - if emb.shape[1] >= 2: - fitted = bundle.niche_feature_reduction.fitted_object - if fitted is not None and hasattr(fitted, "explained_variance_ratio_"): - ev = fitted.explained_variance_ratio_ - x_label = f"PC1 ({ev[0] * 100:.1f}% var)" - y_label = f"PC2 ({ev[1] * 100:.1f}% var)" - else: - x_label, y_label = "PC1", "PC2" - reductions.append(("pca", emb, x_label, y_label)) - if bundle.reduction is not None: - emb = np.asarray(bundle.reduction.embedding) - if emb.shape[1] >= 2: - method = bundle.reduction.method.upper() - reductions.append( - ( - bundle.reduction.method, - emb, - f"{method} 1", - f"{method} 2", - ) - ) - - for red_name, emb, x_label, y_label in reductions: - # Niche scatter (all points) - fig = go.Figure( - data=[ - go.Scattergl( - x=emb[:, 0], - y=emb[:, 1], - mode="markers", - marker={ - "size": 3, - "color": labels, - "colorscale": "Viridis", - "showscale": True, - }, - ) - ] - ) - fig.update_layout( - title=f"{red_name.upper()} — Niche Labels", - xaxis_title=x_label, - yaxis_title=y_label, - template="plotly_white", - ) - figures.append((f"{red_name}_niche_all", fig)) - - # Niche scatter (excluding Background) - if background_mask is not None: - fig = go.Figure( - data=[ - go.Scattergl( - x=emb[background_mask, 0], - y=emb[background_mask, 1], - mode="markers", - marker={ - "size": 3, - "color": labels[background_mask], - "colorscale": "Viridis", - "showscale": True, - }, - ) - ] - ) - fig.update_layout( - title=f"{red_name.upper()} — Niche Labels (excluding Background)", - xaxis_title=x_label, - yaxis_title=y_label, - template="plotly_white", - ) - figures.append((f"{red_name}_niche_filtered", fig)) - - # Extra column scatters - for col, cscale in extra_columns.items(): - values, colorbar_extra = _resolve_color_values(table[col]) - colorbar = {"title": col, **colorbar_extra} - fig = go.Figure( - data=[ - go.Scattergl( - x=emb[:, 0], - y=emb[:, 1], - mode="markers", - marker={ - "size": 3, - "color": values, - "colorscale": cscale, - "showscale": True, - "colorbar": colorbar, - }, - ) - ] - ) - fig.update_layout( - title=f"{red_name.upper()} — {col}", - xaxis_title=x_label, - yaxis_title=y_label, - template="plotly_white", - ) - figures.append((f"{red_name}_{_safe_slug(col)}", fig)) - - if background_mask is not None: - fig = go.Figure( - data=[ - go.Scattergl( - x=emb[background_mask, 0], - y=emb[background_mask, 1], - mode="markers", - marker={ - "size": 3, - "color": values[background_mask], - "colorscale": cscale, - "showscale": True, - "colorbar": colorbar, - }, - ) - ] - ) - fig.update_layout( - title=f"{red_name.upper()} — {col} (excluding Background)", - xaxis_title=x_label, - yaxis_title=y_label, - template="plotly_white", - ) - figures.append((f"{red_name}_{_safe_slug(col)}_filtered", fig)) - - return figures - - -def build_composite_niche_heatmap(bundle: ReportBundle) -> "go.Figure": - """Composition enrichment beside niche abundance. - - Attention-only panels were removed: interpretation is expressed through - weighted logit margins (the ``margin_attribution`` plugin), not through - attention weights on their own, which say what the model looked at but not - which way it pushed the decision. - """ - go = _require_plotly() - from plotly.subplots import make_subplots - - summary = bundle.niche_summary - if summary is None: - raise ValueError("Niche summary is missing.") - - enr = summary.enrichment - # Niche ids already carry dendrogram order; Background is named. - niche_ids = niche_display_names(enr.index) - feature_names = [str(c) for c in enr.columns] - - fig = make_subplots( - rows=1, - cols=2, - shared_yaxes=True, - column_widths=[13, 1], - horizontal_spacing=0.02, - ) - fig.add_trace( - go.Heatmap( - z=enr.to_numpy(), - x=feature_names, - y=niche_ids, - colorscale="RdBu", - zmid=0.0, - colorbar={"title": "Enrichment", "x": 0.86, "len": 0.9}, - ), - row=1, - col=1, - ) - - counts = summary.niche_counts.reindex(enr.index).to_numpy().astype(float) - log_counts = np.log(np.clip(counts, 1, None)).reshape(-1, 1) - fig.add_trace( - go.Heatmap( - z=log_counts, - x=["Log Abund."], - y=niche_ids, - colorscale="Viridis", - colorbar={"title": "Log N", "x": 1.02, "len": 0.9}, - ), - row=1, - col=2, - ) - - fig.update_layout( - title="Niche Composition and Abundance", - template="plotly_white", - height=max(400, 50 * len(niche_ids) + 200), - ) - return fig - - -def bundle_figures(bundle: ReportBundle) -> List[tuple[str, "go.Figure"]]: - figures: List[tuple[str, go.Figure]] = [] - if bundle.reduction is not None: - figures.append(("reduction_scatter", build_reduction_scatter(bundle))) - if bundle.niche_summary is not None: - figures.append(("niche_enrichment", build_niche_enrichment_heatmap(bundle))) - - # Multi-attribute scatters (PCA + UMAP colored by various attributes) - if bundle.reduction is not None or bundle.niche_feature_reduction is not None: - figures.extend(build_multi_attribute_scatters(bundle)) - - # Composite niche heatmap (enrichment + predictions + attention + abundance) - if bundle.niche_summary is not None: - try: - figures.append(("composite_niche_summary", build_composite_niche_heatmap(bundle))) - except (ValueError, KeyError): - pass - - for name, fig in build_plugin_figures(bundle).items(): - figures.append((name, fig)) - return figures diff --git a/src/grass_mil/interpretability/reporting/render.py b/src/grass_mil/interpretability/reporting/render.py deleted file mode 100644 index 065a207..0000000 --- a/src/grass_mil/interpretability/reporting/render.py +++ /dev/null @@ -1,61 +0,0 @@ -from __future__ import annotations - -import logging -from pathlib import Path -from typing import Any, Dict - -from grass_mil.contracts import ReportBundle -from grass_mil.interpretability.reporting.html import render_html_report -from grass_mil.interpretability.reporting.pdf import render_pdf_via_playwright -from grass_mil.interpretability.reporting.plotly_builders import bundle_figures -from grass_mil.interpretability.reporting.sections import build_report_sections -from grass_mil.interpretability.reporting.snapshot import export_plotly_snapshots - -log = logging.getLogger(__name__) - - -def render_interpretability_report( - bundle: ReportBundle, - *, - output_dir: Path, - html_enabled: bool = True, - pdf_enabled: bool = True, - snapshot_dpi_scale: float = 2.0, -) -> Dict[str, Any]: - output_dir = Path(output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - - figure_pairs = bundle_figures(bundle) - section_list = build_report_sections(bundle) - figure_map = {name: fig for name, fig in figure_pairs} - - html_path = None - if html_enabled: - html_path = render_html_report(section_list, figure_map, output_dir / "report.html") - - snapshot_dir = output_dir / "artifacts" / "snapshots" - try: - snapshot_paths = export_plotly_snapshots( - figure_pairs, - snapshot_dir, - scale=snapshot_dpi_scale, - ) - except Exception as exc: # pragma: no cover - optional runtime deps - # Snapshot export is auxiliary; report rendering should remain usable without kaleido. - log.warning("Skipping plot snapshot export: %s", exc) - snapshot_paths = {} - - pdf_path = None - if pdf_enabled: - if html_path is None: - html_path = render_html_report(section_list, figure_map, output_dir / "report.html") - pdf_path = render_pdf_via_playwright( - html_path=html_path, - pdf_path=output_dir / "report.pdf", - ) - - return { - "html_path": str(html_path) if html_path else None, - "pdf_path": str(pdf_path) if pdf_path else None, - "snapshot_paths": {k: str(v) for k, v in snapshot_paths.items()}, - } diff --git a/src/grass_mil/interpretability/reporting/sections.py b/src/grass_mil/interpretability/reporting/sections.py deleted file mode 100644 index 9797347..0000000 --- a/src/grass_mil/interpretability/reporting/sections.py +++ /dev/null @@ -1,32 +0,0 @@ -from __future__ import annotations - -from typing import List - -from grass_mil.contracts import ReportBundle, ReportSection - - -def build_report_sections(bundle: ReportBundle) -> List[ReportSection]: - sections: List[ReportSection] = [] - - sections.append( - ReportSection( - title="Executive Summary", - description="Interpretability analysis output bundle summary.", - metadata=bundle.metadata, - ) - ) - if bundle.niche_summary is not None: - sections.append( - ReportSection( - title="Niche Summary", - description="Niche counts and enrichment tables.", - tables={ - "niche_counts": bundle.niche_summary.niche_counts.to_frame(name="count"), - "enrichment": bundle.niche_summary.enrichment, - }, - ) - ) - - for plugin_result in bundle.plugin_results.values(): - sections.extend(plugin_result.sections) - return sections diff --git a/src/grass_mil/interpretability/reporting/snapshot.py b/src/grass_mil/interpretability/reporting/snapshot.py deleted file mode 100644 index 1a62db3..0000000 --- a/src/grass_mil/interpretability/reporting/snapshot.py +++ /dev/null @@ -1,27 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from typing import Any, Iterable - - -def export_plotly_snapshots( - figures: Iterable[tuple[str, Any]], - output_dir: Path, - *, - width: int = 1600, - height: int = 900, - scale: float = 2.0, -) -> dict[str, Path]: - output_dir = Path(output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - out: dict[str, Path] = {} - for name, fig in figures: - path = output_dir / f"{name}.png" - try: - fig.write_image(path, width=width, height=height, scale=scale) - except Exception as exc: # pragma: no cover - requires kaleido runtime - raise RuntimeError( - "Failed to export plot snapshot. Ensure `kaleido` is installed and working." - ) from exc - out[name] = path - return out diff --git a/src/grass_mil/interpretability/tier2/__init__.py b/src/grass_mil/interpretability/tier2/__init__.py deleted file mode 100644 index dafdaff..0000000 --- a/src/grass_mil/interpretability/tier2/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -from grass_mil.interpretability.tier2.filtration import compute_filtration_curves -from grass_mil.interpretability.tier2.neighborhood import ( - run_diff_neighborhood_enrichment, - run_neighborhood_enrichment, -) -from grass_mil.interpretability.tier2.tissue_graph import ( - build_tissue_graph_figure, - prepare_tissue_graph_view, -) - -__all__ = [ - "run_neighborhood_enrichment", - "run_diff_neighborhood_enrichment", - "compute_filtration_curves", - "prepare_tissue_graph_view", - "build_tissue_graph_figure", -] diff --git a/src/grass_mil/interpretability/tier2/autocorrelation.py b/src/grass_mil/interpretability/tier2/autocorrelation.py deleted file mode 100644 index 87254c1..0000000 --- a/src/grass_mil/interpretability/tier2/autocorrelation.py +++ /dev/null @@ -1,213 +0,0 @@ -"""Global spatial autocorrelation (Moran's I). - -Moran's I over a binary connectivity matrix :math:`W`: - -.. math:: I = \\frac{N}{S_0} \\cdot \\frac{z^\\top W z}{z^\\top z}, - \\qquad S_0 = \\sum_{u,v} w_{uv} - -with :math:`z` the mean-centred feature. Significance comes from a permutation -null: the feature is shuffled over nodes while the graph is held fixed. - -Implemented natively over ``scipy.sparse`` rather than through squidpy. The -package stays installable without scanpy/numba, ``tier2.neighborhood`` already -sets that precedent, and the statistic is a few lines over a sparse matrix. - -**Scope difference worth knowing:** the legacy report computes Moran's I over -*cells* inside pooled ego-graphs. Here the spatial table is root-instance level, -so this measures autocorrelation over the *instance* graph. That is a defensible -and arguably cleaner object, but it is not numerically the same statistic. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Dict, Sequence - -import numpy as np -import pandas as pd - -__all__ = ["MoranResult", "morans_i", "run_morans_i"] - - -@dataclass(frozen=True) -class MoranResult: - """Per-group Moran's I with its permutation null. - - Attributes: - statistic: group x feature observed I. - expected: group x feature permutation mean. - variance: group x feature permutation variance. - pvalue: two-sided permutation p-value. - qvalue: Benjamini-Hochberg adjusted p-value across all tests. - n_nodes: node count per group. - """ - - statistic: pd.DataFrame - expected: pd.DataFrame - variance: pd.DataFrame - pvalue: pd.DataFrame - qvalue: pd.DataFrame - n_nodes: pd.Series - - -def _benjamini_hochberg(pvalues: np.ndarray) -> np.ndarray: - """BH-FDR adjustment, ignoring NaNs.""" - flat = np.asarray(pvalues, dtype=float).ravel() - out = np.full_like(flat, np.nan) - finite = np.isfinite(flat) - if not finite.any(): - return out.reshape(np.shape(pvalues)) - try: - from scipy.stats import false_discovery_control - - out[finite] = false_discovery_control(flat[finite], method="bh") - except Exception: - values = flat[finite] - order = np.argsort(values) - ranked = values[order] - n = ranked.size - adjusted = ranked * n / np.arange(1, n + 1) - adjusted = np.minimum.accumulate(adjusted[::-1])[::-1] - restored = np.empty_like(adjusted) - restored[order] = np.clip(adjusted, 0.0, 1.0) - out[finite] = restored - return out.reshape(np.shape(pvalues)) - - -def morans_i(values: np.ndarray, adjacency) -> float: - """Global Moran's I of ``values`` over a sparse connectivity matrix.""" - values = np.asarray(values, dtype=float).ravel() - n = values.size - if n < 2: - return float("nan") - s0 = float(adjacency.sum()) - if s0 <= 0: - return float("nan") - z = values - values.mean() - denom = float(z @ z) - if denom <= 0: - # A constant feature has no variance to correlate. - return float("nan") - return float((n / s0) * float(z @ (adjacency @ z)) / denom) - - -def _adjacency_from_edges( - edges: pd.DataFrame, - node_ids: Sequence[str], - *, - source_column: str, - target_column: str, -): - from scipy.sparse import coo_matrix - - index = {str(node): position for position, node in enumerate(node_ids)} - src = edges[source_column].astype(str).map(index) - dst = edges[target_column].astype(str).map(index) - keep = src.notna() & dst.notna() - src = src[keep].to_numpy(dtype=np.int64) - dst = dst[keep].to_numpy(dtype=np.int64) - n = len(node_ids) - if src.size == 0: - return coo_matrix((n, n), dtype=float).tocsr() - rows = np.concatenate([src, dst]) - cols = np.concatenate([dst, src]) - data = np.ones(rows.size, dtype=float) - adjacency = coo_matrix((data, (rows, cols)), shape=(n, n)).tocsr() - adjacency.data[:] = 1.0 # binary weights; collapse any duplicate edges - return adjacency - - -def run_morans_i( - node_table: pd.DataFrame, - spatial_table: pd.DataFrame, - *, - feature_columns: Sequence[str], - group_column: str, - id_column: str = "instance_id", - source_column: str = "source_id", - target_column: str = "target_id", - n_perms: int = 100, - random_state: int = 0, - min_nodes: int = 3, -) -> MoranResult: - """Moran's I per group, with a permutation null over node labels.""" - missing = [c for c in (*feature_columns, group_column, id_column) if c not in node_table] - if missing: - raise ValueError(f"node_table is missing required column(s): {missing}.") - - rng = np.random.default_rng(random_state) - groups = sorted(node_table[group_column].unique()) - statistic: Dict[object, Dict[str, float]] = {} - expected: Dict[object, Dict[str, float]] = {} - variance: Dict[object, Dict[str, float]] = {} - pvalue: Dict[object, Dict[str, float]] = {} - counts: Dict[object, int] = {} - - for group in groups: - subset = node_table[node_table[group_column] == group] - counts[group] = int(len(subset)) - statistic[group] = {} - expected[group] = {} - variance[group] = {} - pvalue[group] = {} - if len(subset) < min_nodes: - for feature in feature_columns: - statistic[group][feature] = float("nan") - expected[group][feature] = float("nan") - variance[group][feature] = float("nan") - pvalue[group][feature] = float("nan") - continue - - node_ids = subset[id_column].astype(str).tolist() - keep_ids = set(node_ids) - edges = spatial_table[ - spatial_table[source_column].astype(str).isin(keep_ids) - & spatial_table[target_column].astype(str).isin(keep_ids) - ] - adjacency = _adjacency_from_edges( - edges, node_ids, source_column=source_column, target_column=target_column - ) - - for feature in feature_columns: - values = subset[feature].to_numpy(dtype=float) - observed = morans_i(values, adjacency) - statistic[group][feature] = observed - - if n_perms <= 0 or not np.isfinite(observed): - expected[group][feature] = float("nan") - variance[group][feature] = float("nan") - pvalue[group][feature] = float("nan") - continue - - null = np.array( - [morans_i(rng.permutation(values), adjacency) for _ in range(int(n_perms))], - dtype=float, - ) - null = null[np.isfinite(null)] - if null.size == 0: - expected[group][feature] = float("nan") - variance[group][feature] = float("nan") - pvalue[group][feature] = float("nan") - continue - expected[group][feature] = float(null.mean()) - variance[group][feature] = float(null.var(ddof=1)) if null.size > 1 else 0.0 - # Two-sided, with the observed value included so p is never zero. - extreme = int(np.sum(np.abs(null - null.mean()) >= abs(observed - null.mean()))) - pvalue[group][feature] = float((extreme + 1) / (null.size + 1)) - - columns = list(feature_columns) - stat_frame = pd.DataFrame(statistic).T.reindex(columns=columns) - exp_frame = pd.DataFrame(expected).T.reindex(columns=columns) - var_frame = pd.DataFrame(variance).T.reindex(columns=columns) - p_frame = pd.DataFrame(pvalue).T.reindex(columns=columns) - q_frame = pd.DataFrame( - _benjamini_hochberg(p_frame.to_numpy()), index=p_frame.index, columns=p_frame.columns - ) - return MoranResult( - statistic=stat_frame, - expected=exp_frame, - variance=var_frame, - pvalue=p_frame, - qvalue=q_frame, - n_nodes=pd.Series(counts).reindex(stat_frame.index), - ) diff --git a/src/grass_mil/interpretability/tier2/cell_level.py b/src/grass_mil/interpretability/tier2/cell_level.py deleted file mode 100644 index 28e390b..0000000 --- a/src/grass_mil/interpretability/tier2/cell_level.py +++ /dev/null @@ -1,311 +0,0 @@ -"""Cell-level tier-2 analyses, run within each niche. - -The instance-level analyses treat one ego-graph as one node, so they measure -relationships between neighbourhood summaries. These measure the cells: every -node of every sampled ego-graph is a row, carrying its own type, and each cell -inherits the niche of the subgraph it came from. Within a niche you then get -cell-type x cell-type statistics. - -That is the unit the legacy NSCLC reports use: `run_per_cluster_nhood_enrichment` -pools the member ego-graphs and runs cell-type enrichment per cluster, -`run_moran_per_cluster` computes Moran's I per cell type over the same pooled -graph, and `run_filtration_curves` counts both endpoints of every edge by type. - -The numerics are the existing tier-2 primitives -- these functions only assemble -the right inputs, because "cell-level" is a statement about the table, not about -the estimator. -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Sequence - -import numpy as np -import pandas as pd - -from grass_mil.interpretability.tier2.autocorrelation import MoranResult, run_morans_i -from grass_mil.interpretability.tier2.neighborhood import ( - run_neighborhood_enrichment, -) - -BACKGROUND_NICHE_ID = -1 - - -@dataclass(frozen=True) -class PerNicheEnrichmentResult: - """Cell-type x cell-type enrichment, one matrix per niche.""" - - enrichment: Dict[Any, pd.DataFrame] - pvalues: Dict[Any, pd.DataFrame] - n_cells: Dict[Any, int] - n_edges: Dict[Any, int] - - -def attach_niche_labels_to_cells( - cell_table: pd.DataFrame, - niche_labels: Sequence[Any], - instance_ids: Sequence[Any], - *, - instance_id_column: str = "instance_id", - niche_column: str = "niche_label", -) -> pd.DataFrame: - """Give every cell the niche of the ego-graph it was sampled into. - - Niches are assigned to instances, and a cell belongs to whichever ego-graph - it appeared in, so the label is joined through `instance_id`. This mirrors - the legacy `cluster_labels[batch.batch]` step. - """ - if instance_id_column not in cell_table.columns: - raise ValueError(f"Cell table is missing {instance_id_column!r}.") - if len(niche_labels) != len(instance_ids): - raise ValueError( - "niche_labels and instance_ids must align: " - f"{len(niche_labels)} vs {len(instance_ids)}." - ) - mapping = pd.DataFrame( - {instance_id_column: list(instance_ids), niche_column: list(niche_labels)} - ).drop_duplicates(subset=[instance_id_column]) - merged = cell_table.merge(mapping, on=instance_id_column, how="left") - return merged - - -def add_cell_type_indicators( - cell_table: pd.DataFrame, - *, - cell_type_column: str = "cell_type", - prefix: str = "ct_", - categories: Optional[Sequence[str]] = None, -) -> tuple[pd.DataFrame, List[str]]: - """One-hot the cell type, so each type can be treated as a signal. - - Moran's I asks whether a value is spatially clustered; for a categorical - type that question is asked per type, over its indicator. - """ - if cell_type_column not in cell_table.columns: - raise ValueError(f"Cell table is missing {cell_type_column!r}.") - values = cell_table[cell_type_column].astype(str) - cats = list(categories) if categories is not None else sorted(values.unique()) - out = cell_table.copy() - columns: List[str] = [] - for cat in cats: - col = f"{prefix}{cat}" - out[col] = (values == cat).astype(float) - columns.append(col) - return out, columns - - -def _subset_edges(edges: pd.DataFrame, ids: set, *, source: str, target: str) -> pd.DataFrame: - return edges[edges[source].isin(ids) & edges[target].isin(ids)] - - -def per_niche_cell_type_enrichment( - cell_table: pd.DataFrame, - cell_edge_table: pd.DataFrame, - *, - cell_type_column: str = "cell_type", - niche_column: str = "niche_label", - cell_id_column: str = "cell_id", - source_column: str = "source_id", - target_column: str = "target_id", - n_perms: int = 0, - random_state: int = 42, - undirected: bool = False, - enrichment_mode: str = "zscore", - min_cells: int = 10, - skip_background: bool = False, -) -> PerNicheEnrichmentResult: - """Cell-type neighbourhood enrichment computed separately within each niche. - - Only edges whose both endpoints lie in the niche are used, so each matrix - describes the internal wiring of that niche rather than its border. - """ - for col in (cell_type_column, niche_column, cell_id_column): - if col not in cell_table.columns: - raise ValueError(f"Cell table is missing {col!r}.") - - # A shared category list keeps the matrices aligned across niches, which is - # what makes them comparable (and differenceable) afterwards. - categories = sorted(cell_table[cell_type_column].astype(str).unique()) - - enrichment: Dict[Any, pd.DataFrame] = {} - pvalues: Dict[Any, pd.DataFrame] = {} - n_cells: Dict[Any, int] = {} - n_edges: Dict[Any, int] = {} - - for niche, group in cell_table.groupby(niche_column, dropna=True): - if skip_background and int(niche) == BACKGROUND_NICHE_ID: - continue - if len(group) < int(min_cells): - continue - ids = set(group[cell_id_column]) - sub_edges = _subset_edges(cell_edge_table, ids, source=source_column, target=target_column) - if sub_edges.empty: - continue - result = run_neighborhood_enrichment( - group, - sub_edges, - label_column=cell_type_column, - id_column=cell_id_column, - n_perms=int(n_perms), - random_state=int(random_state), - undirected=bool(undirected), - categories=categories, - warn_analytical=False, - enrichment_mode=str(enrichment_mode), - ) - enrichment[niche] = result.enrichment - pvalues[niche] = result.pvalues - n_cells[niche] = int(len(group)) - n_edges[niche] = int(len(sub_edges)) - - return PerNicheEnrichmentResult( - enrichment=enrichment, pvalues=pvalues, n_cells=n_cells, n_edges=n_edges - ) - - -def differential_cell_type_enrichment( - result: PerNicheEnrichmentResult, - *, - reference: Any = BACKGROUND_NICHE_ID, -) -> Dict[Any, pd.DataFrame]: - """Each niche's enrichment minus a reference niche's. - - The legacy reports difference every cluster against the noise cluster, which - separates structure specific to a niche from structure present everywhere. - """ - if reference not in result.enrichment: - return {} - ref = result.enrichment[reference] - return { - niche: (matrix - ref) for niche, matrix in result.enrichment.items() if niche != reference - } - - -def per_niche_cell_type_moran( - cell_table: pd.DataFrame, - cell_edge_table: pd.DataFrame, - *, - cell_type_column: str = "cell_type", - niche_column: str = "niche_label", - cell_id_column: str = "cell_id", - n_perms: int = 100, - random_state: int = 0, - min_nodes: int = 3, -) -> MoranResult: - """Moran's I of each cell-type indicator, per niche. - - Yields a niche x cell-type table: whether cells of a given type are - spatially clustered within that niche's graph. - """ - table, indicator_columns = add_cell_type_indicators( - cell_table, cell_type_column=cell_type_column - ) - return run_morans_i( - table, - cell_edge_table, - feature_columns=indicator_columns, - group_column=niche_column, - id_column=cell_id_column, - n_perms=int(n_perms), - random_state=int(random_state), - min_nodes=int(min_nodes), - ) - - -def niche_label_moran( - instance_table: pd.DataFrame, - spatial_table: pd.DataFrame, - *, - niche_column: str = "niche_label", - id_column: str = "instance_id", - n_perms: int = 100, - random_state: int = 0, - group_column: Optional[str] = None, -) -> MoranResult: - """Moran's I of the niche assignment itself, over the instance graph. - - The global counterpart to the per-niche cell-type analyses: it asks whether - niches form contiguous territories rather than being interleaved. The - legacy equivalent is `run_cluster_self_moran`. - """ - table = instance_table.copy() - if group_column is None: - group_column = "_all" - table[group_column] = "all" - table, indicator_columns = add_cell_type_indicators( - table, cell_type_column=niche_column, prefix="niche_" - ) - return run_morans_i( - table, - spatial_table, - feature_columns=indicator_columns, - group_column=group_column, - id_column=id_column, - n_perms=int(n_perms), - random_state=int(random_state), - ) - - -def cell_filtration_curves( - cell_table: pd.DataFrame, - cell_edge_table: pd.DataFrame, - *, - thresholds: np.ndarray, - cell_type_column: str = "cell_type", - niche_column: str = "niche_label", - cell_id_column: str = "cell_id", - distance_column: str = "distance", - scale_within_niche: bool = True, -): - """Filtration curves over cell-cell edges, counting both endpoints. - - The instance-level version walks edges between ego-graphs and has one label - per node, so its curves describe neighbourhoods. This walks the edges - *inside* the ego-graphs, where each endpoint is a cell with its own type -- - the legacy `run_filtration_curves` behaviour. - """ - from grass_mil.interpretability.tier2.filtration import compute_filtration_curves - - return compute_filtration_curves( - cell_table, - cell_edge_table, - thresholds=np.asarray(thresholds, dtype=float), - niche_column=niche_column, - cell_type_column=cell_type_column, - id_column=cell_id_column, - distance_column=distance_column, - scale_within_cluster=bool(scale_within_niche), - ) - - -def per_niche_cell_type_ripley( - cell_table: pd.DataFrame, - *, - cell_type_column: str = "cell_type", - niche_column: str = "niche_label", - x_column: str = "x", - y_column: str = "y", - n_radii: int = 50, - max_fraction: float = 0.25, - min_count: int = 5, - radius_source: str = "median", -): - """Ripley cross-L between cell types, with each niche a region. - - Uses the cells' own coordinates rather than ego-graph centroids, so the - curves describe how cell types are arranged relative to one another. - """ - from grass_mil.interpretability.tier2.ripley import aggregate_ripley - - return aggregate_ripley( - cell_table, - label_column=cell_type_column, - group_column=niche_column, - x_column=x_column, - y_column=y_column, - n_radii=int(n_radii), - max_fraction=float(max_fraction), - min_count=int(min_count), - radius_source=str(radius_source), - ) diff --git a/src/grass_mil/interpretability/tier2/filtration.py b/src/grass_mil/interpretability/tier2/filtration.py deleted file mode 100644 index 65f8132..0000000 --- a/src/grass_mil/interpretability/tier2/filtration.py +++ /dev/null @@ -1,110 +0,0 @@ -from __future__ import annotations - -import numpy as np -import pandas as pd - -from grass_mil.contracts import FiltrationCurvesResult - - -def compute_filtration_curves( - node_table: pd.DataFrame, - spatial_table: pd.DataFrame, - *, - thresholds: np.ndarray, - niche_column: str = "niche_label", - cell_type_column: str = "cell_type", - id_column: str = "instance_id", - distance_column: str = "distance", - scale_within_cluster: bool = True, -) -> FiltrationCurvesResult: - if niche_column not in node_table.columns: - raise ValueError(f"Missing niche column {niche_column!r}.") - if cell_type_column not in node_table.columns: - raise ValueError(f"Missing cell_type column {cell_type_column!r}.") - if id_column not in node_table.columns: - raise ValueError(f"Missing id column {id_column!r}.") - if distance_column not in spatial_table.columns: - raise ValueError(f"Missing distance column {distance_column!r} in spatial table.") - if {"source_id", "target_id"} - set(spatial_table.columns): - raise ValueError("spatial_table must include source_id/target_id columns.") - - nodes = node_table[[id_column, niche_column, cell_type_column]].copy() - nodes[id_column] = nodes[id_column].astype(str) - if nodes[id_column].duplicated().any(): - raise ValueError( - f"Node table contains duplicate IDs in {id_column!r}; " - "filtration requires unique node identifiers." - ) - nodes[niche_column] = nodes[niche_column].astype(str) - nodes[cell_type_column] = nodes[cell_type_column].astype(str) - node_to_cell_type = nodes.set_index(id_column)[cell_type_column].to_dict() - node_to_cluster = nodes.set_index(id_column)[niche_column].to_dict() - - edges = spatial_table.copy() - edges["source_id"] = edges["source_id"].astype(str) - edges["target_id"] = edges["target_id"].astype(str) - edges = edges[edges["source_id"].isin(node_to_cluster.keys())].copy() - edges[distance_column] = edges[distance_column].astype(float) - edges["source_cluster"] = edges["source_id"].map(node_to_cluster) - - niches = sorted(nodes[niche_column].astype(str).unique()) - cell_types = sorted(nodes[cell_type_column].astype(str).unique()) - cell_type_to_idx = {name: idx for idx, name in enumerate(cell_types)} - thresholds = np.asarray(thresholds, dtype=float) - if thresholds.ndim != 1: - raise ValueError("thresholds must be a 1D array.") - if thresholds.size == 0: - raise ValueError("thresholds must include at least one value.") - sort_order = np.argsort(thresholds) - thresholds_sorted = thresholds[sort_order] - curves: dict[str, dict[str, np.ndarray]] = { - c: {ct: np.zeros_like(thresholds, dtype=float) for ct in cell_types} for c in niches - } - - for niche in niches: - c_edges = edges[edges["source_cluster"] == str(niche)] - per_ct_diff = np.zeros((len(cell_types), thresholds_sorted.shape[0] + 1), dtype=float) - for _, subgraph_edges in c_edges.groupby("source_id", sort=False): - source_ids = subgraph_edges["source_id"].to_numpy(dtype=str) - target_ids = subgraph_edges["target_id"].to_numpy(dtype=str) - distances = subgraph_edges[distance_column].to_numpy(dtype=float) - if distances.size == 0: - continue - edge_order = np.argsort(distances) - source_ids = source_ids[edge_order] - target_ids = target_ids[edge_order] - distances = distances[edge_order] - - first_threshold_by_node: dict[str, int] = {} - for src_id, tgt_id, distance in zip(source_ids, target_ids, distances): - start_idx = int(np.searchsorted(thresholds_sorted, distance, side="left")) - if start_idx >= int(thresholds_sorted.shape[0]): - continue - prev_src = first_threshold_by_node.get(src_id) - if prev_src is None or start_idx < prev_src: - first_threshold_by_node[src_id] = start_idx - prev_tgt = first_threshold_by_node.get(tgt_id) - if prev_tgt is None or start_idx < prev_tgt: - first_threshold_by_node[tgt_id] = start_idx - - for node_id, start_idx in first_threshold_by_node.items(): - cell_type = node_to_cell_type.get(node_id) - if cell_type is None: - continue - ct_idx = cell_type_to_idx.get(cell_type) - if ct_idx is None: - continue - per_ct_diff[ct_idx, start_idx] += 1.0 - - per_ct_sorted = np.cumsum(per_ct_diff[:, :-1], axis=1) - per_ct_original = np.zeros_like(per_ct_sorted) - per_ct_original[:, sort_order] = per_ct_sorted - for ct_idx, ct in enumerate(cell_types): - curves[niche][ct] = per_ct_original[ct_idx].astype(float, copy=False) - if scale_within_cluster: - for ct in cell_types: - vmax = float(curves[niche][ct].max()) - if vmax > 0: - curves[niche][ct] = curves[niche][ct] / vmax - - return FiltrationCurvesResult(thresholds=thresholds, curves=curves) diff --git a/src/grass_mil/interpretability/tier2/neighborhood.py b/src/grass_mil/interpretability/tier2/neighborhood.py deleted file mode 100644 index c8cb6f2..0000000 --- a/src/grass_mil/interpretability/tier2/neighborhood.py +++ /dev/null @@ -1,421 +0,0 @@ -from __future__ import annotations - -import warnings -from itertools import combinations -from typing import Optional - -import numpy as np -import pandas as pd - -from grass_mil.contracts import NeighborhoodEnrichmentResult - - -def _coerce_edges(spatial_table: pd.DataFrame) -> pd.DataFrame: - required = {"source_id", "target_id"} - if not required.issubset(spatial_table.columns): - raise ValueError("spatial_table must include ['source_id', 'target_id'].") - return spatial_table.copy() - - -def _prepare_edges_with_labels( - node_table: pd.DataFrame, - spatial_table: pd.DataFrame, - *, - label_column: str, - id_column: str, - undirected: bool, -) -> tuple[pd.DataFrame, pd.Series]: - if id_column not in node_table.columns: - raise ValueError(f"Missing id column {id_column!r}.") - nodes = node_table[[id_column, label_column]].copy() - nodes[id_column] = nodes[id_column].astype(str) - if nodes[id_column].duplicated().any(): - raise ValueError( - f"Node table contains duplicate IDs in {id_column!r}; " - "neighborhood enrichment requires unique node identifiers." - ) - label_map = nodes.set_index(id_column)[label_column].astype(str) - - edges = _coerce_edges(spatial_table) - edges["source_id"] = edges["source_id"].astype(str) - edges["target_id"] = edges["target_id"].astype(str) - - if undirected: - src = edges["source_id"].to_numpy(dtype=str) - dst = edges["target_id"].to_numpy(dtype=str) - src_first = np.where(src <= dst, src, dst) - dst_second = np.where(src <= dst, dst, src) - canonical = pd.DataFrame( - { - "source_id": src_first, - "target_id": dst_second, - } - ).drop_duplicates(subset=["source_id", "target_id"], keep="first") - rev = canonical.rename(columns={"source_id": "target_id", "target_id": "source_id"}) - edges = pd.concat([canonical, rev], axis=0, ignore_index=True) - - edges = edges[ - edges["source_id"].isin(label_map.index) & edges["target_id"].isin(label_map.index) - ] - edges = edges.copy() - edges["source_label"] = edges["source_id"].map(label_map) - edges["target_label"] = edges["target_id"].map(label_map) - return edges, label_map - - -def _observed_matrix(edges: pd.DataFrame, categories: list[str]) -> pd.DataFrame: - observed = edges.groupby(["source_label", "target_label"]).size().unstack(fill_value=0.0) - return observed.reindex(index=categories, columns=categories, fill_value=0.0) - - -def _analytical_expected_and_std( - n_edges: int, - node_labels: pd.Series, - categories: list[str], -) -> tuple[pd.DataFrame, pd.DataFrame]: - n_nodes = int(node_labels.shape[0]) - if n_nodes == 0 or n_edges == 0: - zeros = pd.DataFrame(0.0, index=categories, columns=categories) - return zeros, zeros - niche_counts = node_labels.value_counts().reindex(categories, fill_value=0.0).astype(float) - probs = niche_counts.values / float(n_nodes) - pair_probs = np.outer(probs, probs) - expected = float(n_edges) * pair_probs - var = float(n_edges) * pair_probs * (1.0 - pair_probs) - std = np.sqrt(np.clip(var, a_min=0.0, a_max=None)) - return ( - pd.DataFrame(expected, index=categories, columns=categories), - pd.DataFrame(std, index=categories, columns=categories), - ) - - -def _zscore_enrichment( - observed: pd.DataFrame, - expected: pd.DataFrame, - std: pd.DataFrame, -) -> pd.DataFrame: - z = (observed - expected) / std.replace(0.0, np.nan) - return z.replace([np.inf, -np.inf], np.nan).fillna(0.0) - - -def _resolve_enrichment_mode(mode: str) -> str: - key = str(mode).strip().lower() - if key in {"zscore", "z"}: - return "zscore" - if key in {"obs-exp", "obs_minus_exp", "difference", "diff"}: - return "obs-exp" - if key in {"log2fc", "log2_fold_change", "log2"}: - return "log2fc" - raise ValueError( - "Unsupported enrichment_mode. Expected one of: " "'zscore', 'obs-exp', 'log2fc'." - ) - - -def _compute_enrichment( - observed: pd.DataFrame, - expected: pd.DataFrame, - std: pd.DataFrame, - *, - enrichment_mode: str, -) -> pd.DataFrame: - mode = _resolve_enrichment_mode(enrichment_mode) - if mode == "zscore": - return _zscore_enrichment(observed, expected, std) - if mode == "obs-exp": - return observed - expected - ratio = observed.clip(lower=1e-12) / expected.clip(lower=1e-12) - out = np.log2(ratio) - return out.replace([np.inf, -np.inf], np.nan).fillna(0.0) - - -def _perm_counts( - edge_stub: pd.DataFrame, - node_ids: np.ndarray, - labels: np.ndarray, - categories: list[str], - *, - n_perms: int, - random_state: int, -) -> np.ndarray: - rng = np.random.default_rng(int(random_state)) - n_labels = len(categories) - out = np.zeros((int(n_perms), n_labels, n_labels), dtype=float) - for i in range(int(n_perms)): - shuffled_labels = rng.permutation(labels) - shuffled_map = pd.Series(shuffled_labels, index=node_ids) - perm_edges = edge_stub.copy() - perm_edges["source_label"] = perm_edges["source_id"].map(shuffled_map) - perm_edges["target_label"] = perm_edges["target_id"].map(shuffled_map) - out[i] = _observed_matrix(perm_edges, categories).values - return out - - -def _two_sided_empirical_pvalues( - observed: pd.DataFrame, - permuted_samples: np.ndarray, -) -> pd.DataFrame: - if permuted_samples.ndim != 3: - raise ValueError("Expected permutation samples with shape (n_perms, n_labels, n_labels).") - if permuted_samples.shape[1:] != observed.shape: - raise ValueError( - "Permutation sample shape mismatch for empirical p-value computation: " - f"{permuted_samples.shape[1:]} vs {observed.shape}." - ) - center = permuted_samples.mean(axis=0) - obs_dev = np.abs(observed.values - center) - perm_dev = np.abs(permuted_samples - center[None, :, :]) - extreme = np.sum(perm_dev >= obs_dev[None, :, :], axis=0) - p = (extreme + 1.0) / (permuted_samples.shape[0] + 1.0) - return pd.DataFrame(np.clip(p, 0.0, 1.0), index=observed.index, columns=observed.columns) - - -def _condition_assignments_by_group( - node_table: pd.DataFrame, - *, - condition_column: str, - group_column: str, -) -> tuple[np.ndarray, np.ndarray]: - if group_column not in node_table.columns: - raise ValueError( - f"Missing permutation group column {group_column!r}. " - "Notebook-consistent differential permutations require a library/sample grouping key." - ) - if node_table[group_column].isna().any(): - raise ValueError( - f"Permutation group column {group_column!r} contains null values; " - "all rows must map to a valid library/sample id." - ) - - group_frame = node_table[[group_column, condition_column]].copy() - group_frame[group_column] = group_frame[group_column].astype(str) - group_frame[condition_column] = group_frame[condition_column].astype(str) - nunique = group_frame.groupby(group_column)[condition_column].nunique() - bad = nunique[nunique > 1] - if not bad.empty: - preview = ", ".join(bad.index.astype(str).tolist()[:5]) - raise ValueError( - "Permutation group ids must belong to exactly one condition. " - f"Found mixed-condition groups in {group_column!r}, e.g. {preview}." - ) - - assignments = group_frame.drop_duplicates(subset=[group_column]).reset_index(drop=True) - return ( - assignments[group_column].to_numpy(dtype=object), - assignments[condition_column].to_numpy(dtype=object), - ) - - -def run_neighborhood_enrichment( - node_table: pd.DataFrame, - spatial_table: pd.DataFrame, - *, - label_column: str, - id_column: str = "instance_id", - n_perms: int = 0, - random_state: int = 42, - undirected: bool = False, - categories: Optional[list[str]] = None, - warn_analytical: bool = True, - enrichment_mode: str = "zscore", -) -> NeighborhoodEnrichmentResult: - """Compute unweighted neighborhood enrichment with configurable enrichment score. - - Behavior: - - Edge `weight` is ignored. Observed values are raw edge counts by label pair. - - If `undirected=True`, reverse edges are explicitly mirrored (`u->v` and `v->u`). - - If `n_perms > 0`, expected/std are estimated from label permutations and `pvalues` - are two-sided empirical p-values from the same permutation null. - - If `n_perms == 0`, expected/std use an analytical approximation and a warning is emitted. - """ - if label_column not in node_table.columns: - raise ValueError(f"Missing label column {label_column!r}.") - perms = int(n_perms) - if perms < 0: - raise ValueError("n_perms must be >= 0.") - if perms == 0 and warn_analytical: - warnings.warn( - "run_neighborhood_enrichment is using analytical expected/std (n_perms=0). " - "Set n_perms>0 for Squidpy-style permutation expected/std and empirical p-values.", - UserWarning, - stacklevel=2, - ) - edges, node_labels = _prepare_edges_with_labels( - node_table, - spatial_table, - label_column=label_column, - id_column=id_column, - undirected=undirected, - ) - if categories is None: - category_list = sorted(node_labels.astype(str).unique().tolist()) - else: - category_list = [str(c) for c in categories] - - observed = _observed_matrix(edges, category_list) - n_edges = int(edges.shape[0]) - expected, std = _analytical_expected_and_std(n_edges, node_labels, category_list) - enrichment = _compute_enrichment( - observed, - expected, - std, - enrichment_mode=enrichment_mode, - ) - - pvalues: Optional[pd.DataFrame] = None - if perms > 0: - node_ids = node_labels.index.to_numpy() - base_labels = node_labels.to_numpy() - edge_stub = edges[["source_id", "target_id"]].copy() - perm_counts = _perm_counts( - edge_stub, - node_ids, - base_labels, - category_list, - n_perms=perms, - random_state=int(random_state), - ) - perm_mean = perm_counts.mean(axis=0) - perm_std = perm_counts.std(axis=0, ddof=0) - expected = pd.DataFrame(perm_mean, index=category_list, columns=category_list) - std = pd.DataFrame(perm_std, index=category_list, columns=category_list) - enrichment = _compute_enrichment( - observed, - expected, - std, - enrichment_mode=enrichment_mode, - ) - pvalues = _two_sided_empirical_pvalues(observed, perm_counts) - - return NeighborhoodEnrichmentResult( - enrichment=enrichment, - observed=observed, - expected=expected, - pvalues=pvalues, - ) - - -def run_diff_neighborhood_enrichment( - node_table: pd.DataFrame, - spatial_table: pd.DataFrame, - *, - label_column: str, - condition_column: str, - permutation_group_column: str = "sample_id", - id_column: str = "instance_id", - n_perms: int = 0, - random_state: int = 42, - undirected: bool = False, - warn_analytical: bool = True, - enrichment_mode: str = "zscore", -) -> dict[str, NeighborhoodEnrichmentResult]: - """Compute pairwise differential neighborhood enrichment between conditions. - - Differential enrichment is - `metric(condition_a) - metric(condition_b)` where metric is controlled by - `enrichment_mode`. - Per-condition enrichments are computed in analytical mode for stability/speed; when - `n_perms > 0`, differential p-values are estimated by permuting condition labels at the - library/sample group level (`permutation_group_column`), consistent with the notebook. - """ - if condition_column not in node_table.columns: - raise ValueError(f"Missing condition column {condition_column!r}.") - perms = int(n_perms) - if perms < 0: - raise ValueError("n_perms must be >= 0.") - if warn_analytical: - warnings.warn( - "run_diff_neighborhood_enrichment computes per-condition baselines with the " - "analytical neighborhood model; n_perms controls only differential p-value permutations.", - UserWarning, - stacklevel=2, - ) - category_list = sorted(node_table[label_column].astype(str).unique().tolist()) - per_condition: dict[str, NeighborhoodEnrichmentResult] = {} - for cond, sub_nodes in node_table.groupby(condition_column): - per_condition[str(cond)] = run_neighborhood_enrichment( - sub_nodes, - spatial_table, - label_column=label_column, - id_column=id_column, - n_perms=0, - random_state=random_state, - undirected=undirected, - categories=category_list, - warn_analytical=False, - enrichment_mode=enrichment_mode, - ) - - cond_keys = sorted(per_condition.keys()) - outputs: dict[str, NeighborhoodEnrichmentResult] = {} - for left, right in combinations(cond_keys, 2): - left_res = per_condition[left] - right_res = per_condition[right] - pair_nodes = node_table[ - node_table[condition_column].astype(str).isin([left, right]) - ].copy() - diff_enrichment = left_res.enrichment - right_res.enrichment - diff_observed = left_res.observed - right_res.observed - diff_expected = left_res.expected - right_res.expected - - pvalues: Optional[pd.DataFrame] = None - if perms > 0: - group_ids, group_conditions = _condition_assignments_by_group( - pair_nodes, - condition_column=condition_column, - group_column=permutation_group_column, - ) - rng = np.random.default_rng(int(random_state)) - n_labels = len(category_list) - perm_scores = np.zeros((perms, n_labels, n_labels), dtype=float) - pair_group_ids = pair_nodes[permutation_group_column].astype(str) - for i in range(perms): - permuted_group_conditions = rng.permutation(group_conditions) - group_to_condition = dict(zip(group_ids, permuted_group_conditions)) - perm_nodes = pair_nodes.copy() - perm_nodes[condition_column] = pair_group_ids.map(group_to_condition).astype(str) - perm_left = perm_nodes[perm_nodes[condition_column].astype(str) == left] - perm_right = perm_nodes[perm_nodes[condition_column].astype(str) == right] - left_perm_res = run_neighborhood_enrichment( - perm_left, - spatial_table, - label_column=label_column, - id_column=id_column, - n_perms=0, - random_state=random_state, - undirected=undirected, - categories=category_list, - warn_analytical=False, - enrichment_mode=enrichment_mode, - ) - right_perm_res = run_neighborhood_enrichment( - perm_right, - spatial_table, - label_column=label_column, - id_column=id_column, - n_perms=0, - random_state=random_state, - undirected=undirected, - categories=category_list, - warn_analytical=False, - enrichment_mode=enrichment_mode, - ) - perm_scores[i] = (left_perm_res.enrichment - right_perm_res.enrichment).values - obs = diff_enrichment.values - perm_mean = perm_scores.mean(axis=0) - obs_dev = np.abs(obs - perm_mean) - perm_dev = np.abs(perm_scores - perm_mean[None, :, :]) - extreme = np.sum(perm_dev >= obs_dev[None, :, :], axis=0) - pvals = (extreme + 1.0) / (perm_scores.shape[0] + 1.0) - pvalues = pd.DataFrame( - np.clip(pvals, 0.0, 1.0), - index=diff_enrichment.index, - columns=diff_enrichment.columns, - ) - - outputs[f"{left}_{right}"] = NeighborhoodEnrichmentResult( - enrichment=diff_enrichment, - observed=diff_observed, - expected=diff_expected, - pvalues=pvalues, - ) - return outputs diff --git a/src/grass_mil/interpretability/tier2/ripley.py b/src/grass_mil/interpretability/tier2/ripley.py deleted file mode 100644 index 14a1e79..0000000 --- a/src/grass_mil/interpretability/tier2/ripley.py +++ /dev/null @@ -1,227 +0,0 @@ -"""Ripley's K/L and the bivariate cross-K. - -For a pattern of intensity :math:`\\lambda` on a region of area :math:`|A|`, - -.. math:: K(r) = \\frac{1}{\\lambda^2 |A|} \\sum_i \\sum_{j \\neq i} - \\mathbf{1}[d_{ij} \\le r], \\qquad L(r) = \\sqrt{K(r)/\\pi} - -reported in the centred form :math:`L(r) - r`, which is zero under complete -spatial randomness, positive under clustering and negative under regularity. -The bivariate form replaces the pair count with cross-type pairs. - -Ported from ``_cross_L_for_region`` / ``_aggregate_cross_L``, keeping the -self-pair correction (a point is not its own neighbour, so the count loses -``n_a`` and the denominator is ``n_a(n_a - 1)``) and the region-equal -aggregation. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Dict, Iterable, List, Optional, Sequence, Tuple - -import numpy as np -import pandas as pd - -__all__ = ["RipleyResult", "resolve_radii", "ripley_cross_l", "aggregate_ripley"] - -PairKey = Tuple[str, str] - - -@dataclass(frozen=True) -class RipleyResult: - """Centred cross-L curves per label pair. - - Attributes: - radii: Radii the curves are evaluated at. - curves: ``(label_a, label_b) -> L(r) - r``. - pair_counts: ``(label_a, label_b) -> (n_a, n_b)``. - n_groups: Number of groups contributing to each pair. - """ - - radii: np.ndarray - curves: Dict[PairKey, np.ndarray] - pair_counts: Dict[PairKey, Tuple[int, int]] = field(default_factory=dict) - n_groups: Dict[PairKey, int] = field(default_factory=dict) - - def to_frame(self) -> pd.DataFrame: - """Curves as a radius-indexed frame, one column per pair.""" - return pd.DataFrame( - {f"{a}__{b}": curve for (a, b), curve in self.curves.items()}, - index=pd.Index(self.radii, name="radius"), - ) - - -def resolve_radii( - coords: np.ndarray, - *, - n_radii: int = 50, - max_fraction: float = 0.25, -) -> np.ndarray: - """Radius grid up to ``max_fraction`` of the bounding-box diagonal.""" - coords = np.asarray(coords, dtype=float) - if coords.size == 0: - return np.linspace(0.0, 1.0, n_radii) - spans = coords.max(axis=0) - coords.min(axis=0) - diagonal = float(np.sqrt(float(spans @ spans))) - limit = max(diagonal * float(max_fraction), 1e-9) - return np.linspace(limit / n_radii, limit, n_radii) - - -def ripley_cross_l( - coords: np.ndarray, - labels: Sequence[str], - *, - pairs: Iterable[PairKey], - radii: np.ndarray, -) -> RipleyResult: - """Centred cross-L for each label pair within one group.""" - from scipy.spatial import cKDTree - - coords = np.asarray(coords, dtype=float) - labels = np.asarray([str(v) for v in labels]) - radii = np.asarray(radii, dtype=float) - - spans = coords.max(axis=0) - coords.min(axis=0) if coords.size else np.zeros(2) - area = max(float(np.prod(spans)), 1e-10) - - index_map: Dict[str, np.ndarray] = {} - tree_map = {} - for label in set(labels.tolist()): - idx = np.where(labels == label)[0] - if idx.size >= 2: - index_map[label] = idx - tree_map[label] = cKDTree(coords[idx]) - - curves: Dict[PairKey, np.ndarray] = {} - counts: Dict[PairKey, Tuple[int, int]] = {} - for a, b in pairs: - a, b = str(a), str(b) - if a not in tree_map or b not in tree_map: - curves[(a, b)] = np.full(radii.shape, np.nan) - counts[(a, b)] = (0, 0) - continue - n_a, n_b = int(index_map[a].size), int(index_map[b].size) - pair_counts = tree_map[a].count_neighbors(tree_map[b], radii).astype(float) - if a == b: - # A point is not its own neighbour. - pair_counts = pair_counts - n_a - denom = n_a * (n_a - 1) - else: - denom = n_a * n_b - k = (area / max(denom, 1)) * pair_counts - curves[(a, b)] = np.sqrt(np.maximum(k, 0.0) / np.pi) - radii - counts[(a, b)] = (n_a, n_b) - return RipleyResult(radii=radii, curves=curves, pair_counts=counts) - - -def aggregate_ripley( - node_table: pd.DataFrame, - *, - label_column: str, - group_column: Optional[str] = None, - x_column: str = "center_x", - y_column: str = "center_y", - pairs: Optional[Iterable[PairKey]] = None, - n_radii: int = 50, - max_fraction: float = 0.25, - min_count: int = 5, - radius_source: str = "median", - radii: Optional[np.ndarray] = None, -) -> RipleyResult: - """Region-equal mean of the centred cross-L curves. - - Args: - radius_source: How to fix the shared radius grid across groups -- - ``median`` (of per-group extents), ``max``, or ``first``. The - legacy implementation used the first group encountered, which made - the grid depend on iteration order; ``median`` is the default here. - min_count: Groups with fewer than this many points of a label - contribute NaN for that pair rather than a noisy curve. - """ - for column in (label_column, x_column, y_column): - if column not in node_table.columns: - raise ValueError(f"node_table is missing required column {column!r}.") - - if group_column is None: - groups: List[Tuple[object, pd.DataFrame]] = [("__all__", node_table)] - else: - if group_column not in node_table.columns: - raise ValueError(f"node_table is missing group column {group_column!r}.") - groups = list(node_table.groupby(group_column)) - - if pairs is None: - labels = sorted({str(v) for v in node_table[label_column]}) - pairs = [(a, b) for a in labels for b in labels] - pairs = [(str(a), str(b)) for a, b in pairs] - - if radii is None: - radii = _shared_radii( - groups, - x_column=x_column, - y_column=y_column, - n_radii=n_radii, - max_fraction=max_fraction, - radius_source=radius_source, - ) - radii = np.asarray(radii, dtype=float) - - stacked: Dict[PairKey, List[np.ndarray]] = {pair: [] for pair in pairs} - totals: Dict[PairKey, Tuple[int, int]] = {pair: (0, 0) for pair in pairs} - for _, frame in groups: - coords = frame[[x_column, y_column]].to_numpy(dtype=float) - if coords.shape[0] < 2: - continue - result = ripley_cross_l(coords, frame[label_column].tolist(), pairs=pairs, radii=radii) - for pair, curve in result.curves.items(): - n_a, n_b = result.pair_counts[pair] - if n_a < min_count or n_b < min_count or not np.isfinite(curve).any(): - continue - stacked[pair].append(curve) - totals[pair] = (totals[pair][0] + n_a, totals[pair][1] + n_b) - - curves: Dict[PairKey, np.ndarray] = {} - n_groups: Dict[PairKey, int] = {} - for pair, values in stacked.items(): - n_groups[pair] = len(values) - if not values: - curves[pair] = np.full(radii.shape, np.nan) - continue - # Region-equal: every group counts once regardless of its size. - curves[pair] = np.nanmean(np.vstack(values), axis=0) - return RipleyResult(radii=radii, curves=curves, pair_counts=totals, n_groups=n_groups) - - -def _shared_radii( - groups: Sequence[Tuple[object, pd.DataFrame]], - *, - x_column: str, - y_column: str, - n_radii: int, - max_fraction: float, - radius_source: str, -) -> np.ndarray: - extents: List[float] = [] - for _, frame in groups: - coords = frame[[x_column, y_column]].to_numpy(dtype=float) - if coords.shape[0] < 2: - continue - spans = coords.max(axis=0) - coords.min(axis=0) - extents.append(float(np.sqrt(float(spans @ spans)))) - if not extents: - return np.linspace(0.0, 1.0, n_radii) - - source = str(radius_source).strip().lower() - if source == "median": - diagonal = float(np.median(extents)) - elif source == "max": - diagonal = float(np.max(extents)) - elif source == "first": - diagonal = float(extents[0]) - else: - raise ValueError( - f"Unsupported radius_source '{radius_source}'. " - "Expected one of: 'median', 'max', 'first'." - ) - limit = max(diagonal * float(max_fraction), 1e-9) - return np.linspace(limit / n_radii, limit, n_radii) diff --git a/src/grass_mil/interpretability/tier2/tissue_graph.py b/src/grass_mil/interpretability/tier2/tissue_graph.py deleted file mode 100644 index c29659e..0000000 --- a/src/grass_mil/interpretability/tier2/tissue_graph.py +++ /dev/null @@ -1,207 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Dict - -import numpy as np -import pandas as pd - -if TYPE_CHECKING: - import plotly.graph_objects as go - - -@dataclass(frozen=True) -class TissueGraphView: - nodes: pd.DataFrame - edges: pd.DataFrame - metadata: Dict[str, Any] = field(default_factory=dict) - - -def _require_columns(frame: pd.DataFrame, required: list[str], frame_name: str) -> None: - missing = [col for col in required if str(col) not in frame.columns] - if missing: - missing_text = ", ".join(sorted(missing)) - raise ValueError(f"{frame_name} is missing required columns: {missing_text}.") - - -def prepare_tissue_graph_view( - node_table: pd.DataFrame, - spatial_table: pd.DataFrame, - *, - sample_column: str, - sample_value: Any, - id_column: str, - x_column: str = "center_x", - y_column: str = "center_y", - label_column: str = "niche_label", - coerce_ids_to_str: bool = True, - include_edge_distances: bool = True, -) -> TissueGraphView: - if sample_value is None or str(sample_value).strip() == "": - raise ValueError("sample_value is required for tissue graph selection.") - - _require_columns( - node_table, - [str(sample_column), str(id_column), str(x_column), str(y_column), str(label_column)], - "node_table", - ) - _require_columns(spatial_table, ["source_id", "target_id"], "spatial_table") - - selected = node_table[node_table[str(sample_column)].astype(str) == str(sample_value)].copy() - if selected.empty: - raise ValueError( - "No rows found for tissue graph selection: " f"{sample_column}={sample_value!r}." - ) - - node_cols = [ - str(id_column), - str(sample_column), - str(x_column), - str(y_column), - str(label_column), - ] - nodes = selected[node_cols].copy() - if coerce_ids_to_str: - nodes[str(id_column)] = nodes[str(id_column)].astype(str) - - edges = spatial_table[["source_id", "target_id"]].copy() - if include_edge_distances and "distance" in spatial_table.columns: - edges["distance"] = spatial_table["distance"] - if coerce_ids_to_str: - edges["source_id"] = edges["source_id"].astype(str) - edges["target_id"] = edges["target_id"].astype(str) - - node_ids = set(nodes[str(id_column)].astype(str).tolist()) - edges = edges[ - edges["source_id"].astype(str).isin(node_ids) - & edges["target_id"].astype(str).isin(node_ids) - ].copy() - - coords = nodes.set_index(str(id_column))[[str(x_column), str(y_column)]] - edges = ( - edges.join( - coords.rename(columns={str(x_column): "x0", str(y_column): "y0"}), on="source_id" - ) - .join(coords.rename(columns={str(x_column): "x1", str(y_column): "y1"}), on="target_id") - .dropna(subset=["x0", "y0", "x1", "y1"]) - ) - - edge_cols = ["source_id", "target_id", "x0", "y0", "x1", "y1"] - if "distance" in edges.columns: - edge_cols.append("distance") - edges = edges[edge_cols].reset_index(drop=True) - - niche_counts = ( - nodes[str(label_column)].astype(str).value_counts().sort_index().astype(int).to_dict() - ) - metadata: Dict[str, Any] = { - "sample_column": str(sample_column), - "sample_value": str(sample_value), - "id_column": str(id_column), - "x_column": str(x_column), - "y_column": str(y_column), - "label_column": str(label_column), - "node_count": int(len(nodes)), - "edge_count": int(len(edges)), - "niche_counts": niche_counts, - } - return TissueGraphView(nodes=nodes.reset_index(drop=True), edges=edges, metadata=metadata) - - -def build_tissue_graph_figure( - view: TissueGraphView, - *, - show_edges: bool = True, - node_size: float = 5.0, - edge_width: float = 0.5, - edge_opacity: float = 0.25, - colorscale: str = "Viridis", - reverse_y: bool = True, - title_prefix: str = "Tissue Graph", -) -> "go.Figure": - try: - import plotly.graph_objects as go - except ModuleNotFoundError as exc: - raise ModuleNotFoundError( - "plotly is required to build tissue-graph figures. " - "Install it with: pip install plotly" - ) from exc - - nodes = view.nodes.copy() - metadata = dict(view.metadata) - - label_column = str(metadata.get("label_column", "niche_label")) - id_column = str(metadata.get("id_column", "instance_id")) - x_column = str(metadata.get("x_column", "center_x")) - y_column = str(metadata.get("y_column", "center_y")) - sample_column = str(metadata.get("sample_column", "sample_id")) - sample_value = str(metadata.get("sample_value", "unknown")) - - _require_columns(nodes, [id_column, x_column, y_column, label_column], "view.nodes") - - label_series = nodes[label_column].astype(str) - label_codes = pd.Categorical(label_series).codes - - fig = go.Figure() - if show_edges and not view.edges.empty: - edges = view.edges - _require_columns(edges, ["x0", "y0", "x1", "y1"], "view.edges") - line_x = np.empty(len(edges) * 3) - line_y = np.empty(len(edges) * 3) - line_x[0::3] = edges["x0"].to_numpy(dtype=float) - line_x[1::3] = edges["x1"].to_numpy(dtype=float) - line_x[2::3] = np.nan - line_y[0::3] = edges["y0"].to_numpy(dtype=float) - line_y[1::3] = edges["y1"].to_numpy(dtype=float) - line_y[2::3] = np.nan - fig.add_trace( - go.Scattergl( - x=line_x, - y=line_y, - mode="lines", - line={ - "width": float(edge_width), - "color": f"rgba(120,120,120,{float(edge_opacity)})", - }, - hoverinfo="skip", - name="edges", - ) - ) - - fig.add_trace( - go.Scattergl( - x=nodes[x_column].to_numpy(dtype=float), - y=nodes[y_column].to_numpy(dtype=float), - mode="markers", - marker={ - "size": float(node_size), - "color": label_codes, - "colorscale": colorscale, - "showscale": True, - "colorbar": {"title": "niche"}, - }, - customdata=np.stack( - [ - nodes[id_column].astype(str).to_numpy(), - label_series.to_numpy(), - ], - axis=1, - ), - hovertemplate=( - f"{id_column}: %{{customdata[0]}}
" - f"{label_column}: %{{customdata[1]}}" - "" - ), - name="cells", - ) - ) - - fig.update_layout( - title=f"{title_prefix} ({sample_column}={sample_value})", - xaxis_title=x_column, - yaxis_title=y_column, - template="plotly_white", - ) - if reverse_y: - fig.update_yaxes(autorange="reversed") - return fig diff --git a/src/grass_mil/models/augmentations.py b/src/grass_mil/models/augmentations.py index 8a3a10d..470e4c6 100644 --- a/src/grass_mil/models/augmentations.py +++ b/src/grass_mil/models/augmentations.py @@ -77,7 +77,9 @@ def drop_importance(batch, *, mu: float, p_lambda: float, unassigned_index: int out = batch.clone() imp = degree_importance(out.edge_index, int(out.num_nodes)) p_node = torch.clamp((1.0 - imp) * mu, max=p_lambda) - out = mask_nodes(out, _rand(p_node.shape, generator, p_node.device) <= p_node, unassigned_index) + out = mask_nodes( + out, _rand(p_node.shape, generator, p_node.device) <= p_node, unassigned_index + ) if out.edge_index.numel(): e = (imp[out.edge_index[0]] + imp[out.edge_index[1]]) / 2.0 e = torch.clamp((e - e.mean()) / (e.max() - e.mean() + 1e-8), min=0.0) @@ -143,7 +145,13 @@ def build(self, *, feature_names: Sequence[str], unassigned_index: int | None) - columns = _columns(self.noise_features, feature_names) def view(batch, generator): - out = drop_importance(batch, mu=self.mu, p_lambda=self.p_lambda, unassigned_index=unassigned_index, generator=generator) + out = drop_importance( + batch, + mu=self.mu, + p_lambda=self.p_lambda, + unassigned_index=unassigned_index, + generator=generator, + ) return feature_noise(out, std=self.noise_std, columns=columns, generator=generator) return view diff --git a/src/grass_mil/models/encoder.py b/src/grass_mil/models/encoder.py index f8c1355..b21d868 100644 --- a/src/grass_mil/models/encoder.py +++ b/src/grass_mil/models/encoder.py @@ -98,7 +98,10 @@ def __init__(self, feature_dim: int, cell_types: tuple[str, ...], cfg: EncoderCo if cfg.unassigned_cell_type is None: rows, self.unassigned_index = len(cell_types) + 1, len(cell_types) elif cfg.unassigned_cell_type in cell_types: - rows, self.unassigned_index = len(cell_types), cell_types.index(cfg.unassigned_cell_type) + rows, self.unassigned_index = ( + len(cell_types), + cell_types.index(cfg.unassigned_cell_type), + ) else: raise ValueError( f"model.encoder.unassigned_cell_type '{cfg.unassigned_cell_type}' is not a " @@ -121,25 +124,36 @@ def forward(self, x: torch.Tensor, cell_type: torch.Tensor | None) -> torch.Tens class GINEEncoder(nn.Module): - def __init__(self, cfg: EncoderConfig, *, feature_dim: int, cell_types: tuple[str, ...], edge_dim: int) -> None: + def __init__( + self, cfg: EncoderConfig, *, feature_dim: int, cell_types: tuple[str, ...], edge_dim: int + ) -> None: super().__init__() self.cfg = cfg if cfg.edge_feature is not None and not 0 <= cfg.edge_feature < edge_dim: - raise ValueError(f"model.encoder.edge_feature={cfg.edge_feature} but edges have {edge_dim} feature(s).") + raise ValueError( + f"model.encoder.edge_feature={cfg.edge_feature} but edges have {edge_dim} feature(s)." + ) conv_edge_dim = 1 if cfg.edge_feature is not None else edge_dim h, k = cfg.hidden_dim, cfg.mlp_ratio * cfg.hidden_dim self.input_proj = NodeInput(feature_dim, cell_types, cfg) self.layers = nn.ModuleList( - GINEConv(nn.Sequential(nn.Linear(h, k), nn.ReLU(), nn.Linear(k, h)), edge_dim=conv_edge_dim, train_eps=cfg.train_eps) + GINEConv( + nn.Sequential(nn.Linear(h, k), nn.ReLU(), nn.Linear(k, h)), + edge_dim=conv_edge_dim, + train_eps=cfg.train_eps, + ) for _ in range(cfg.num_layers) ) self.norms = nn.ModuleList( - _NORM[cfg.norm](h) if cfg.norm != "none" else nn.Identity() for _ in range(cfg.num_layers) + _NORM[cfg.norm](h) if cfg.norm != "none" else nn.Identity() + for _ in range(cfg.num_layers) ) self.act = _ACT[cfg.act]() self.dropout = nn.Dropout(cfg.dropout) jk_dim = h * (cfg.num_layers + 1) if cfg.jk == "concat" else h - self.output_proj = nn.Identity() if jk_dim == cfg.out_dim else nn.Linear(jk_dim, cfg.out_dim) + self.output_proj = ( + nn.Identity() if jk_dim == cfg.out_dim else nn.Linear(jk_dim, cfg.out_dim) + ) self.output_dim = cfg.out_dim @property diff --git a/src/grass_mil/models/mil.py b/src/grass_mil/models/mil.py index 9998650..bda616f 100644 --- a/src/grass_mil/models/mil.py +++ b/src/grass_mil/models/mil.py @@ -17,7 +17,14 @@ import torch from torch import nn -__all__ = ["HeadConfig", "AttentionConfig", "InstanceHead", "GatedAttention", "attention_pool", "pool_by_bag"] +__all__ = [ + "HeadConfig", + "AttentionConfig", + "InstanceHead", + "GatedAttention", + "attention_pool", + "pool_by_bag", +] @dataclass(frozen=True) @@ -96,7 +103,9 @@ class GatedAttention(nn.Module): def __init__(self, input_dim: int, num_classes: int, cfg: AttentionConfig) -> None: super().__init__() p, h = cfg.projection_dim, cfg.hidden_dim - self.projection = nn.Sequential(nn.Linear(input_dim, p), nn.LeakyReLU(), nn.Dropout(cfg.dropout)) + self.projection = nn.Sequential( + nn.Linear(input_dim, p), nn.LeakyReLU(), nn.Dropout(cfg.dropout) + ) self.attention_a = nn.Sequential(nn.Linear(p, h), nn.Tanh(), nn.Dropout(cfg.dropout)) self.attention_b = nn.Sequential(nn.Linear(p, h), nn.Sigmoid(), nn.Dropout(cfg.dropout)) self.attention_c = nn.Linear(h, num_classes) @@ -117,7 +126,9 @@ def space(self, x: torch.Tensor) -> torch.Tensor: return self.attention_a[0](self.projection[1](self.projection[0](x))) -def attention_pool(logits: torch.Tensor, scores: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: +def attention_pool( + logits: torch.Tensor, scores: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: """Pool one bag: ``(bag_logits [C], attention weights [n, C])``.""" if scores.shape != logits.shape: raise ValueError( diff --git a/src/grass_mil/models/objectives.py b/src/grass_mil/models/objectives.py index 66f20c9..49853d3 100644 --- a/src/grass_mil/models/objectives.py +++ b/src/grass_mil/models/objectives.py @@ -20,7 +20,14 @@ import torch import torch.nn.functional as F -__all__ = ["TARGET_TYPES", "output_dim", "Objective", "region_metrics", "concordance_index", "class_weights_from_labels"] +__all__ = [ + "TARGET_TYPES", + "output_dim", + "Objective", + "region_metrics", + "concordance_index", + "class_weights_from_labels", +] TARGET_TYPES = ("classification", "regression", "survival") @@ -56,18 +63,26 @@ def __init__(self, task: Mapping[str, Any], class_weight: torch.Tensor | None = if self.kind == "survival" and len(task.get("target_columns") or []) != 2: raise ValueError("Survival tasks need task.target_columns: [time, event].") - def __call__(self, logits: torch.Tensor, targets: torch.Tensor, weights: torch.Tensor | None = None) -> torch.Tensor: + def __call__( + self, logits: torch.Tensor, targets: torch.Tensor, weights: torch.Tensor | None = None + ) -> torch.Tensor: logits = logits.float() if self.kind == "classification": weight = self.class_weight.to(logits.device) if self.class_weight is not None else None - per_bag = F.cross_entropy(logits, targets[:, 0].long(), weight=weight, reduction="none") + per_bag = F.cross_entropy( + logits, targets[:, 0].long(), weight=weight, reduction="none" + ) elif self.kind == "regression": if self.loss_name == "huber": - per_bag = F.huber_loss(logits, targets, delta=self.huber_delta, reduction="none").mean(1) + per_bag = F.huber_loss( + logits, targets, delta=self.huber_delta, reduction="none" + ).mean(1) else: per_bag = F.mse_loss(logits, targets, reduction="none").mean(1) else: - return cox_loss(logits[:, 0], targets[:, 0], targets[:, 1], weights, self.cox_regularization) + return cox_loss( + logits[:, 0], targets[:, 0], targets[:, 1], weights, self.cox_regularization + ) if weights is not None: per_bag = per_bag * weights.view(-1).to(per_bag.device) return per_bag.mean() @@ -87,7 +102,11 @@ def cox_loss( """ order = torch.argsort(time, descending=True) risk, event = risk[order], event[order].float() - w = torch.ones_like(risk) if weights is None else weights.view(-1)[order].float().to(risk.device) + w = ( + torch.ones_like(risk) + if weights is None + else weights.view(-1)[order].float().to(risk.device) + ) # Cumulative log-sum-exp over the risk set (everyone with time >= t_i). log_risk_set = torch.logcumsumexp(risk + torch.log(w), dim=0) n_events = (event * w).sum() @@ -142,7 +161,13 @@ def region_metrics(kind: str, logits: torch.Tensor, targets: torch.Tensor) -> di out["auroc"] = float( roc_auc_score(y, prob[:, 1]) if prob.shape[1] == 2 - else roc_auc_score(y, prob, multi_class="ovr", average="macro", labels=list(range(prob.shape[1]))) + else roc_auc_score( + y, + prob, + multi_class="ovr", + average="macro", + labels=list(range(prob.shape[1])), + ) ) except ValueError: pass @@ -155,8 +180,17 @@ def region_metrics(kind: str, logits: torch.Tensor, targets: torch.Tensor) -> di out["r2"] = out["pearson"] = float("nan") if bool((ss_tot > 0).all()): out["r2"] = float((1 - (err**2).sum(0) / ss_tot).mean()) - out["pearson"] = float(np.mean([np.corrcoef(logits[:, j], targets[:, j])[0, 1] for j in range(targets.shape[1])])) + out["pearson"] = float( + np.mean( + [ + np.corrcoef(logits[:, j], targets[:, j])[0, 1] + for j in range(targets.shape[1]) + ] + ) + ) else: out["loss"] = float(cox_loss(logits[:, 0], targets[:, 0], targets[:, 1])) - out["c_index"] = concordance_index(logits[:, 0].numpy(), targets[:, 0].numpy(), targets[:, 1].numpy()) + out["c_index"] = concordance_index( + logits[:, 0].numpy(), targets[:, 0].numpy(), targets[:, 1].numpy() + ) return out diff --git a/src/grass_mil/models/optimization.py b/src/grass_mil/models/optimization.py index 3cd31c3..9e03c0a 100644 --- a/src/grass_mil/models/optimization.py +++ b/src/grass_mil/models/optimization.py @@ -12,7 +12,9 @@ def _plain(cfg: Any) -> dict: - return dict(OmegaConf.to_container(cfg, resolve=True)) if OmegaConf.is_config(cfg) else dict(cfg) + return ( + dict(OmegaConf.to_container(cfg, resolve=True)) if OmegaConf.is_config(cfg) else dict(cfg) + ) def build_optimizer(cfg: Mapping[str, Any], groups: list[dict]) -> torch.optim.Optimizer: @@ -51,4 +53,6 @@ def build_scheduler( warmup = torch.optim.lr_scheduler.LinearLR( optimizer, start_factor=1.0 / (warmup_steps + 1), end_factor=1.0, total_iters=warmup_steps ) - return torch.optim.lr_scheduler.SequentialLR(optimizer, [warmup, main], milestones=[warmup_steps]) + return torch.optim.lr_scheduler.SequentialLR( + optimizer, [warmup, main], milestones=[warmup_steps] + ) diff --git a/src/grass_mil/models/pretraining.py b/src/grass_mil/models/pretraining.py index d12093b..48d6c1f 100644 --- a/src/grass_mil/models/pretraining.py +++ b/src/grass_mil/models/pretraining.py @@ -62,7 +62,9 @@ def from_dict(cls, cfg: Mapping[str, Any] | None) -> PretrainConfig: def momentum(self, step: int) -> float: progress = min(max(step, 0), self.total_steps) / max(1, self.total_steps) - return self.momentum_end - (self.momentum_end - self.momentum_start) * 0.5 * (1 + math.cos(math.pi * progress)) + return self.momentum_end - (self.momentum_end - self.momentum_start) * 0.5 * ( + 1 + math.cos(math.pi * progress) + ) class Predictor(nn.Module): @@ -103,7 +105,10 @@ def __init__( def build(self, spec: InputSpec, feature_names: list[str] | None = None) -> None: self.input_spec = spec self.encoder = GINEEncoder( - self.encoder_cfg, feature_dim=spec.feature_dim, cell_types=spec.cell_types, edge_dim=spec.edge_dim + self.encoder_cfg, + feature_dim=spec.feature_dim, + cell_types=spec.cell_types, + edge_dim=spec.edge_dim, ) self.predictor = Predictor(self.encoder.output_dim, self.pretrain.predictor_hidden_dim) self.target_encoder = copy.deepcopy(self.encoder) @@ -120,7 +125,10 @@ def setup(self, stage: str | None = None) -> None: if self.hparams.init_from_ckpt: load_encoder_weights(self.encoder, self.hparams.init_from_ckpt) self.target_encoder.load_state_dict(self.encoder.state_dict()) - if self.trainer.max_steps not in (-1, None) and self.trainer.max_steps != self.pretrain.total_steps: + if ( + self.trainer.max_steps not in (-1, None) + and self.trainer.max_steps != self.pretrain.total_steps + ): rank_zero_warn( f"trainer.max_steps={self.trainer.max_steps} differs from pretrain.total_steps=" f"{self.pretrain.total_steps}; the schedules follow pretrain.total_steps." @@ -136,11 +144,22 @@ def _loss(self, batch) -> torch.Tensor: with torch.no_grad(): y1 = self.target_encoder.node_embeddings(view1) y2 = self.target_encoder.node_embeddings(view2) - return 2.0 - F.cosine_similarity(q1, y2, dim=-1).mean() - F.cosine_similarity(q2, y1, dim=-1).mean() + return ( + 2.0 + - F.cosine_similarity(q1, y2, dim=-1).mean() + - F.cosine_similarity(q2, y1, dim=-1).mean() + ) def training_step(self, batch, batch_idx: int) -> torch.Tensor: loss = self._loss(batch) - self.log("train/loss", loss, on_step=True, on_epoch=True, prog_bar=True, batch_size=int(batch.num_graphs)) + self.log( + "train/loss", + loss, + on_step=True, + on_epoch=True, + prog_bar=True, + batch_size=int(batch.num_graphs), + ) return loss def validation_step(self, batch, batch_idx: int) -> None: @@ -169,4 +188,7 @@ def configure_optimizers(self): ) if scheduler is None: return optimizer - return {"optimizer": optimizer, "lr_scheduler": {"scheduler": scheduler, "interval": "step"}} + return { + "optimizer": optimizer, + "lr_scheduler": {"scheduler": scheduler, "interval": "step"}, + } diff --git a/src/grass_mil/models/probe.py b/src/grass_mil/models/probe.py index 6ccd87a..a5c58e3 100644 --- a/src/grass_mil/models/probe.py +++ b/src/grass_mil/models/probe.py @@ -37,7 +37,9 @@ def from_dict(cls, cfg: Mapping[str, Any] | None) -> ProbeConfig: @torch.no_grad() -def collect_embeddings(encoder: torch.nn.Module, loader, device: torch.device) -> tuple[np.ndarray, np.ndarray]: +def collect_embeddings( + encoder: torch.nn.Module, loader, device: torch.device +) -> tuple[np.ndarray, np.ndarray]: """Embeddings and inherited targets of every ego-graph of ``loader``.""" was_training = encoder.training encoder.eval() @@ -81,7 +83,9 @@ def fit_linear_head( w = coef / scaler.scale_[None, :] b = intercept - (coef * scaler.mean_[None, :] / scaler.scale_[None, :]).sum(axis=1) if coef.shape[0] not in (1, num_logits): - raise ValueError(f"The probe saw {coef.shape[0]} classes but the head has {num_logits} logits.") + raise ValueError( + f"The probe saw {coef.shape[0]} classes but the head has {num_logits} logits." + ) weight = torch.zeros(num_logits, embeddings.shape[1]) bias = torch.zeros(num_logits) if kind == "classification" and coef.shape[0] == 1: diff --git a/src/grass_mil/models/supervised.py b/src/grass_mil/models/supervised.py index 2cde1b9..edb5d1a 100644 --- a/src/grass_mil/models/supervised.py +++ b/src/grass_mil/models/supervised.py @@ -25,7 +25,14 @@ from .augmentations import FinetuneAugmentation from .checkpoints import load_encoder_weights from .encoder import EncoderConfig, GINEEncoder -from .mil import AttentionConfig, GatedAttention, HeadConfig, InstanceHead, attention_pool, pool_by_bag +from .mil import ( + AttentionConfig, + GatedAttention, + HeadConfig, + InstanceHead, + attention_pool, + pool_by_bag, +) from .objectives import Objective, class_weights_from_labels, output_dim, region_metrics from .optimization import build_optimizer, build_scheduler from .probe import ProbeConfig, collect_embeddings, fit_linear_head @@ -71,8 +78,12 @@ def from_dict(cls, cfg: Mapping[str, Any]) -> FinetuneConfig: if unknown: raise ValueError(f"Unknown finetune keys: {sorted(unknown)}") scales = dict(cfg.get("lr_scales") or {}) - if set(scales) != {"encoder", "head", "attention"} or any(v is None for v in scales.values()): - raise ValueError("finetune.lr_scales needs numeric encoder, head and attention scales.") + if set(scales) != {"encoder", "head", "attention"} or any( + v is None for v in scales.values() + ): + raise ValueError( + "finetune.lr_scales needs numeric encoder, head and attention scales." + ) cfg["lr_scales"] = {k: float(v) for k, v in scales.items()} cfg["augmentation"] = FinetuneAugmentation.from_dict(cfg.get("augmentation")) cfg["probe"] = ProbeConfig.from_dict(cfg.get("probe")) @@ -129,13 +140,18 @@ def build(self, spec: InputSpec) -> None: """Create the layers for the given node/edge inputs.""" self.input_spec = spec self.encoder = GINEEncoder( - self.encoder_cfg, feature_dim=spec.feature_dim, cell_types=spec.cell_types, edge_dim=spec.edge_dim + self.encoder_cfg, + feature_dim=spec.feature_dim, + cell_types=spec.cell_types, + edge_dim=spec.edge_dim, ) width = self.encoder.output_dim self.head = InstanceHead(width, self.num_logits, self.head_cfg) self.attention = GatedAttention(width, self.num_logits, self.attention_cfg) if self.finetune.augmentation.enabled: - self._augment = self.finetune.augmentation.build(unassigned_index=self.encoder.unassigned_index) + self._augment = self.finetune.augmentation.build( + unassigned_index=self.encoder.unassigned_index + ) def setup(self, stage: str | None = None) -> None: if self.built: @@ -162,23 +178,35 @@ def _configure_class_weights(self, dataset) -> None: def _fit_probe_head(self, datamodule) -> None: if self.task["target_type"] == "survival": - log.info("No probe for survival targets; the head starts from a random initialisation.") + log.info( + "No probe for survival targets; the head starts from a random initialisation." + ) return if not self.head.is_linear: - raise ValueError("finetune.head_init=probe needs a linear head (model.head.num_layers=1).") + raise ValueError( + "finetune.head_init=probe needs a linear head (model.head.num_layers=1)." + ) if not self.hparams.init_from_ckpt: - log.warning("Fitting the probe head on a randomly initialised encoder (no model.init_from_ckpt).") + log.warning( + "Fitting the probe head on a randomly initialised encoder (no model.init_from_ckpt)." + ) cfg = self.finetune.probe - loader = datamodule.loader(datamodule.dataset_train, roots_per_unit=cfg.roots_per_unit, seed=cfg.seed) + loader = datamodule.loader( + datamodule.dataset_train, roots_per_unit=cfg.roots_per_unit, seed=cfg.seed + ) device = self.trainer.strategy.root_device self.encoder.to(device) embeddings, targets = collect_embeddings(self.encoder, loader, device) self.encoder.cpu() - weight, bias = fit_linear_head(embeddings, targets, kind=self.task["target_type"], num_logits=self.num_logits, cfg=cfg) + weight, bias = fit_linear_head( + embeddings, targets, kind=self.task["target_type"], num_logits=self.num_logits, cfg=cfg + ) with torch.no_grad(): self.head.net[0].weight.copy_(weight) self.head.net[0].bias.copy_(bias) - log.info(f"Probe head fitted on {len(embeddings)} ego-graphs of {len(datamodule.dataset_train)} units.") + log.info( + f"Probe head fitted on {len(embeddings)} ego-graphs of {len(datamodule.dataset_train)} units." + ) def _apply_freezing(self) -> None: scales = self.finetune.lr_scales @@ -202,7 +230,9 @@ def on_save_checkpoint(self, checkpoint: dict[str, Any]) -> None: checkpoint["input_spec"] = self.input_spec.as_dict() @classmethod - def from_checkpoint(cls, path: str, map_location: str | torch.device = "cpu") -> SupervisedModule: + def from_checkpoint( + cls, path: str, map_location: str | torch.device = "cpu" + ) -> SupervisedModule: """Rebuild a trained module from its checkpoint alone (eval mode).""" ckpt = torch.load(path, map_location=map_location, weights_only=False) hparams = dict(ckpt["hyper_parameters"]) @@ -230,7 +260,9 @@ def forward(self, batch) -> torch.Tensor: def training_step(self, batch, batch_idx: int) -> None: bag_logits = self(batch).unsqueeze(0) weights = getattr(batch, "graph_w", None) - self._bag_buffer.append((bag_logits, batch.graph_y[:1].float(), None if weights is None else weights[:1])) + self._bag_buffer.append( + (bag_logits, batch.graph_y[:1].float(), None if weights is None else weights[:1]) + ) if len(self._bag_buffer) >= self.finetune.hyperbatch_size: self._optimizer_step(on_step=True) @@ -258,7 +290,14 @@ def _optimizer_step(self, *, on_step: bool) -> None: scheduler = self.lr_schedulers() if scheduler is not None: scheduler.step() - self.log("train/loss", loss.detach(), on_step=on_step, on_epoch=True, prog_bar=True, batch_size=len(targets)) + self.log( + "train/loss", + loss.detach(), + on_step=on_step, + on_epoch=True, + prog_bar=True, + batch_size=len(targets), + ) def on_before_optimizer_step(self, optimizer) -> None: # Called after mixed-precision unscaling, so the norm is the true one. @@ -270,7 +309,11 @@ def configure_optimizers(self): scales = self.finetune.lr_scales lr = float(self.hparams.optim["lr"]) groups = [] - for name, module in (("encoder", self.encoder), ("head", self.head), ("attention", self.attention)): + for name, module in ( + ("encoder", self.encoder), + ("head", self.head), + ("attention", self.attention), + ): params = [p for p in module.parameters() if p.requires_grad] if params: group = {"params": params, "lr": lr * scales[name], "name": name} @@ -283,11 +326,17 @@ def configure_optimizers(self): n_units = len(self.trainer.datamodule.dataset_train) steps = int(self.trainer.max_epochs) * math.ceil(n_units / self.finetune.hyperbatch_size) scheduler = build_scheduler( - self.hparams.scheduler, optimizer, total_steps=steps, warmup_steps=self.finetune.warmup_steps + self.hparams.scheduler, + optimizer, + total_steps=steps, + warmup_steps=self.finetune.warmup_steps, ) if scheduler is None: return optimizer - return {"optimizer": optimizer, "lr_scheduler": {"scheduler": scheduler, "interval": "step"}} + return { + "optimizer": optimizer, + "lr_scheduler": {"scheduler": scheduler, "interval": "step"}, + } # ----------------------------------------------------------- evaluation def _eval_step(self, batch) -> None: @@ -325,10 +374,14 @@ def _finish_eval(self, stage: str) -> None: "scores": [t for g in gathered for t in g["scores"]], "targets": {k: v for g in gathered for k, v in g["targets"].items()}, } - units, bag_logits = pool_by_bag(torch.cat(buf["logits"]), torch.cat(buf["scores"]), buf["unit_ids"]) + units, bag_logits = pool_by_bag( + torch.cat(buf["logits"]), torch.cat(buf["scores"]), buf["unit_ids"] + ) targets = torch.stack([buf["targets"][u] for u in units]) metrics = region_metrics(self.task["target_type"], bag_logits, targets) for name, value in metrics.items(): - self.log(f"{stage}/{name}", value, prog_bar=name in ("loss", "auroc", "c_index", "rmse")) + self.log( + f"{stage}/{name}", value, prog_bar=name in ("loss", "auroc", "c_index", "rmse") + ) if stage == "test": self.test_predictions = {"unit_ids": units, "logits": bag_logits, "targets": targets} diff --git a/src/grass_mil/nested_cv.py b/src/grass_mil/nested_cv.py index 2c868ed..26b5b2a 100644 --- a/src/grass_mil/nested_cv.py +++ b/src/grass_mil/nested_cv.py @@ -56,7 +56,9 @@ def _reduce(epochs: Sequence[int], how: str, minimum: int) -> int: return max(int(minimum), int(round(reducers[how](epochs)))) -def _run(cfg: DictConfig, *, train_ids, val_ids, test_ids, out: Path, max_epochs: int | None = None) -> dict[str, Any]: +def _run( + cfg: DictConfig, *, train_ids, val_ids, test_ids, out: Path, max_epochs: int | None = None +) -> dict[str, Any]: """One training run on explicit unit lists (resumable through ``done.json``).""" done = out / "done.json" if done.exists() and cfg.nested_cv.resume: @@ -73,8 +75,17 @@ def _run(cfg: DictConfig, *, train_ids, val_ids, test_ids, out: Path, max_epochs run_cfg.trainer.max_epochs = run_cfg.trainer.min_epochs = int(max_epochs) out.mkdir(parents=True, exist_ok=True) metrics, _ = train(run_cfg) - epochs, used = best_epoch(read_epoch_metrics(out / "metrics_per_epoch.jsonl"), cfg.task.selection.metric, cfg.task.selection.mode) - result = {"metrics": metrics, "best_epoch": epochs, "selection_metric_used": used, "output_dir": str(out)} + epochs, used = best_epoch( + read_epoch_metrics(out / "metrics_per_epoch.jsonl"), + cfg.task.selection.metric, + cfg.task.selection.mode, + ) + result = { + "metrics": metrics, + "best_epoch": epochs, + "selection_metric_used": used, + "output_dir": str(out), + } done.write_text(json.dumps(result, indent=1)) return result @@ -92,13 +103,26 @@ def run_nested_cv(cfg: DictConfig) -> dict[str, Any]: if datamodule.read_labels() is None: raise ValueError("Cross-validation needs labels (data.labels.file).") meta = datamodule.unit_table(units) - labelled = {e.unit_id: units.label_id(e) for e in units.entries if units.label_id(e) is not None} - groups = {u: (meta.loc[u, "patient_id"] if "patient_id" in meta.columns else meta.loc[u, "region_id"]) for u in labelled} - classes = {u: class_index(units.labels.target(i)[0]) for u, i in labelled.items()} if kind == "classification" else None + labelled = { + e.unit_id: units.label_id(e) for e in units.entries if units.label_id(e) is not None + } + groups = { + u: ( + meta.loc[u, "patient_id"] if "patient_id" in meta.columns else meta.loc[u, "region_id"] + ) + for u in labelled + } + classes = ( + {u: class_index(units.labels.target(i)[0]) for u, i in labelled.items()} + if kind == "classification" + else None + ) def _units_of(ids: Sequence[str]) -> list[str]: wanted = set(map(str, ids)) - return sorted(u for u in labelled if {u, meta.loc[u, "region_id"], meta.loc[u, "sample_id"]} & wanted) + return sorted( + u for u in labelled if {u, meta.loc[u, "region_id"], meta.loc[u, "sample_id"]} & wanted + ) folds = [_units_of(f) for f in load_folds(str(ncfg.folds_file))] if any(not f for f in folds): @@ -116,15 +140,26 @@ def _units_of(ids: Sequence[str]) -> list[str]: continue test = folds[i] dev = sorted(u for j, f in enumerate(folds) if j != i for u in f) - inner = inner_folds(dev, groups=groups, classes=classes, k=int(ncfg.inner_k), seed=int(ncfg.inner_seed)) + inner = inner_folds( + dev, groups=groups, classes=classes, k=int(ncfg.inner_k), seed=int(ncfg.inner_seed) + ) epochs = [] for j, (tr, va) in enumerate(inner): res = _run(cfg, train_ids=tr, val_ids=va, test_ids=[], out=fold_dir / f"inner_{j}") if res["selection_metric_used"] != cfg.task.selection.metric: - log.warning(f"[outer {i} inner {j}] epoch chosen by {res['selection_metric_used']}") + log.warning( + f"[outer {i} inner {j}] epoch chosen by {res['selection_metric_used']}" + ) epochs.append(int(res["best_epoch"])) n_epochs = _reduce(epochs, str(ncfg.epoch_reduce), int(ncfg.min_refit_epochs)) - refit = _run(cfg, train_ids=dev, val_ids=[], test_ids=test, out=fold_dir / "refit", max_epochs=n_epochs) + refit = _run( + cfg, + train_ids=dev, + val_ids=[], + test_ids=test, + out=fold_dir / "refit", + max_epochs=n_epochs, + ) preds = pd.read_csv(fold_dir / "refit" / "test_predictions.csv") preds.insert(0, "outer", i) preds.to_csv(fold_dir / "oof_predictions.csv", index=False) @@ -140,7 +175,9 @@ def _units_of(ids: Sequence[str]) -> list[str]: result_path.write_text(json.dumps(fold, indent=1)) results.append(fold) oof.append(preds) - log.info(f"[outer {i}] {n_epochs} epochs, held-out {metric} {fold['metrics'].get(metric, float('nan')):.4f}") + log.info( + f"[outer {i}] {n_epochs} epochs, held-out {metric} {fold['metrics'].get(metric, float('nan')):.4f}" + ) oof_frame = pd.concat(oof, ignore_index=True) oof_frame.to_csv(out_dir / "oof_predictions.csv", index=False) @@ -162,7 +199,9 @@ def _units_of(ids: Sequence[str]) -> list[str]: patients.to_csv(out_dir / "oof_patient_predictions.csv", index=False) summary["patient"] = frame_metrics(patients, kind) (out_dir / "nested_cv_summary.json").write_text(json.dumps(summary, indent=1, default=float)) - log.info(f"{metric}: {summary['mean']:.4f} ± {summary['sd']:.4f} over {len(finite)} outer folds") + log.info( + f"{metric}: {summary['mean']:.4f} ± {summary['sd']:.4f} over {len(finite)} outer folds" + ) return summary diff --git a/src/grass_mil/niches/__init__.py b/src/grass_mil/niches/__init__.py new file mode 100644 index 0000000..8ee3b56 --- /dev/null +++ b/src/grass_mil/niches/__init__.py @@ -0,0 +1,18 @@ +"""Niche discovery and description on the per-cell outputs of a trained model.""" + +from .analysis import NicheAnalysis, analyse_niches, write_niche_outputs +from .attribution import cell_margins +from .cells import Cells, load_cells +from .discovery import ClusteringConfig, NicheModel, fit_niches + +__all__ = [ + "Cells", + "ClusteringConfig", + "NicheAnalysis", + "NicheModel", + "analyse_niches", + "cell_margins", + "fit_niches", + "load_cells", + "write_niche_outputs", +] diff --git a/src/grass_mil/niches/analysis.py b/src/grass_mil/niches/analysis.py new file mode 100644 index 0000000..08cbf91 --- /dev/null +++ b/src/grass_mil/niches/analysis.py @@ -0,0 +1,301 @@ +"""Niche analysis of one model's cells: discovery, attribution, descriptors, figures.""" + +from __future__ import annotations + +import json +import logging +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import joblib +import numpy as np +import pandas as pd + +from grass_mil.inference.cell_cache import CellCache +from grass_mil.spatial.graph_stats import colocalisation + +from . import figures +from .attribution import cell_margins +from .cells import Cells, load_cells +from .descriptors import DescriptorConfig, niche_descriptors +from .discovery import BACKGROUND, ClusteringConfig, NicheModel, fit_niches, niche_name +from .survival import SurvivalConfig, niche_survival +from .tables import niche_table, unit_niche_shares + +__all__ = ["NicheAnalysis", "analyse_niches", "write_niche_outputs"] + +log = logging.getLogger(__name__) + + +@dataclass +class NicheAnalysis: + cells: Cells + model: NicheModel + niches: np.ndarray # niche id per cell (0 = Background) + margin: np.ndarray + table: pd.DataFrame + unit_shares: pd.DataFrame + descriptors: dict[str, pd.DataFrame] + colocalisation: pd.DataFrame + positive: dict[str, bool] + survival: tuple[pd.DataFrame, pd.DataFrame | None] | None = None + settings: dict[str, Any] = field(default_factory=dict) + + @property + def niche_names(self) -> np.ndarray: + return np.array([niche_name(n) for n in self.niches], dtype=object) + + +def analyse_niches( + cache: CellCache, + units: list[str], + *, + clustering: ClusteringConfig, + descriptors: DescriptorConfig, + survival: SurvivalConfig, + positive_class: int = 1, + present_threshold: float = 0.02, + enrichment_threshold: float = 1.5, + n_boot: int = 2000, +) -> NicheAnalysis: + cells = load_cells(cache, units) + missing = [u for u in units if u not in cells.targets] + if missing: + raise ValueError( + f"{len(missing)} unit(s) have no label in the cache (first: {missing[:5]})." + ) + margin = cell_margins( + cells.instance_logits, + cells.attention_logits, + cells.head_bias, + cells.unit_of_cell, + positive_class=positive_class, + ) + model, niches = fit_niches(cells.space[clustering.space], margin, clustering) + positive = cells.positive(positive_class) + table = niche_table( + niches, + cells, + margin, + positive=positive, + present_threshold=present_threshold, + enrichment_threshold=enrichment_threshold, + n_boot=n_boot, + seed=clustering.seed, + ) + shares = unit_niche_shares(niches, cells) + named = [n for n in table["niche"] if n != "Background"] + coloc = colocalisation(np.where(niches == BACKGROUND, -1, niches - 1), cells.edges, len(named)) + result = NicheAnalysis( + cells=cells, + model=model, + niches=niches, + margin=margin, + table=table, + unit_shares=shares, + descriptors=niche_descriptors(niches, cells, descriptors), + colocalisation=pd.DataFrame(coloc, index=named, columns=named), + positive=positive, + settings={ + "units": units, + "n_cells": cells.n, + "positive_class": positive_class, + "present_threshold": present_threshold, + "enrichment_threshold": enrichment_threshold, + "clustering": model.settings(), + }, + ) + if survival.file: + outcomes = pd.read_csv(survival.file, dtype={survival.id_column: str}).set_index( + survival.id_column + ) + outcomes = outcomes[[survival.time_column, survival.event_column]].set_axis( + ["time", "event"], axis=1 + ) + patients = shares.groupby([cells.patients.get(u, u) for u in shares.index]).mean() + result.survival = niche_survival(patients, outcomes, survival) + return result + + +def write_niche_outputs( + result: NicheAnalysis, + out_dir: str | Path, + *, + title: str, + positive_name: str, + render_units: int, + figures_enabled: bool = True, +) -> Path: + out = Path(out_dir) + out.mkdir(parents=True, exist_ok=True) + cells = result.cells + result.table.to_csv(out / "niches.csv", index=False) + result.unit_shares.to_csv(out / "unit_niche_shares.csv") + pd.DataFrame( + { + "unit_id": cells.unit_of_cell, + "cell_id": cells.cell_ids, + "x": cells.pos[:, 0], + "y": cells.pos[:, 1], + "cell_type": [cells.cell_type_names[c] if c >= 0 else "" for c in cells.cell_type], + "niche": result.niche_names, + "margin": result.margin, + } + ).to_csv(out / "cells.csv", index=False) + result.colocalisation.to_csv(out / "niche_colocalisation.csv") + desc_dir = out / "descriptors" + desc_dir.mkdir(exist_ok=True) + for name, frame in result.descriptors.items(): + frame.to_csv(desc_dir / f"{name}.csv", index=False) + if result.survival is not None: + result.survival[0].to_csv(out / "survival_cox.csv", index=False) + if result.survival[1] is not None: + result.survival[1].to_csv(out / "survival_groups.csv") + joblib.dump(result.model, out / "niche_model.joblib") + (out / "settings.json").write_text(json.dumps(result.settings, indent=1, default=str)) + (out / "report.md").write_text(_report(result, title=title, positive_name=positive_name)) + if figures_enabled: + _figures(result, out, positive_name=positive_name, render_units=render_units) + return out + + +def _figures(result: NicheAnalysis, out: Path, *, positive_name: str, render_units: int) -> None: + from sklearn.decomposition import PCA + + cells = result.cells + names = cells.cell_type_names + fig_dir = out / "figures" + fig_dir.mkdir(exist_ok=True) + figures.niche_panel( + result.table, names, fig_dir / "niche_panel.png", positive_name=positive_name + ) + sample = result.model.sample + x = result.model.coordinates(cells.space[result.model.cfg.space][sample]) + pca = PCA(n_components=2, random_state=0).fit(x) + ev = pca.explained_variance_ratio_ + is_pos = np.array([result.positive.get(u, False) for u in cells.unit_of_cell[sample]]) + figures.embedding_panel( + pca.transform(x), + result.niche_names[sample], + result.margin[sample], + is_pos, + fig_dir / "embedding.png", + positive_name=positive_name, + axis_labels=(f"PC 1 ({100 * ev[0]:.0f}%)", f"PC 2 ({100 * ev[1]:.0f}%)"), + ) + figures.unit_composition( + result.unit_shares, + result.positive, + fig_dir / "unit_composition.png", + positive_name=positive_name, + ) + if len(result.colocalisation): + figures.colocalisation_heatmap( + result.colocalisation.to_numpy(), + list(result.colocalisation.index), + fig_dir / "niche_colocalisation.png", + ) + desc = result.descriptors + if "moran" in desc: + figures.moran_heatmap(desc["moran"], fig_dir / "moran.png") + if "enrichment" in desc: + figures.pair_heatmaps( + desc["enrichment"], + "z", + names, + fig_dir / "enrichment.png", + label="neighbourhood enrichment z", + ) + if "cross_l" in desc: + figures.pair_heatmaps( + desc["cross_l"], "centred_l", names, fig_dir / "cross_l.png", label="L(r*) - r*" + ) + if "contacts" in desc: + figures.pair_heatmaps( + desc["contacts"], + "log2_enrichment", + names, + fig_dir / "contacts.png", + label="log2 contact enrichment", + ) + if result.survival is not None and result.survival[1] is not None: + figures.survival_curves( + result.survival[1], fig_dir / "survival.png", score_label="niche score tertiles" + ) + if render_units > 0: + _unit_maps(result, fig_dir / "units", render_units) + + +def _unit_maps(result: NicheAnalysis, out: Path, n: int) -> None: + """Maps of the units the model scores most and least positive.""" + cells = result.cells + out.mkdir(exist_ok=True) + unit_margin = pd.Series(result.margin).groupby(cells.unit_of_cell).mean() + order = unit_margin.sort_values().index.tolist() + chosen = list(dict.fromkeys(order[-((n + 1) // 2) :][::-1] + order[: n // 2])) + for unit in chosen: + m = cells.unit_of_cell == unit + figures.unit_maps( + cells.pos[m], + cells.cell_type[m], + cells.cell_type_names, + result.margin[m], + result.niche_names[m], + out / f"{''.join(c if c.isalnum() or c in '._-' else '_' for c in unit)}.png", + title=unit, + ) + + +def _fmt(value: float, spec: str = "+.2f") -> str: + return "" if value is None or not np.isfinite(value) else format(value, spec) + + +def _report(result: NicheAnalysis, *, title: str, positive_name: str) -> str: + s = result.settings + c = s["clustering"] + t = result.table + base = t.attrs.get("base_rate", float("nan")) + lines = [ + f"# {title}", + "", + f"{len(s['units'])} units, {s['n_cells']} cells. Clustering: {c['method']} in the {c['space']} space " + f"(PCA {c['pca_components'] or 'none'}, cosine {c['cosine']}, fitted on {c['n_fitted']} cells, " + f"{c['assign_k']}-NN assignment): {c['n_niches']} niches. Fraction of cells from {positive_name} units: {base:.3f}.", + "", + f"| Niche | Cells | Units (>{100 * s['present_threshold']:.0f}%) | Composition | Enriched (>{s['enrichment_threshold']}x) | " + f"Association | Margin [95% CI] | Sign consistency |", + "|---|---|---|---|---|---|---|---|", + ] + for _, r in t.iterrows(): + lines.append( + f"| {r['niche']} | {100 * r['share']:.1f}% | {r['n_units']} | {r.get('composition', '')} | {r.get('enriched', '')} | " + f"{r['association']:.2f} | {_fmt(r['margin'])} [{_fmt(r['margin_lo'])}, {_fmt(r['margin_hi'])}] | {_fmt(r['sign_consistency'], '.2f')} |" + ) + moran = result.descriptors.get("moran") + if moran is not None and len(moran): + top = moran[moran["q"] < 0.05].sort_values("morans_i", ascending=False).head(10) + if len(top): + lines += ["", "## Clustered cell types within niches (Moran's I, q < 0.05)", ""] + lines += [ + f"- {r.niche}: {r.cell_type} I = {r.morans_i:.2f} ± {r.sem:.2f} ({r.n_units} units)" + for r in top.itertuples() + ] + if result.survival is not None: + cox = result.survival[0].sort_values("p").head(10) + lines += [ + "", + "## Niche share and survival (Cox, per SD of log share)", + "", + "| Niche | HR [95% CI] | p | q |", + "|---|---|---|---|", + ] + lines += [ + f"| {r.niche} | {r.hazard_ratio:.2f} [{r.hr_lo:.2f}, {r.hr_hi:.2f}] | {r.p:.3g} | {r.q:.2g} |" + for r in cox.itertuples() + ] + if result.survival[1] is not None: + lines.append( + f"\nScore tertiles: log-rank p = {result.survival[1].attrs['logrank_p']:.3g}" + ) + return "\n".join(lines) + "\n" diff --git a/src/grass_mil/niches/attribution.py b/src/grass_mil/niches/attribution.py new file mode 100644 index 0000000..8ad33ab --- /dev/null +++ b/src/grass_mil/niches/attribution.py @@ -0,0 +1,50 @@ +"""Per-cell attention-weighted margins. + +Within a unit, the attention of class ``c`` is a softmax over the unit's cells, +so a cell's contribution ``A[i, c] (l[i, c] - b_c)`` sums over the unit to its +bag logit minus the head bias ``b_c``. The margin of a cell is its +contribution to the positive class minus its contribution to the other classes +(their mean beyond two classes), multiplied by the unit's cell count ``N_u`` so +that units of different sizes are comparable (a uniformly attended cell gets +weight 1). For regression and survival models it is the contribution to the +first output. +""" + +from __future__ import annotations + +import numpy as np + +__all__ = ["cell_margins", "unit_attention"] + + +def unit_attention(scores: np.ndarray, unit_of_cell: np.ndarray) -> np.ndarray: + """Attention weights: a softmax over each unit's cells, per class.""" + weights = np.zeros_like(scores, dtype=np.float64) + for unit in np.unique(unit_of_cell): + m = unit_of_cell == unit + s = scores[m] - scores[m].max(axis=0, keepdims=True) + e = np.exp(s) + weights[m] = e / e.sum(axis=0, keepdims=True) + return weights + + +def cell_margins( + logits: np.ndarray, + scores: np.ndarray, + bias: np.ndarray, + unit_of_cell: np.ndarray, + *, + positive_class: int = 1, +) -> np.ndarray: + """Size-scaled, attention-weighted margin of every cell toward the positive output.""" + logits = np.asarray(logits, dtype=np.float64) + weights = unit_attention(np.asarray(scores, dtype=np.float64), unit_of_cell) + contrib = weights * (logits - np.asarray(bias, dtype=np.float64)[None, :]) + n_classes = contrib.shape[1] + if n_classes == 1: + margin = contrib[:, 0] + else: + rest = (contrib.sum(axis=1) - contrib[:, positive_class]) / (n_classes - 1) + margin = contrib[:, positive_class] - rest + _, inverse, counts = np.unique(unit_of_cell, return_inverse=True, return_counts=True) + return (margin * counts[inverse]).astype(np.float32) diff --git a/src/grass_mil/niches/cells.py b/src/grass_mil/niches/cells.py new file mode 100644 index 0000000..4ddcb72 --- /dev/null +++ b/src/grass_mil/niches/cells.py @@ -0,0 +1,127 @@ +"""The cells of selected units of a cell cache, concatenated, with their labels.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Sequence + +import numpy as np + +from grass_mil.data.labels import class_index +from grass_mil.inference.cell_cache import CellCache + +__all__ = ["Cells", "load_cells", "resolve_units"] + +SPACES = ("attention", "embedding") + + +@dataclass +class Cells: + units: list[str] + unit_of_cell: np.ndarray # (N,) unit id of every cell + space: dict[str, np.ndarray] # "attention" / "embedding" -> (N, d) + instance_logits: np.ndarray # (N, C) + attention_logits: np.ndarray # (N, C) + head_bias: np.ndarray # (C,) + cell_type: np.ndarray # (N,) int (-1 without cell types) + cell_type_names: list[str] + pos: np.ndarray # (N, 2) + cell_ids: np.ndarray # (N,) + edges: np.ndarray # (2, M) undirected, global cell indices + target_type: str + targets: dict[str, list[float]] = field(default_factory=dict) # unit -> targets + patients: dict[str, str] = field(default_factory=dict) # unit -> patient + + @property + def n(self) -> int: + return int(self.unit_of_cell.size) + + def unit_class(self, unit: str) -> int: + return class_index(self.targets[unit][0]) + + def unit_value(self, unit: str) -> float: + """The value that defines a positive unit: the class, or the event (survival).""" + t = self.targets[unit] + return float(t[1]) if self.target_type == "survival" else float(t[0]) + + def positive(self, positive_class: int) -> dict[str, bool]: + """Which units count as positive (class ``positive_class``; event for survival).""" + if self.target_type == "classification": + return {u: self.unit_class(u) == positive_class for u in self.targets} + if self.target_type == "survival": + return {u: bool(self.targets[u][1] > 0) for u in self.targets} + median = float(np.median([t[0] for t in self.targets.values()])) + return {u: t[0] > median for u, t in self.targets.items()} + + +def resolve_units(selection, cache: CellCache) -> list[str]: + """``held_out`` (the units the model did not train on), ``all``, or explicit ids.""" + from grass_mil.data.splits import read_ids + + if selection in (None, "held_out"): + if not cache.held_out: + raise ValueError( + f"{cache.model_dir} records no held-out units; pass units=all or a list of unit ids." + ) + return sorted(cache.held_out) + if selection == "all": + return sorted(cache.units) + return sorted(read_ids(selection)) + + +def load_cells(cache: CellCache, units: Sequence[str]) -> Cells: + parts: dict[str, list] = { + k: [] + for k in ( + "attention", + "embedding", + "logits", + "scores", + "ct", + "pos", + "ids", + "edges", + "unit", + ) + } + offset = 0 + targets, patients = {}, {} + for unit in units: + info = cache.unit_info(unit) + d = cache.load(unit) + n = int(d["pos"].shape[0]) + parts["attention"].append(d["attention_space"].astype(np.float32)) + parts["embedding"].append(d["embedding"].astype(np.float32)) + parts["logits"].append(d["instance_logits"]) + parts["scores"].append(d["attention_logits"]) + parts["ct"].append(d["cell_type"] if "cell_type" in d else np.full(n, -1, dtype=np.int32)) + parts["pos"].append(d["pos"]) + parts["ids"].append(d["cell_ids"]) + e = d["edge_index"] + e = e[:, e[0] < e[1]] + parts["edges"].append(e + offset) + parts["unit"].append(np.full(n, unit, dtype=object)) + if info.get("targets") is not None: + targets[unit] = list(info["targets"]) + if info.get("patient_id") is not None: + patients[unit] = str(info["patient_id"]) + offset += n + return Cells( + units=list(units), + unit_of_cell=np.concatenate(parts["unit"]), + space={ + "attention": np.concatenate(parts["attention"]), + "embedding": np.concatenate(parts["embedding"]), + }, + instance_logits=np.concatenate(parts["logits"]), + attention_logits=np.concatenate(parts["scores"]), + head_bias=cache.head_bias, + cell_type=np.concatenate(parts["ct"]).astype(int), + cell_type_names=cache.cell_types, + pos=np.concatenate(parts["pos"]), + cell_ids=np.concatenate(parts["ids"]), + edges=np.concatenate(parts["edges"], axis=1), + target_type=str(cache.index.get("target_type", "classification")), + targets=targets, + patients=patients, + ) diff --git a/src/grass_mil/niches/cli.py b/src/grass_mil/niches/cli.py new file mode 100644 index 0000000..2b8a692 --- /dev/null +++ b/src/grass_mil/niches/cli.py @@ -0,0 +1,62 @@ +"""``grass-mil-niches``: niche analysis of one model's cell cache.""" + +from __future__ import annotations + +import logging +from pathlib import Path + +import hydra +from omegaconf import DictConfig, OmegaConf + +from grass_mil.inference.cell_cache import CellCache + +from .analysis import analyse_niches, write_niche_outputs +from .cells import resolve_units +from .descriptors import DescriptorConfig +from .discovery import ClusteringConfig +from .survival import SurvivalConfig + +__all__ = ["run_niches", "main"] + +log = logging.getLogger(__name__) + + +def run_niches(cfg: DictConfig) -> Path: + cache = CellCache.open(str(cfg.cells_dir)) + + def plain(key: str): + value = cfg.get(key) + return OmegaConf.to_container(value, resolve=True) if OmegaConf.is_config(value) else value + + result = analyse_niches( + cache, + resolve_units(plain("units"), cache), + clustering=ClusteringConfig.from_dict(plain("clustering")), + descriptors=DescriptorConfig.from_dict(plain("descriptors")), + survival=SurvivalConfig.from_dict(plain("survival")), + positive_class=int(cfg.positive_class), + present_threshold=float(cfg.present_threshold), + enrichment_threshold=float(cfg.enrichment_threshold), + n_boot=int(cfg.n_boot), + ) + out = write_niche_outputs( + result, + str(cfg.output_dir), + title=str(cfg.title or f"Niches of {cache.name}"), + positive_name=str(cfg.positive_name), + render_units=int(cfg.render_units), + figures_enabled=bool(cfg.figures), + ) + log.info( + f"{result.model.n_niches} niches in {result.cells.n} cells; report: {out / 'report.md'}" + ) + return out + + +@hydra.main(version_base="1.3", config_path="../configs", config_name="niches.yaml") +def main(cfg: DictConfig) -> None: + run_niches(cfg) + + +if __name__ == "__main__": + main() diff --git a/src/grass_mil/niches/descriptors.py b/src/grass_mil/niches/descriptors.py new file mode 100644 index 0000000..73ef816 --- /dev/null +++ b/src/grass_mil/niches/descriptors.py @@ -0,0 +1,254 @@ +"""Classical spatial descriptors of every niche. + +Each statistic is computed on the cells of one niche within one unit (the +niche's subgraph: its cells and the cell-graph edges between them) and then +averaged over the units in which the niche has at least ``min_cells`` cells: + +* ``moran``: Moran's I of every cell-type indicator (are cells of a type + clustered within the niche?), mean and s.e.m. over units, with a + label-permutation p-value of the mean and BH correction over niche x type; +* ``enrichment``: neighbourhood-enrichment z-score of every cell-type pair; +* ``cross_l``: Ripley's centred cross-L of every cell-type pair on a shared + radius grid; ``r_star`` is the radius maximising the mean absolute value + over pairs, and the table reports ``L(r*) - r*``; +* ``contacts``: ``log2`` of a pair's share of the contacts inside the niche + over its share of all contacts of the analysed cells (all units pooled), + masked below ``min_contacts`` contacts inside the niche. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import Any, Mapping + +import numpy as np +import pandas as pd + +from grass_mil.spatial.graph_stats import contact_counts, cross_l, enrichment_zscores, morans_i +from grass_mil.spatial.stats import bh + +from .cells import Cells +from .discovery import BACKGROUND, niche_name + +__all__ = ["DescriptorConfig", "niche_descriptors"] + +ANALYSES = ("moran", "enrichment", "cross_l", "contacts") + + +@dataclass(frozen=True) +class DescriptorConfig: + analyses: tuple[str, ...] = ANALYSES + min_cells: int = 20 + n_perms: int = 100 + n_radii: int = 20 + max_radius: float | None = None # null = a quarter of the median unit diagonal + min_contacts: int = 100 + max_edge_length: float | None = None # drop longer cell-graph edges first + seed: int = 0 + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any] | None) -> DescriptorConfig: + cfg = dict(cfg or {}) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown descriptors keys: {sorted(unknown)}") + cfg["analyses"] = tuple(cfg.get("analyses", ANALYSES)) + bad = set(cfg["analyses"]) - set(ANALYSES) + if bad: + raise ValueError( + f"Unknown descriptor analyses {sorted(bad)}; expected {list(ANALYSES)}." + ) + return cls(**cfg) + + +def _subgraphs(cells: Cells, niches: np.ndarray, edges: np.ndarray, min_cells: int): + """``(niche, unit, cell indices, local edges)`` of every niche subgraph with enough cells.""" + _, unit_index = np.unique(cells.unit_of_cell, return_inverse=True) + key = niches.astype(np.int64) * (int(unit_index.max()) + 1) + unit_index + order = np.argsort(key, kind="stable") + starts = np.flatnonzero(np.r_[True, np.diff(key[order]) != 0]) + # Edges inside one subgraph (same niche; units never share edges), grouped by key. + inside = edges[:, key[edges[0]] == key[edges[1]]] + edge_key = key[inside[0]] + edge_order = np.argsort(edge_key, kind="stable") + inside, edge_key = inside[:, edge_order], edge_key[edge_order] + local = np.empty(cells.n, dtype=np.int64) + for idx in np.split(order, starts[1:]): + if idx.size < min_cells or niches[idx[0]] == BACKGROUND: + continue + k = key[idx[0]] + lo, hi = np.searchsorted(edge_key, [k, k + 1]) + local[idx] = np.arange(idx.size) + yield int(niches[idx[0]]), str(cells.unit_of_cell[idx[0]]), idx, local[inside[:, lo:hi]] + + +def niche_descriptors( + niches: np.ndarray, cells: Cells, cfg: DescriptorConfig +) -> dict[str, pd.DataFrame]: + if (cells.cell_type < 0).any(): + return {} + names = cells.cell_type_names + k = len(names) + edges = cells.edges + if cfg.max_edge_length is not None: + length = np.linalg.norm(cells.pos[edges[0]] - cells.pos[edges[1]], axis=1) + edges = edges[:, length <= cfg.max_edge_length] + rng = np.random.default_rng(cfg.seed) + radii = _radii(cells, cfg) + pairs = [(a, b) for a in range(k) for b in range(a, k)] + collected: dict[str, list] = {a: [] for a in cfg.analyses} + for niche, unit, idx, local_edges in _subgraphs(cells, niches, edges, cfg.min_cells): + types = cells.cell_type[idx] + if "moran" in cfg.analyses: + onehot = np.eye(k)[types] + observed = morans_i(onehot, local_edges) + null = np.stack( + [ + morans_i(onehot[rng.permutation(idx.size)], local_edges) + for _ in range(cfg.n_perms) + ] + ) + collected["moran"].append((niche, observed, null)) + if "enrichment" in cfg.analyses: + z = enrichment_zscores(types, local_edges, k, n_perms=cfg.n_perms, rng=rng) + collected["enrichment"].append((niche, z)) + if "cross_l" in cfg.analyses: + coords = cells.pos[idx] + curves = np.full((len(pairs), radii.size), np.nan) + counts = np.bincount(types, minlength=k) + for p, (a, b) in enumerate(pairs): + if counts[a] >= 2 and counts[b] >= 2: + curves[p] = cross_l(coords, types == a, types == b, radii) + collected["cross_l"].append((niche, curves)) + out: dict[str, pd.DataFrame] = {} + if "moran" in cfg.analyses: + out["moran"] = _moran_table(collected["moran"], names) + if "enrichment" in cfg.analyses: + out["enrichment"] = _pair_table(collected["enrichment"], names, "z") + if "cross_l" in cfg.analyses: + out["cross_l"] = _cross_l_table(collected["cross_l"], names, pairs, radii) + if "contacts" in cfg.analyses: + out["contacts"] = _contact_table(niches, cells.cell_type, edges, names, cfg.min_contacts) + return out + + +def _radii(cells: Cells, cfg: DescriptorConfig) -> np.ndarray: + if cfg.max_radius is not None: + top = float(cfg.max_radius) + else: + diagonals = [ + float(np.linalg.norm(np.ptp(cells.pos[cells.unit_of_cell == u], axis=0))) + for u in cells.units + ] + top = 0.25 * float(np.median(diagonals)) + return np.linspace(top / cfg.n_radii, top, cfg.n_radii) + + +def _moran_table(records, names) -> pd.DataFrame: + rows = [] + for niche in sorted({r[0] for r in records}): + obs = np.stack([r[1] for r in records if r[0] == niche]) # (units, k) + null = np.stack([r[2] for r in records if r[0] == niche]) # (units, perms, k) + for j, name in enumerate(names): + values = obs[:, j] + finite = np.isfinite(values) + if not finite.any(): + continue + mean = float(values[finite].mean()) + null_mean = np.nanmean(null[finite][:, :, j], axis=0) + centre = float(np.nanmean(null_mean)) + extreme = np.abs(null_mean - centre) >= abs(mean - centre) + rows.append( + { + "niche": niche_name(niche), + "cell_type": name, + "morans_i": mean, + "sem": float(values[finite].std(ddof=1) / np.sqrt(finite.sum())) + if finite.sum() > 1 + else float("nan"), + "n_units": int(finite.sum()), + "p": float((1 + extreme.sum()) / (1 + null_mean.size)), + } + ) + table = pd.DataFrame(rows, columns=["niche", "cell_type", "morans_i", "sem", "n_units", "p"]) + table["q"] = bh(table["p"].to_numpy()) + return table + + +def _pair_table(records, names, value: str) -> pd.DataFrame: + rows = [] + for niche in sorted({r[0] for r in records}): + stack = np.stack([r[1] for r in records if r[0] == niche]) + with np.errstate(all="ignore"): + mean = np.nanmean(stack, axis=0) + n = np.isfinite(stack).sum(0) + for a in range(len(names)): + for b in range(a, len(names)): + if n[a, b]: + rows.append( + { + "niche": niche_name(niche), + "type_a": names[a], + "type_b": names[b], + value: float(mean[a, b]), + "n_units": int(n[a, b]), + } + ) + return pd.DataFrame(rows, columns=["niche", "type_a", "type_b", value, "n_units"]) + + +def _cross_l_table(records, names, pairs, radii) -> pd.DataFrame: + rows = [] + for niche in sorted({r[0] for r in records}): + stack = np.stack([r[1] for r in records if r[0] == niche]) # (units, pairs, radii) + with np.errstate(all="ignore"): + mean = np.nanmean(stack, axis=0) + strength = np.nanmean(np.abs(mean), axis=0) + if not np.isfinite(strength).any(): + continue + r_idx = int(np.nanargmax(strength)) + for p, (a, b) in enumerate(pairs): + if np.isfinite(mean[p, r_idx]): + rows.append( + { + "niche": niche_name(niche), + "type_a": names[a], + "type_b": names[b], + "r_star": float(radii[r_idx]), + "centred_l": float(mean[p, r_idx]), + "n_units": int(np.isfinite(stack[:, p, r_idx]).sum()), + } + ) + return pd.DataFrame( + rows, columns=["niche", "type_a", "type_b", "r_star", "centred_l", "n_units"] + ) + + +def _contact_table(niches, types, edges, names, min_contacts: int) -> pd.DataFrame: + k = len(names) + tissue = contact_counts(types, edges, k) + tissue_frac = tissue / max(tissue.sum(), 1.0) + rows = [] + same = niches[edges[0]] == niches[edges[1]] + for niche in sorted(set(niches.tolist()) - {BACKGROUND}): + inside = edges[:, same & (niches[edges[0]] == niche)] + counts = contact_counts(types, inside, k) + frac = counts / max(counts.sum(), 1.0) + for a in range(k): + for b in range(a, k): + n_contacts = int(counts[a, b] if a != b else counts[a, b] / 2) + enough = n_contacts >= min_contacts and tissue_frac[a, b] > 0 + rows.append( + { + "niche": niche_name(niche), + "type_a": names[a], + "type_b": names[b], + "n_contacts": n_contacts, + "log2_enrichment": float(np.log2(frac[a, b] / tissue_frac[a, b])) + if enough + else float("nan"), + } + ) + return pd.DataFrame( + rows, columns=["niche", "type_a", "type_b", "n_contacts", "log2_enrichment"] + ) diff --git a/src/grass_mil/niches/discovery.py b/src/grass_mil/niches/discovery.py new file mode 100644 index 0000000..462ce3e --- /dev/null +++ b/src/grass_mil/niches/discovery.py @@ -0,0 +1,147 @@ +"""Clustering cells into niches. + +The clustering is fitted on a random sample of at most ``sample_size`` cells: +an optional PCA, L2 normalisation (``cosine``, so euclidean distances follow +the cosine geometry), then HDBSCAN (or k-means). The remaining cells are +assigned by a k-nearest-neighbour vote on the same coordinates. Cells HDBSCAN +leaves unassigned form the ``Background``. + +Niches are numbered by their attention-weighted margin, ``Niche 1`` being the +most positive; ``0`` is the Background. +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, fields +from typing import Any, Mapping + +import numpy as np + +__all__ = ["ClusteringConfig", "NicheModel", "fit_niches", "niche_name", "BACKGROUND"] + +BACKGROUND = 0 + + +def niche_name(niche_id: int) -> str: + return "Background" if int(niche_id) == BACKGROUND else f"Niche {int(niche_id)}" + + +@dataclass(frozen=True) +class ClusteringConfig: + """Attributes: + space: ``attention`` (the attention head's pre-tanh space) or ``embedding``. + pca_components: PCA before clustering (``None``: none). + cosine: L2-normalise before clustering. + method: ``hdbscan`` or ``kmeans``. + min_cluster_size, min_samples, selection: HDBSCAN settings. + n_clusters: k-means clusters. + sample_size: Cells used to fit (``None``: all). + assign_k: Neighbours voting for the cells outside the sample. + seed: Sampling and k-means seed. + """ + + space: str = "attention" + pca_components: int | None = None + cosine: bool = True + method: str = "hdbscan" + min_cluster_size: int = 300 + min_samples: int = 1 + selection: str = "eom" + n_clusters: int = 10 + sample_size: int | None = 100_000 + assign_k: int = 15 + seed: int = 0 + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any] | None) -> ClusteringConfig: + cfg = dict(cfg or {}) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown clustering keys: {sorted(unknown)}") + out = cls(**cfg) + if out.space not in ("attention", "embedding"): + raise ValueError("clustering.space must be 'attention' or 'embedding'.") + if out.method not in ("hdbscan", "kmeans"): + raise ValueError("clustering.method must be 'hdbscan' or 'kmeans'.") + return out + + +class NicheModel: + """A fitted niche clustering that labels any cells of the same space.""" + + def __init__( + self, cfg: ClusteringConfig, pca, assigner, mapping: dict[int, int], sample: np.ndarray + ) -> None: + self.cfg = cfg + self.pca = pca + self.assigner = assigner + self.mapping = mapping # raw cluster id -> niche id + self.sample = sample + + @property + def n_niches(self) -> int: + return len([v for v in self.mapping.values() if v != BACKGROUND]) + + def coordinates(self, x: np.ndarray) -> np.ndarray: + from sklearn.preprocessing import normalize + + z = np.asarray(x, dtype=np.float32) + if self.pca is not None: + z = self.pca.transform(z) + return normalize(z) if self.cfg.cosine else z + + def predict(self, x: np.ndarray) -> np.ndarray: + """Niche ids of the cells (``0`` = Background).""" + raw = self.assigner.predict(self.coordinates(x)) + return np.vectorize(self.mapping.get)(raw).astype(int) + + def settings(self) -> dict[str, Any]: + return {**asdict(self.cfg), "n_niches": self.n_niches, "n_fitted": int(self.sample.size)} + + +def fit_niches( + x: np.ndarray, margin: np.ndarray, cfg: ClusteringConfig +) -> tuple[NicheModel, np.ndarray]: + """Fit the clustering on ``x`` and return the model and every cell's niche id. + + ``margin`` (per cell) orders the niches. + """ + from sklearn.cluster import HDBSCAN, KMeans + from sklearn.decomposition import PCA + from sklearn.neighbors import KNeighborsClassifier + from sklearn.preprocessing import normalize + + x = np.asarray(x, dtype=np.float32) + n = x.shape[0] + rng = np.random.default_rng(cfg.seed) + sample = ( + np.arange(n) + if cfg.sample_size is None or cfg.sample_size >= n + else np.sort(rng.choice(n, cfg.sample_size, replace=False)) + ) + pca = None + z = x[sample] + if cfg.pca_components: + pca = PCA( + n_components=min(cfg.pca_components, x.shape[1], sample.size), random_state=cfg.seed + ).fit(z) + z = pca.transform(z) + if cfg.cosine: + z = normalize(z) + if cfg.method == "hdbscan": + raw = HDBSCAN( + min_cluster_size=int(min(cfg.min_cluster_size, max(2, sample.size))), + min_samples=cfg.min_samples, + cluster_selection_method=cfg.selection, + ).fit_predict(z) + else: + raw = KMeans(n_clusters=cfg.n_clusters, random_state=cfg.seed, n_init=10).fit_predict(z) + assigner = KNeighborsClassifier(n_neighbors=max(1, min(cfg.assign_k, sample.size))).fit(z, raw) + model = NicheModel(cfg, pca, assigner, {}, sample) + labels = assigner.predict(model.coordinates(x)) + labels[sample] = raw + # Number the clusters by mean margin, most positive first; -1 is Background. + clusters = [c for c in np.unique(labels) if c >= 0] + order = sorted(clusters, key=lambda c: -float(np.mean(margin[labels == c]))) + model.mapping = {-1: BACKGROUND, **{int(c): rank + 1 for rank, c in enumerate(order)}} + return model, np.vectorize(model.mapping.get)(labels).astype(int) diff --git a/src/grass_mil/niches/figures.py b/src/grass_mil/niches/figures.py new file mode 100644 index 0000000..67ea62b --- /dev/null +++ b/src/grass_mil/niches/figures.py @@ -0,0 +1,283 @@ +"""Figures of a niche analysis (matplotlib, written as PNG).""" + +from __future__ import annotations + +from pathlib import Path +from typing import Sequence + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +__all__ = [ + "niche_colours", + "niche_panel", + "embedding_panel", + "unit_composition", + "unit_maps", + "pair_heatmaps", + "moran_heatmap", + "colocalisation_heatmap", + "survival_curves", +] + +BACKGROUND_COLOUR = "#d9d9d9" + + +def niche_colours(names: Sequence[str]) -> dict[str, tuple]: + cmap = plt.get_cmap("tab20") + out, i = {}, 0 + for name in names: + if name == "Background": + out[name] = BACKGROUND_COLOUR + else: + out[name] = cmap(i % 20) + i += 1 + return out + + +def _save(fig, path: Path) -> Path: + fig.tight_layout() + fig.savefig(path, dpi=150) + plt.close(fig) + return path + + +def niche_panel( + table: pd.DataFrame, cell_types: Sequence[str], path: Path, *, positive_name: str +) -> Path: + """Composition | association | margin with 95% CI | share, one row per niche.""" + t = table[table["niche"] != "Background"].reset_index(drop=True) + if t.empty: + return path + base = table.attrs.get("base_rate", np.nan) + frac = [f"frac_{c}" for c in cell_types if f"frac_{c}" in t.columns] + y = np.arange(len(t))[::-1] + widths = [max(2.0, 0.28 * len(frac)), 1.1, 1.3, 0.9] if frac else [1.1, 1.3, 0.9] + fig, axes = plt.subplots( + 1, + len(widths), + figsize=(sum(widths) + 2.5, 0.42 * len(t) + 1.8), + gridspec_kw={"width_ratios": widths}, + ) + axes = list(np.atleast_1d(axes)) + if frac: + ax = axes.pop(0) + im = ax.imshow(t[frac].to_numpy() * 100, aspect="auto", cmap="viridis", vmin=0) + ax.set_xticks(range(len(frac)), [c[5:] for c in frac], rotation=90, fontsize=7) + ax.set_yticks(range(len(t)), t["niche"], fontsize=8) + plt.colorbar(im, ax=ax, fraction=0.04, label="% of niche cells") + ax.set_title("composition", fontsize=9) + ax_assoc, ax_margin, ax_share = axes + ax_assoc.barh( + y, + t["association"], + color=["#b2182b" if v > base else "#2166ac" for v in t["association"]], + height=0.6, + ) + ax_assoc.axvline(base, color="k", ls=":", lw=1) + ax_assoc.set_xlim(0, 1) + ax_assoc.set_xlabel(f"fraction from {positive_name} units") + ax_margin.errorbar( + t["margin"], + y, + xerr=[t["margin"] - t["margin_lo"], t["margin_hi"] - t["margin"]], + fmt="o", + color="#1b9e77", + ms=4, + capsize=2, + ) + ax_margin.axvline(0, color="k", lw=0.8) + ax_margin.set_xlabel(f"attention-weighted margin (+ = {positive_name})") + ax_share.barh(y, 100 * t["share"], color="#7f7f7f", height=0.6) + ax_share.set_xlabel("% of cells") + for ax in (ax_assoc, ax_margin, ax_share): + ax.set_yticks(y, [] if frac else t["niche"]) + ax.spines[["top", "right"]].set_visible(False) + return _save(fig, path) + + +def embedding_panel( + coords: np.ndarray, + niche_names: np.ndarray, + margin: np.ndarray, + positive: np.ndarray, + path: Path, + *, + positive_name: str, + axis_labels=("PC 1", "PC 2"), +) -> Path: + """The clustering space (first two PCs): niches | margin | unit class.""" + fig, axes = plt.subplots(1, 3, figsize=(15, 4.6)) + colours = niche_colours(sorted(set(niche_names), key=lambda n: (n == "Background", n))) + for name, colour in colours.items(): + m = niche_names == name + axes[0].scatter(coords[m, 0], coords[m, 1], s=2, color=colour, lw=0, label=name) + axes[0].legend(markerscale=6, fontsize=6, ncol=2) + lim = float(np.nanpercentile(np.abs(margin), 99)) or 1.0 + sc = axes[1].scatter( + coords[:, 0], coords[:, 1], s=2, c=margin, cmap="RdBu_r", vmin=-lim, vmax=lim, lw=0 + ) + plt.colorbar(sc, ax=axes[1], label=f"margin (+ = {positive_name})") + hb = axes[2].hexbin( + coords[:, 0], + coords[:, 1], + C=positive.astype(float), + gridsize=40, + cmap="RdBu_r", + vmin=0, + vmax=1, + mincnt=3, + ) + plt.colorbar(hb, ax=axes[2], label=f"fraction from {positive_name} units") + for ax, title in zip(axes, ("niches", "attention-weighted margin", "empirical association")): + ax.set_title(title, fontsize=9) + ax.set_xlabel(axis_labels[0]) + ax.set_ylabel(axis_labels[1]) + ax.set_xticks([]) + ax.set_yticks([]) + return _save(fig, path) + + +def unit_composition( + shares: pd.DataFrame, positive: dict[str, bool], path: Path, *, positive_name: str +) -> Path: + """Stacked niche shares of every unit, negative units first.""" + order = sorted( + shares.index, + key=lambda u: ( + positive.get(u, False), + -shares.loc[u].drop("Background", errors="ignore").max(), + ), + ) + s = shares.loc[order] + colours = niche_colours(list(s.columns)) + fig, ax = plt.subplots(figsize=(max(6, 0.12 * len(s) + 3), 3.6)) + bottom = np.zeros(len(s)) + for name in s.columns: + ax.bar(range(len(s)), s[name], bottom=bottom, color=colours[name], width=0.9, label=name) + bottom += s[name].to_numpy() + n_neg = sum(not positive.get(u, False) for u in order) + ax.axvline(n_neg - 0.5, color="k", lw=1) + ax.set_xticks([]) + ax.set_xlabel(f"units (left: other, right: {positive_name})") + ax.set_ylabel("share of cells") + ax.set_ylim(0, 1) + ax.legend(fontsize=6, ncol=2, bbox_to_anchor=(1.01, 1), loc="upper left") + return _save(fig, path) + + +def unit_maps( + pos: np.ndarray, + cell_types: np.ndarray, + type_names: Sequence[str], + margin: np.ndarray, + niche_names: np.ndarray, + path: Path, + *, + title: str, +) -> Path: + """One unit in the tissue: cell types | cell margins | niches.""" + fig, axes = plt.subplots(1, 3, figsize=(15, 5)) + cmap = plt.get_cmap("tab20") + for j, name in enumerate(type_names): + m = cell_types == j + if m.any(): + axes[0].scatter(pos[m, 0], pos[m, 1], s=4, color=cmap(j % 20), lw=0, label=name) + axes[0].legend(markerscale=3, fontsize=6, ncol=2) + lim = float(np.nanpercentile(np.abs(margin), 99)) or 1.0 + sc = axes[1].scatter( + pos[:, 0], pos[:, 1], s=4, c=margin, cmap="RdBu_r", vmin=-lim, vmax=lim, lw=0 + ) + plt.colorbar(sc, ax=axes[1], label="margin") + colours = niche_colours(sorted(set(niche_names), key=lambda n: (n == "Background", n))) + for name, colour in colours.items(): + m = niche_names == name + axes[2].scatter(pos[m, 0], pos[m, 1], s=4, color=colour, lw=0, label=name) + axes[2].legend(markerscale=3, fontsize=6) + for ax, sub in zip(axes, ("cell types", "attention-weighted margin", "niches")): + ax.set_aspect("equal") + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_title(f"{title}: {sub}", fontsize=9) + return _save(fig, path) + + +def pair_heatmaps( + table: pd.DataFrame, value: str, cell_types: Sequence[str], path: Path, *, label: str +) -> Path: + """One symmetric cell-type x cell-type heatmap per niche.""" + niches = list(dict.fromkeys(table["niche"])) + if not niches: + return path + k = len(cell_types) + index = {n: i for i, n in enumerate(cell_types)} + cols = min(4, len(niches)) + rows = int(np.ceil(len(niches) / cols)) + fig, axes = plt.subplots(rows, cols, figsize=(3.2 * cols + 1, 3.0 * rows + 0.5), squeeze=False) + lim = float(np.nanpercentile(np.abs(table[value]), 98)) if table[value].notna().any() else 1.0 + for ax, niche in zip(axes.flat, niches): + m = np.full((k, k), np.nan) + for _, r in table[table["niche"] == niche].iterrows(): + a, b = index[r["type_a"]], index[r["type_b"]] + m[a, b] = m[b, a] = r[value] + im = ax.imshow(m, cmap="RdBu_r", vmin=-lim, vmax=lim) + ax.set_title(niche, fontsize=8) + ax.set_xticks(range(k), cell_types, rotation=90, fontsize=5) + ax.set_yticks(range(k), cell_types, fontsize=5) + for ax in list(axes.flat)[len(niches) :]: + ax.axis("off") + fig.colorbar(im, ax=axes, shrink=0.6, label=label) + fig.savefig(path, dpi=150, bbox_inches="tight") + plt.close(fig) + return path + + +def moran_heatmap(table: pd.DataFrame, path: Path) -> Path: + """Moran's I per niche (rows) and cell type (columns); dots mark q < 0.05.""" + if table.empty: + return path + m = table.pivot(index="niche", columns="cell_type", values="morans_i") + q = table.pivot(index="niche", columns="cell_type", values="q").reindex_like(m) + fig, ax = plt.subplots(figsize=(0.35 * m.shape[1] + 3, 0.4 * m.shape[0] + 1.5)) + lim = float(np.nanmax(np.abs(m.to_numpy()))) or 1.0 + im = ax.imshow(m.to_numpy(), cmap="RdBu_r", vmin=-lim, vmax=lim, aspect="auto") + yy, xx = np.nonzero((q.to_numpy() < 0.05)) + ax.scatter(xx, yy, s=6, color="k") + ax.set_xticks(range(m.shape[1]), m.columns, rotation=90, fontsize=7) + ax.set_yticks(range(m.shape[0]), m.index, fontsize=8) + plt.colorbar(im, ax=ax, label="Moran's I") + return _save(fig, path) + + +def colocalisation_heatmap(matrix: np.ndarray, names: Sequence[str], path: Path) -> Path: + fig, ax = plt.subplots(figsize=(0.45 * len(names) + 2.5, 0.45 * len(names) + 1.8)) + lim = float(np.nanmax(np.abs(matrix))) or 1.0 + im = ax.imshow(matrix, cmap="RdBu_r", vmin=-lim, vmax=lim) + ax.set_xticks(range(len(names)), names, rotation=90, fontsize=7) + ax.set_yticks(range(len(names)), names, fontsize=7) + plt.colorbar(im, ax=ax, label="log2 observed / expected contacts") + return _save(fig, path) + + +def survival_curves(groups: pd.DataFrame, path: Path, *, score_label: str) -> Path: + from .survival import kaplan_meier + + fig, ax = plt.subplots(figsize=(5, 4)) + for tertile, label, colour in ( + (0, "low", "#2166ac"), + (1, "intermediate", "#7f7f7f"), + (2, "high", "#b2182b"), + ): + g = groups[groups["tertile"] == tertile] + if len(g): + t, s = kaplan_meier(g["time"].to_numpy(float), g["event"].to_numpy(float)) + ax.step(t, s, where="post", color=colour, label=f"{label} (n={len(g)})") + ax.set_xlabel("time") + ax.set_ylabel("survival probability") + ax.set_title( + f"{score_label}; log-rank p = {groups.attrs.get('logrank_p', float('nan')):.3g}", + fontsize=8, + ) + ax.legend(fontsize=7) + return _save(fig, path) diff --git a/src/grass_mil/niches/replication/__init__.py b/src/grass_mil/niches/replication/__init__.py new file mode 100644 index 0000000..9e0bf54 --- /dev/null +++ b/src/grass_mil/niches/replication/__init__.py @@ -0,0 +1 @@ +"""Replication of niches across the models of a cross-validation (optional analysis).""" diff --git a/src/grass_mil/niches/replication/cli.py b/src/grass_mil/niches/replication/cli.py new file mode 100644 index 0000000..3d49834 --- /dev/null +++ b/src/grass_mil/niches/replication/cli.py @@ -0,0 +1,170 @@ +"""``grass-mil-metaclusters``: replication of niches across cross-validation models.""" + +from __future__ import annotations + +import json +import logging +from pathlib import Path + +import hydra +import joblib +import numpy as np +import pandas as pd +from omegaconf import DictConfig, OmegaConf + +from grass_mil.inference.cell_cache import CellCache + +from ..discovery import ClusteringConfig +from .stages import ( + cluster_folds, + held_out_shares, + match_niches, + share_tests, + shared_basis, + transfer_niches, +) + +__all__ = ["run_metaclusters", "main", "open_caches"] + +log = logging.getLogger(__name__) + +STAGES = ("cluster", "match", "transfer", "basis") + + +def open_caches(cells_dir: str | Path) -> dict[str, CellCache]: + """Every model cache under ``cells_dir`` that records held-out units.""" + caches = {} + for sub in sorted(Path(cells_dir).iterdir()): + if (sub / "index.json").exists(): + cache = CellCache.open(sub) + if cache.held_out: + caches[cache.name] = cache + if len(caches) < 2: + raise ValueError(f"{cells_dir} needs at least two model caches with held-out units.") + return caches + + +def run_metaclusters(cfg: DictConfig) -> dict: + stage = str(cfg.stage) + if stage not in ("all", *STAGES): + raise ValueError(f"stage must be 'all' or one of {STAGES}.") + stages = STAGES if stage == "all" else (stage,) + out = Path(str(cfg.output_dir)) + out.mkdir(parents=True, exist_ok=True) + caches = open_caches(str(cfg.cells_dir)) + clustering = ClusteringConfig.from_dict(OmegaConf.to_container(cfg.clustering, resolve=True)) + positive_class = int(cfg.positive_class) + summary: dict = {"models": list(caches)} + folds_path, records_path = out / "fold_niches.joblib", out / "matched_niches.csv" + + if "cluster" in stages: + folds = cluster_folds( + caches, clustering, positive_class=positive_class, n_boot=int(cfg.n_boot) + ) + joblib.dump(folds, folds_path) + folds.records.to_csv(out / "fold_niches.csv", index=False) + if "match" in stages or "transfer" in stages: + if not folds_path.exists(): + raise FileNotFoundError(f"{folds_path} is missing; run stage=cluster first.") + folds = joblib.load(folds_path) + if "match" in stages: + start, stop, n = cfg.match.scan + records, table, threshold, scan = match_niches( + folds, + threshold=cfg.match.threshold, + scan=np.linspace(start, stop, int(n)), + min_models=int(cfg.match.min_models), + ) + records.to_csv(records_path, index=False) + table.to_csv(out / "metaclusters.csv", index=False) + if scan is not None: + scan.to_csv(out / "threshold_scan.csv", index=False) + units, patients = held_out_shares(folds, records) + units.to_csv(out / "unit_metacluster_shares.csv") + positive = {} + for unit, pos in folds.positive.items(): + positive[folds.patients.get(unit, unit)] = ( + positive.get(folds.patients.get(unit, unit), False) or pos + ) + share_tests(patients, positive).to_csv(out / "patient_metacluster_tests.csv", index=False) + summary["match"] = { + "threshold": threshold, + "n_metaclusters": len(table), + "n_replicated": int(table["replicated"].sum()), + } + _overview( + table, + records, + folds.cell_types, + out / "metaclusters.png", + positive_name=str(cfg.positive_name), + ) + if "transfer" in stages: + if not records_path.exists(): + raise FileNotFoundError(f"{records_path} is missing; run stage=match first.") + result = transfer_niches( + caches, folds, pd.read_csv(records_path), positive_class=positive_class + ) + tdir = out / "transfer" + tdir.mkdir(exist_ok=True) + for name, frame in result.items(): + frame.to_csv(tdir / f"{name}.csv") + summary["transfer"] = {"mean_agreement": float(np.nanmean(result["agreement"].to_numpy()))} + if "basis" in stages and cfg.basis.enabled: + result = shared_basis( + caches, + space_model=cfg.basis.model, + space=str(cfg.basis.space), + k=int(cfg.basis.k), + sample_size=int(cfg.basis.sample_size), + seed=int(cfg.basis.seed), + n_boot=int(cfg.n_boot), + positive_class=positive_class, + ) + bdir = out / "basis" + bdir.mkdir(exist_ok=True) + result["groups"].to_csv(bdir / "groups.csv", index=False) + result["pairs"].to_csv(bdir / "attribution_pairs.csv", index=False) + (out / "summary.json").write_text(json.dumps(summary, indent=1, default=str)) + return summary + + +def _overview( + table: pd.DataFrame, records: pd.DataFrame, cell_types, path: Path, *, positive_name: str +) -> None: + """Composition per metacluster and every member's margin.""" + import matplotlib.pyplot as plt + + if table.empty: + return + frac = [f"frac_{c}" for c in cell_types] + y = np.arange(len(table))[::-1] + fig, (ax0, ax1) = plt.subplots( + 1, 2, figsize=(14, 0.45 * len(table) + 2), gridspec_kw={"width_ratios": [2.5, 1.2]} + ) + im = ax0.imshow(table[frac].to_numpy() * 100, aspect="auto", cmap="viridis", vmin=0) + ax0.set_xticks(range(len(frac)), cell_types, rotation=90, fontsize=7) + ax0.set_yticks( + range(len(table)), + [f"{m} ({n} models)" for m, n in zip(table["metacluster"], table["n_models"])], + fontsize=7, + ) + plt.colorbar(im, ax=ax0, fraction=0.03, label="% of cells") + for row, meta in zip(y, table["metacluster"]): + members = records[records["metacluster"] == int(meta[1:])] + ax1.scatter(members["margin"], np.full(len(members), row), s=18, color="#1b9e77") + ax1.axvline(0, color="k", lw=0.8) + ax1.set_yticks(y, []) + ax1.set_xlabel(f"margin of each member niche (+ = {positive_name})") + fig.tight_layout() + fig.savefig(path, dpi=150) + plt.close(fig) + + +@hydra.main(version_base="1.3", config_path="../../configs", config_name="metaclusters.yaml") +def main(cfg: DictConfig) -> None: + run_metaclusters(cfg) + + +if __name__ == "__main__": + main() diff --git a/src/grass_mil/niches/replication/stages.py b/src/grass_mil/niches/replication/stages.py new file mode 100644 index 0000000..d50a64f --- /dev/null +++ b/src/grass_mil/niches/replication/stages.py @@ -0,0 +1,367 @@ +"""Replication of niches across the models of a cross-validation (optional analysis). + +1. ``cluster``: discover the niches of every fold model on the units it held out. +2. ``match``: group the niches of different models into metaclusters by the + Hellinger distance of their cell-type compositions, with average linkage + that never merges two niches of the same model. The distance threshold is + fixed or chosen to maximise the metaclusters found in ``min_models`` models. +3. ``transfer``: every model labels every unit with its own niches (mapped to + metaclusters); reports train vs held-out behaviour, agreement between models + and consensus metacluster shares per patient. +4. ``basis``: one label-free k-means partition of a single space, attributed by + every model, so the margins of different models refer to the same cells. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, field +from typing import Sequence + +import numpy as np +import pandas as pd + +from grass_mil.inference.cell_cache import CellCache +from grass_mil.spatial.stats import bh, bootstrap_mean + +from ..attribution import cell_margins +from ..cells import Cells, load_cells +from ..discovery import BACKGROUND, ClusteringConfig, NicheModel, fit_niches +from ..tables import niche_table + +__all__ = [ + "FoldNiches", + "cluster_folds", + "hellinger", + "constrained_linkage", + "match_niches", + "held_out_shares", + "transfer_niches", + "shared_basis", + "share_tests", +] + +log = logging.getLogger(__name__) + + +@dataclass +class FoldNiches: + records: pd.DataFrame # one row per (model, niche) + labels: dict[str, np.ndarray] = field( + default_factory=dict + ) # unit -> niche ids (held-out model) + model_of_unit: dict[str, str] = field(default_factory=dict) + models: dict[str, NicheModel] = field(default_factory=dict) + cell_types: list[str] = field(default_factory=list) + patients: dict[str, str] = field(default_factory=dict) + positive: dict[str, bool] = field(default_factory=dict) + + +def _margins(cells: Cells, positive_class: int) -> np.ndarray: + return cell_margins( + cells.instance_logits, + cells.attention_logits, + cells.head_bias, + cells.unit_of_cell, + positive_class=positive_class, + ) + + +def cluster_folds( + caches: dict[str, CellCache], cfg: ClusteringConfig, *, positive_class: int, n_boot: int +) -> FoldNiches: + frames, out = [], FoldNiches(records=pd.DataFrame()) + for name, cache in caches.items(): + cells = load_cells(cache, sorted(cache.held_out)) + margin = _margins(cells, positive_class) + model, niches = fit_niches(cells.space[cfg.space], margin, cfg) + positive = cells.positive(positive_class) + table = niche_table( + niches, + cells, + margin, + positive=positive, + present_threshold=0.02, + enrichment_threshold=1.5, + n_boot=n_boot, + seed=cfg.seed, + ) + table = table[table["niche_id"] != BACKGROUND].copy() + table.insert(0, "model", name) + table["id"] = name + ":" + table["niche"] + table["association_vs_base"] = table["association"] - table.attrs.get("base_rate", np.nan) + frames.append(table) + for unit in cells.units: + out.labels[unit] = niches[cells.unit_of_cell == unit] + out.model_of_unit[unit] = name + out.models[name] = model + out.cell_types = cells.cell_type_names + out.patients.update(cells.patients) + out.positive.update(positive) + log.info(f"[{name}] {len(cells.units)} held-out units, {model.n_niches} niches") + out.records = pd.concat(frames, ignore_index=True) + return out + + +def hellinger(compositions: np.ndarray) -> np.ndarray: + """Pairwise Hellinger distance between composition rows (0 to 1).""" + h = np.sqrt(np.clip(compositions, 0.0, None)) + return np.sqrt(0.5 * ((h[:, None, :] - h[None, :, :]) ** 2).sum(-1)) + + +def constrained_linkage( + distance: np.ndarray, model: Sequence[str], threshold: float +) -> list[list[int]]: + """Greedy average linkage below ``threshold``; groups sharing a model never merge.""" + model = np.asarray(model) + groups = [[i] for i in range(distance.shape[0])] + while True: + best = None + for a in range(len(groups)): + models_a = set(model[groups[a]]) + for b in range(a + 1, len(groups)): + if models_a & set(model[groups[b]]): + continue + d = float(distance[np.ix_(groups[a], groups[b])].mean()) + if d < threshold and (best is None or d < best[0]): + best = (d, a, b) + if best is None: + return groups + _, a, b = best + groups[a] += groups.pop(b) + + +def match_niches( + folds: FoldNiches, *, threshold: float | None, scan: Sequence[float], min_models: int +) -> tuple[pd.DataFrame, pd.DataFrame, float, pd.DataFrame | None]: + """Metacluster of every niche record, the metacluster table and the threshold used.""" + records = folds.records.reset_index(drop=True).copy() + comp_cols = [f"frac_{c}" for c in folds.cell_types] + distance = hellinger(records[comp_cols].to_numpy(float)) + scan_table = None + if threshold is None: + rows, best = [], None + for t in scan: + groups = constrained_linkage(distance, records["model"], float(t)) + replicated = sum(len(g) >= min_models for g in groups) + rows.append( + {"threshold": float(t), "n_metaclusters": len(groups), "n_replicated": replicated} + ) + if best is None or (replicated, -len(groups)) > best[0]: + best = ((replicated, -len(groups)), float(t), groups) + _, threshold, groups = best + scan_table = pd.DataFrame(rows) + else: + groups = constrained_linkage(distance, records["model"], float(threshold)) + records["metacluster"] = 0 + for k, members in enumerate(sorted(groups, key=len, reverse=True), 1): + records.loc[members, "metacluster"] = k + rows = [] + for k, g in records.groupby("metacluster"): + comp = np.average(g[comp_cols].to_numpy(float), axis=0, weights=g["n_cells"]) + idx = g.index.to_numpy() + rows.append( + { + "metacluster": f"M{k}", + "n_models": g["model"].nunique(), + "replicated": g["model"].nunique() >= min_models, + "members": ", ".join(g["id"]), + "within_distance": float( + distance[np.ix_(idx, idx)][np.triu_indices(len(idx), 1)].mean() + ) + if len(idx) > 1 + else 0.0, + "share": float(g["share"].mean()), + "association_vs_base": float(g["association_vs_base"].mean()), + "margin": float(g["margin"].mean()), + "margin_sd": float(g["margin"].std(ddof=1)) if len(g) > 1 else float("nan"), + "sign_agreement": float( + (np.sign(g["margin"]) == np.sign(g["margin"].mean())).mean() + ), + "composition": ", ".join( + f"{folds.cell_types[j]} {100 * comp[j]:.0f}%" for j in np.argsort(-comp)[:3] + ), + **{c: float(v) for c, v in zip(comp_cols, comp)}, + } + ) + return records, pd.DataFrame(rows), float(threshold), scan_table + + +def share_tests(shares: pd.DataFrame, positive: dict[str, bool]) -> pd.DataFrame: + """Per column: mean share in positive and other patients, AUROC, Mann-Whitney p and BH q.""" + from scipy.stats import mannwhitneyu + from sklearn.metrics import roc_auc_score + + y = np.array([positive.get(p, False) for p in shares.index]) + rows = [] + for column in shares.columns: + v = shares[column].to_numpy(float) + ok = y.sum() > 2 and (~y).sum() > 2 + rows.append( + { + "metacluster": column, + "share_pos": float(v[y].mean()) if y.any() else float("nan"), + "share_neg": float(v[~y].mean()) if (~y).any() else float("nan"), + "auroc": float(roc_auc_score(y, v)) if ok else float("nan"), + "p": float(mannwhitneyu(v[y], v[~y]).pvalue) if ok else float("nan"), + } + ) + table = pd.DataFrame(rows) + table["q"] = bh(table["p"].to_numpy()) + return table.sort_values("p") + + +def _patient_positive(folds: FoldNiches) -> dict[str, bool]: + out: dict[str, bool] = {} + for unit, pos in folds.positive.items(): + out[folds.patients.get(unit, unit)] = out.get(folds.patients.get(unit, unit), False) or pos + return out + + +def held_out_shares(folds: FoldNiches, records: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]: + """Metacluster shares of every unit (labelled by the model that held it out) and per patient.""" + meta_of = {(r.model, r.niche_id): f"M{r.metacluster}" for r in records.itertuples()} + columns = sorted(set(meta_of.values()), key=lambda m: int(m[1:])) + rows = {} + for unit, labels in folds.labels.items(): + model = folds.model_of_unit[unit] + names = pd.Series([meta_of.get((model, int(n)), "Background") for n in labels]) + rows[unit] = names.value_counts(normalize=True).reindex(columns, fill_value=0.0) + units = pd.DataFrame(rows).T + patients = units.groupby([folds.patients.get(u, u) for u in units.index]).mean() + return units, patients + + +def transfer_niches( + caches: dict[str, CellCache], folds: FoldNiches, records: pd.DataFrame, *, positive_class: int +) -> dict[str, pd.DataFrame]: + """Every model labels every unit with its niches, mapped to metaclusters.""" + from sklearn.metrics import adjusted_rand_score + + meta_of = {(r.model, r.niche_id): int(r.metacluster) for r in records.itertuples()} + names = list(caches) + units = sorted(set().union(*(c.units for c in caches.values()))) + labels, rows = {}, [] + reference = None + for name in names: + cells = load_cells(caches[name], units) + reference = cells if reference is None else reference + niches = folds.models[name].predict(cells.space[folds.models[name].cfg.space]) + meta = np.array([meta_of.get((name, int(n)), 0) for n in niches]) + labels[name] = meta + held = np.isin(cells.unit_of_cell, caches[name].held_out) + positive = np.array([folds.positive.get(u, False) for u in cells.unit_of_cell]) + margin = _margins(cells, positive_class) + for k in sorted(set(meta.tolist()) - {0}): + for split, mask in (("held_out", held), ("train", ~held)): + m = (meta == k) & mask + if mask.any() and m.any(): + rows.append( + { + "metacluster": f"M{k}", + "model": name, + "split": split, + "share": float(m.sum() / mask.sum()), + "association_vs_base": float( + positive[m].mean() - positive[mask].mean() + ), + "margin": float( + pd.Series(margin[m]).groupby(cells.unit_of_cell[m]).mean().mean() + ), + } + ) + ari = pd.DataFrame(np.nan, index=names, columns=names) + for a in names: + for b in names: + both = (labels[a] > 0) & (labels[b] > 0) + if both.sum() > 10: + ari.loc[a, b] = adjusted_rand_score(labels[a][both], labels[b][both]) + metas = sorted(set(np.concatenate(list(labels.values())).tolist()) - {0}) + consensus = pd.DataFrame( + { + f"M{k}": [ + np.mean([np.mean(labels[n][reference.unit_of_cell == u] == k) for n in names]) + for u in units + ] + for k in metas + }, + index=units, + ) + patients = consensus.groupby([folds.patients.get(u, u) for u in consensus.index]).mean() + return { + "by_model": pd.DataFrame(rows), + "agreement": ari, + "unit_shares": consensus, + "patient_shares": patients, + "patient_tests": share_tests(patients, _patient_positive(folds)), + } + + +def shared_basis( + caches: dict[str, CellCache], + *, + space_model: str | None, + space: str, + k: int, + sample_size: int, + seed: int, + n_boot: int, + positive_class: int, +) -> dict[str, pd.DataFrame]: + """A k-means partition of one model's space, attributed by every model.""" + from sklearn.cluster import MiniBatchKMeans + from sklearn.preprocessing import StandardScaler + + names = list(caches) + basis_name = space_model or names[0] + units = sorted(caches[basis_name].units) + base = load_cells(caches[basis_name], units) + x = base.space[space] + rng = np.random.default_rng(seed) + sample = rng.choice(len(x), min(sample_size, len(x)), replace=False) + scaler = StandardScaler().fit(x[sample]) + kmeans = MiniBatchKMeans(k, random_state=seed, n_init=10).fit(scaler.transform(x[sample])) + groups = kmeans.predict(scaler.transform(x)) + 1 + positive = base.positive(positive_class) + is_pos = np.array([positive.get(u, False) for u in base.unit_of_cell]) + pairs = [] + for name in names: + cells = load_cells(caches[name], units) + margin = _margins(cells, positive_class) + held = set(caches[name].held_out) + frame = pd.DataFrame({"unit": cells.unit_of_cell, "group": groups, "margin": margin}) + for (unit, group), g in frame.groupby(["unit", "group"]): + pairs.append( + { + "model": name, + "unit_id": unit, + "group": int(group), + "split": "held_out" if unit in held else "train", + "margin": float(g["margin"].mean()), + } + ) + pairs = pd.DataFrame(pairs) + rows = [] + k_types = len(base.cell_type_names) + for group in range(1, k + 1): + m = groups == group + if not m.any(): + continue + comp = ( + np.bincount(base.cell_type[m], minlength=k_types) / m.sum() if k_types else np.zeros(0) + ) + row = { + "group": f"G{group}", + "share": float(m.mean()), + "association": float(is_pos[m].mean()), + "composition": ", ".join( + f"{base.cell_type_names[j]} {100 * comp[j]:.0f}%" for j in np.argsort(-comp)[:3] + ), + } + for split in ("held_out", "train"): + values = pairs.loc[(pairs.group == group) & (pairs.split == split), "margin"] + row[f"margin_{split}"], row[f"margin_{split}_lo"], row[f"margin_{split}_hi"] = ( + bootstrap_mean(values, n_boot=n_boot, seed=seed) + ) + rows.append(row) + return {"groups": pd.DataFrame(rows), "pairs": pairs} diff --git a/src/grass_mil/niches/survival.py b/src/grass_mil/niches/survival.py new file mode 100644 index 0000000..fd67c90 --- /dev/null +++ b/src/grass_mil/niches/survival.py @@ -0,0 +1,138 @@ +"""Niche shares versus survival. + +Per patient, the share of every niche is averaged over the patient's units and +summarised as the z-scored log share. Each niche gets a univariate Cox hazard +ratio (per standard deviation of the log share) with a Wald test and BH +correction. Optionally, patients are split into tertiles of a niche score (a +signed sum of z-scored log shares, e.g. ``["Niche 3", "-Niche 4"]``) and +compared with Kaplan-Meier curves and a log-rank test. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import Any, Mapping + +import numpy as np +import pandas as pd + +from grass_mil.spatial.stats import bh + +__all__ = ["SurvivalConfig", "niche_survival", "cox_univariate", "kaplan_meier", "logrank"] + + +@dataclass(frozen=True) +class SurvivalConfig: + file: str | None = None + id_column: str = "patient_id" + time_column: str = "time" + event_column: str = "event" + pseudo_share: float = 1e-3 + score: tuple[str, ...] = () + + @classmethod + def from_dict(cls, cfg: Mapping[str, Any] | None) -> SurvivalConfig: + cfg = dict(cfg or {}) + unknown = set(cfg) - {f.name for f in fields(cls)} + if unknown: + raise ValueError(f"Unknown survival keys: {sorted(unknown)}") + cfg["score"] = tuple(cfg.get("score") or ()) + return cls(**cfg) + + +def cox_univariate( + x: np.ndarray, time: np.ndarray, event: np.ndarray, n_iter: int = 50 +) -> tuple[float, float, float, float]: + """Cox model with one covariate (Breslow ties): ``(hr, lo, hi, p)`` with a 95% Wald interval.""" + from scipy.stats import norm + + order = np.argsort(-time, kind="stable") + x, event, t = x[order], event[order].astype(bool), time[order] + # Risk set of cell i: everyone with time >= t_i, including later ties. + last = np.searchsorted(-t, -t, side="right") - 1 + beta = 0.0 + info = np.nan + for _ in range(n_iter): + w = np.exp(beta * x) + s0, s1, s2 = np.cumsum(w)[last], np.cumsum(w * x)[last], np.cumsum(w * x * x)[last] + mean = s1 / s0 + grad = float(np.sum(x[event] - mean[event])) + info = float(np.sum(s2[event] / s0[event] - mean[event] ** 2)) + if info <= 0: + return (float("nan"),) * 4 + step = grad / info + beta += step + if abs(step) < 1e-9: + break + se = 1.0 / np.sqrt(info) + p = 2 * norm.sf(abs(beta / se)) + return ( + float(np.exp(beta)), + float(np.exp(beta - 1.96 * se)), + float(np.exp(beta + 1.96 * se)), + float(p), + ) + + +def kaplan_meier(time: np.ndarray, event: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Step function ``(times, survival)`` starting at ``(0, 1)``.""" + times = np.unique(time[event > 0]) + surv, s = [1.0], 1.0 + for t in times: + at_risk = (time >= t).sum() + s *= 1.0 - ((time == t) & (event > 0)).sum() / at_risk + surv.append(s) + return np.r_[0.0, times], np.asarray(surv) + + +def logrank(time: np.ndarray, event: np.ndarray, group: np.ndarray) -> float: + """p-value of the k-group log-rank test.""" + from scipy.stats import chi2 + + groups = np.unique(group) + o_minus_e = np.zeros(len(groups)) + cov = np.zeros((len(groups), len(groups))) + for t in np.unique(time[event > 0]): + at_risk = time >= t + n = at_risk.sum() + d = ((time == t) & (event > 0)).sum() + n_g = np.array([(at_risk & (group == g)).sum() for g in groups]) + d_g = np.array([((time == t) & (event > 0) & (group == g)).sum() for g in groups]) + o_minus_e += d_g - d * n_g / n + if n > 1: + f = d * (n - d) / (n * n * (n - 1)) + cov += f * (np.diag(n_g * n) - np.outer(n_g, n_g)) + stat = float(o_minus_e[:-1] @ np.linalg.pinv(cov[:-1, :-1]) @ o_minus_e[:-1]) + return float(chi2.sf(stat, len(groups) - 1)) + + +def niche_survival(patient_shares: pd.DataFrame, outcomes: pd.DataFrame, cfg: SurvivalConfig): + """Cox table per niche and, with ``cfg.score``, the patients' score tertiles. + + ``patient_shares``: patients x niches; ``outcomes``: patients x [time, event]. + """ + data = patient_shares.join(outcomes, how="inner") + niches = [c for c in patient_shares.columns if c != "Background"] + z = np.log(data[niches] + cfg.pseudo_share) + z = (z - z.mean()) / z.std(ddof=0).replace(0, np.nan) + time, event = data["time"].to_numpy(float), data["event"].to_numpy(float) + rows = [] + for niche in niches: + hr, lo, hi, p = cox_univariate(z[niche].fillna(0).to_numpy(), time, event) + rows.append({"niche": niche, "hazard_ratio": hr, "hr_lo": lo, "hr_hi": hi, "p": p}) + table = pd.DataFrame(rows) + table["q"] = bh(table["p"].to_numpy()) + groups = None + if cfg.score: + score = np.zeros(len(data)) + for term in cfg.score: + sign, name = (-1.0, term[1:]) if term.startswith("-") else (1.0, term) + if name not in z.columns: + raise ValueError(f"survival.score names an unknown niche '{name}'.") + score += sign * z[name].fillna(0).to_numpy() + tertile = np.digitize(score, np.quantile(score, [1 / 3, 2 / 3])) + groups = pd.DataFrame( + {"score": score, "tertile": tertile, "time": time, "event": event}, index=data.index + ) + groups.attrs["logrank_p"] = logrank(time, event, tertile) + return table, groups diff --git a/src/grass_mil/niches/tables.py b/src/grass_mil/niches/tables.py new file mode 100644 index 0000000..845f3d6 --- /dev/null +++ b/src/grass_mil/niches/tables.py @@ -0,0 +1,93 @@ +"""Per-niche summary and per-unit niche composition.""" + +from __future__ import annotations + +import numpy as np +import pandas as pd + +from grass_mil.spatial.stats import bootstrap_mean + +from .cells import Cells +from .discovery import BACKGROUND, niche_name + +__all__ = ["niche_table", "unit_niche_shares"] + + +def unit_niche_shares(niches: np.ndarray, cells: Cells) -> pd.DataFrame: + """Fraction of every unit's cells in every niche (rows: units).""" + frame = pd.crosstab( + pd.Series(cells.unit_of_cell, name="unit_id"), + pd.Series(niches, name="niche"), + normalize="index", + ) + frame.columns = [niche_name(c) for c in frame.columns] + return frame.reindex(cells.units).fillna(0.0) + + +def niche_table( + niches: np.ndarray, + cells: Cells, + margin: np.ndarray, + *, + positive: dict[str, bool], + present_threshold: float, + enrichment_threshold: float, + n_boot: int, + seed: int, +) -> pd.DataFrame: + """One row per niche (Background last). + + * ``share``: fraction of the analysed cells in the niche; + * ``n_units``: units where it makes up more than ``present_threshold``; + * ``association``: fraction of the niche's cells that come from positive + units (cell-weighted; compare with ``base_rate``); + * ``margin`` [``margin_lo``, ``margin_hi``]: mean over units of the mean + cell margin within the niche, with a 95% bootstrap interval over units; + * ``sign_consistency``: fraction of those units where the niche's margin + has the sign of the overall margin; + * ``frac_``: composition; ``enriched``: cell types over-represented + by more than ``enrichment_threshold`` relative to all analysed cells. + """ + names = cells.cell_type_names + k = len(names) + has_types = k > 0 and (cells.cell_type >= 0).all() + cohort = np.bincount(cells.cell_type, minlength=k) / cells.n if has_types else None + is_pos = np.array([positive.get(u, False) for u in cells.unit_of_cell]) + shares = unit_niche_shares(niches, cells) + rows = [] + for niche in sorted(set(niches.tolist()) - {BACKGROUND}) + ( + [BACKGROUND] if (niches == BACKGROUND).any() else [] + ): + m = niches == niche + per_unit = pd.Series(margin[m]).groupby(cells.unit_of_cell[m]).mean() + mean, lo, hi = bootstrap_mean(per_unit.to_numpy(), n_boot=n_boot, seed=seed) + row = { + "niche": niche_name(niche), + "niche_id": int(niche), + "n_cells": int(m.sum()), + "share": float(m.mean()), + "n_units": int((shares[niche_name(niche)] > present_threshold).sum()), + "association": float(is_pos[m].mean()), + "margin": mean, + "margin_lo": lo, + "margin_hi": hi, + "sign_consistency": float((np.sign(per_unit) == np.sign(mean)).mean()) + if len(per_unit) + else float("nan"), + } + if has_types: + comp = np.bincount(cells.cell_type[m], minlength=k) / m.sum() + ratio = comp / np.maximum(cohort, 1e-12) + row["composition"] = ", ".join( + f"{names[j]} {100 * comp[j]:.0f}%" for j in np.argsort(-comp)[:3] if comp[j] > 0 + ) + row["enriched"] = ", ".join( + f"{names[j]} x{ratio[j]:.1f}" + for j in np.argsort(-ratio) + if ratio[j] > enrichment_threshold + ) + row.update({f"frac_{names[j]}": float(comp[j]) for j in range(k)}) + rows.append(row) + table = pd.DataFrame(rows) + table.attrs["base_rate"] = float(is_pos.mean()) + return table diff --git a/src/grass_mil/predict.py b/src/grass_mil/predict.py index bb377f7..c1079d2 100644 --- a/src/grass_mil/predict.py +++ b/src/grass_mil/predict.py @@ -7,7 +7,6 @@ from pathlib import Path import hydra -import pandas as pd import torch from omegaconf import DictConfig @@ -45,7 +44,8 @@ def run_predict(cfg: DictConfig) -> dict: if wanted is not None: wanted = set(wanted) units = GraphUnitDataset( - [e for e in units.entries if {e.unit_id, e.region_id, e.sample_id} & wanted], units.labels + [e for e in units.entries if {e.unit_id, e.region_id, e.sample_id} & wanted], + units.labels, ) sampler = sampler_of_checkpoint(ckpt, int(pcfg.batch_size)) meta = datamodule.unit_table(units) diff --git a/src/grass_mil/runtime.py b/src/grass_mil/runtime.py index 5979376..007d7df 100644 --- a/src/grass_mil/runtime.py +++ b/src/grass_mil/runtime.py @@ -36,7 +36,9 @@ def instantiate_group(group: DictConfig | None) -> list[Any]: """Instantiate every entry of a ``{name: {_target_: ...}}`` group (callbacks, loggers).""" if not group: return [] - return [hydra.utils.instantiate(conf) for conf in group.values() if conf and "_target_" in conf] + return [ + hydra.utils.instantiate(conf) for conf in group.values() if conf and "_target_" in conf + ] def resolve_device(name: str | None) -> torch.device: @@ -65,7 +67,10 @@ def on_validation_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> # Runs after the module logged its epoch-level validation metrics. if trainer.sanity_checking or not trainer.is_global_zero: return - row: dict[str, Any] = {"epoch": int(trainer.current_epoch), "step": int(trainer.global_step)} + row: dict[str, Any] = { + "epoch": int(trainer.current_epoch), + "step": int(trainer.global_step), + } for key, value in trainer.callback_metrics.items(): if key.startswith(("train/", "val/")) and not key.endswith("_step"): value = float(value) diff --git a/src/grass_mil/spatial/__init__.py b/src/grass_mil/spatial/__init__.py new file mode 100644 index 0000000..f2d4cb2 --- /dev/null +++ b/src/grass_mil/spatial/__init__.py @@ -0,0 +1,2 @@ +"""Classical spatial statistics of labelled cells, shared by the niche descriptors +and the model-free ``grass-mil-spatial-stats`` comparison.""" diff --git a/src/grass_mil/spatial/cli.py b/src/grass_mil/spatial/cli.py new file mode 100644 index 0000000..720bd56 --- /dev/null +++ b/src/grass_mil/spatial/cli.py @@ -0,0 +1,126 @@ +"""``grass-mil-spatial-stats``: model-free spatial statistics, compared between classes.""" + +from __future__ import annotations + +import logging +from pathlib import Path + +import hydra +import numpy as np +import pandas as pd +from omegaconf import DictConfig, OmegaConf + +from grass_mil.inference.cell_cache import CellCache +from grass_mil.niches.cells import load_cells, resolve_units + +from .graph_stats import contact_counts, cross_l, enrichment_zscores +from .neighbourhoods import fit_neighbourhoods, window_composition +from .stats import clr, mann_whitney_table, patient_mean + +__all__ = ["run_spatial_stats", "main"] + +log = logging.getLogger(__name__) + +ANALYSES = ("enrichment", "neighbourhoods", "contacts", "cross_l") + + +def run_spatial_stats(cfg: DictConfig) -> Path: + bad = set(cfg.analyses) - set(ANALYSES) + if bad: + raise ValueError(f"Unknown analyses {sorted(bad)}; expected {list(ANALYSES)}.") + out = Path(str(cfg.output_dir)) + out.mkdir(parents=True, exist_ok=True) + cache = CellCache.open(str(cfg.cells_dir)) + units_sel = cfg.units if isinstance(cfg.units, str) else OmegaConf.to_container(cfg.units) + cells = load_cells(cache, resolve_units(units_sel, cache)) + if (cells.cell_type < 0).any(): + raise ValueError("Spatial statistics need cell types.") + mapping = dict(cfg.group_map or {}) + groups = sorted({mapping.get(n, n) for n in cells.cell_type_names}) + to_group = np.array([groups.index(mapping.get(n, n)) for n in cells.cell_type_names]) + labels = to_group[cells.cell_type] + g = len(groups) + reference = [groups.index(r) for r in cfg.reference] + partners = [j for j in range(g) if j not in reference] + rng = np.random.default_rng(int(cfg.seed)) + positive = cells.positive(int(cfg.positive_class)) + patient_of = {u: cells.patients.get(u, u) for u in cells.units} + y = pd.Series({patient_of[u]: int(positive[u]) for u in cells.units}) + + per_unit: dict[str, dict[str, dict[str, float]]] = {a: {} for a in cfg.analyses} + windows = [] + for unit in cells.units: + m = np.flatnonzero(cells.unit_of_cell == unit) + local = np.full(cells.n, -1) + local[m] = np.arange(m.size) + e = cells.edges[:, np.isin(cells.edges[0], m)] + e = local[e] + lab, pos = labels[m], cells.pos[m] + if "enrichment" in cfg.analyses: + z = enrichment_zscores(lab, e, g, n_perms=int(cfg.n_perms), rng=rng) + per_unit["enrichment"][unit] = { + f"{groups[a]}|{groups[b]}": z[a, b] for a in range(g) for b in range(a, g) + } + if "neighbourhoods" in cfg.analyses: + windows.append(window_composition(pos, lab, g, int(cfg.neighbourhoods.window))) + if "contacts" in cfg.analyses and reference: + counts = contact_counts(lab, e, g) + ref_ends = counts[reference].sum() + per_unit["contacts"][unit] = { + f"{groups[r]}->{groups[p]}": counts[r, p] / ref_ends if ref_ends else np.nan + for r in reference + for p in range(g) + } + if "cross_l" in cfg.analyses and reference: + row = {} + for r in reference: + for p in [*reference, *partners]: + if (lab == r).sum() >= 2 and (lab == p).sum() >= 2: + curve = cross_l( + pos, lab == r, lab == p, np.asarray(cfg.radii, dtype=float) + ) + row.update( + { + f"{groups[r]}->{groups[p]}@{rad:g}": v + for rad, v in zip(cfg.radii, curve) + } + ) + per_unit["cross_l"][unit] = row + if "neighbourhoods" in cfg.analyses: + n_cn = int(cfg.neighbourhoods.n) + cn_labels, centres = fit_neighbourhoods( + windows, n_cn, int(cfg.neighbourhoods.sample_size), int(cfg.seed) + ) + per_unit["neighbourhoods"] = { + u: {f"CN{k + 1}": float(np.mean(lab == k)) for k in range(n_cn)} + for u, lab in zip(cells.units, cn_labels) + } + pd.DataFrame(centres, columns=groups, index=[f"CN{k + 1}" for k in range(n_cn)]).to_csv( + out / "neighbourhood_centres.csv" + ) + + for analysis, rows in per_unit.items(): + if not rows: + continue + frame = pd.DataFrame.from_dict(rows, orient="index") + frame.index.name = "unit_id" + frame.to_csv(out / f"{analysis}_per_unit.csv") + patients = patient_mean(frame, patient_of) + if analysis == "neighbourhoods": + patients = pd.DataFrame( + clr(patients.to_numpy()), index=patients.index, columns=patients.columns + ) + mann_whitney_table(patients, y, min_per_group=int(cfg.min_patients)).to_csv( + out / f"{analysis}_tests.csv", index=False + ) + log.info(f"Spatial statistics of {len(cells.units)} units written to {out}") + return out + + +@hydra.main(version_base="1.3", config_path="../configs", config_name="spatial_stats.yaml") +def main(cfg: DictConfig) -> None: + run_spatial_stats(cfg) + + +if __name__ == "__main__": + main() diff --git a/src/grass_mil/spatial/graph_stats.py b/src/grass_mil/spatial/graph_stats.py new file mode 100644 index 0000000..24010f0 --- /dev/null +++ b/src/grass_mil/spatial/graph_stats.py @@ -0,0 +1,114 @@ +"""Spatial statistics of labelled cells on a cell graph or in the plane. + +* :func:`enrichment_zscores`: neighbourhood enrichment, the z-score of the + observed count of edges between two labels against label permutations over + the fixed graph (the statistic of ``squidpy.gr.nhood_enrichment``). +* :func:`morans_i`: Moran's I of label indicators over the binary adjacency. +* :func:`cross_l`: Ripley's centred cross-L, ``L_ab(r) - r``. +* :func:`contact_counts`: symmetric label-pair edge counts. +* :func:`colocalisation`: ``log2`` observed/expected edges between labels. +""" + +from __future__ import annotations + +import numpy as np + +__all__ = [ + "undirected_edges", + "contact_counts", + "enrichment_zscores", + "morans_i", + "cross_l", + "colocalisation", +] + + +def undirected_edges(edge_index: np.ndarray) -> np.ndarray: + """Each undirected edge once (``src < dst``).""" + a, b = np.asarray(edge_index) + keep = a < b + return np.stack([a[keep], b[keep]]) + + +def contact_counts(labels: np.ndarray, edges: np.ndarray, n_labels: int) -> np.ndarray: + """Symmetric ``[n_labels, n_labels]`` edge-end counts of undirected edges. + + Entry ``(a, b)`` counts the edges between labels ``a`` and ``b``; the + diagonal counts same-label edges twice (once per end), so a row sums to the + number of edge ends carrying that label. + """ + counts = np.zeros((n_labels, n_labels)) + np.add.at(counts, (labels[edges[0]], labels[edges[1]]), 1.0) + return counts + counts.T + + +def enrichment_zscores( + labels: np.ndarray, edges: np.ndarray, n_labels: int, *, n_perms: int, rng: np.random.Generator +) -> np.ndarray: + """Neighbourhood enrichment z-scores; labels absent from the graph get NaN.""" + labels = np.asarray(labels, dtype=int) + observed = contact_counts(labels, edges, n_labels) + null = np.stack( + [contact_counts(rng.permutation(labels), edges, n_labels) for _ in range(n_perms)] + ) + with np.errstate(divide="ignore", invalid="ignore"): + z = (observed - null.mean(0)) / null.std(0) + z[~np.isfinite(z)] = 0.0 + absent = np.setdiff1d(np.arange(n_labels), labels) + z[absent, :] = np.nan + z[:, absent] = np.nan + return z + + +def morans_i(indicators: np.ndarray, edges: np.ndarray) -> np.ndarray: + """Moran's I of every column of ``indicators`` ``[n, k]`` over the undirected ``edges``. + + ``I = (n / S0) * z' W z / z' z`` with binary symmetric ``W`` and the + centred column ``z``; NaN for constant columns or a graph without edges. + """ + from scipy.sparse import coo_matrix + + x = np.asarray(indicators, dtype=float) + n = x.shape[0] + if n < 2 or edges.shape[1] == 0: + return np.full(x.shape[1], np.nan) + rows = np.concatenate([edges[0], edges[1]]) + cols = np.concatenate([edges[1], edges[0]]) + w = coo_matrix((np.ones(rows.size), (rows, cols)), shape=(n, n)).tocsr() + z = x - x.mean(0) + with np.errstate(divide="ignore", invalid="ignore"): + return (n / w.sum()) * (z * (w @ z)).sum(0) / (z * z).sum(0) + + +def cross_l( + coords: np.ndarray, in_a: np.ndarray, in_b: np.ndarray, radii: np.ndarray +) -> np.ndarray: + """Centred cross-L ``L_ab(r) - r`` of two point sets on their bounding box. + + ``K_ab(r) = |A| / (n_a n_b) * #{(i in a, j in b): d_ij <= r, i != j}`` with + ``n_a (n_a - 1)`` in the denominator when ``a`` and ``b`` are the same set. + """ + from scipy.spatial import cKDTree + + a, b = coords[in_a], coords[in_b] + span = coords.max(0) - coords.min(0) + area = max(float(np.prod(span)), 1e-12) + same = bool(np.array_equal(in_a, in_b)) + counts = cKDTree(a).count_neighbors(cKDTree(b), radii).astype(float) + if same: + counts -= len(a) + pairs = len(a) * (len(a) - 1) + else: + pairs = len(a) * len(b) + k = area * counts / max(pairs, 1) + return np.sqrt(np.maximum(k, 0.0) / np.pi) - radii + + +def colocalisation(labels: np.ndarray, edges: np.ndarray, n_labels: int) -> np.ndarray: + """``log2((observed + 1) / (expected + 1))`` edges between labels, with the + expectation ``d_a d_b / sum(d)`` from the label degrees (negative labels ignored).""" + keep = (labels[edges[0]] >= 0) & (labels[edges[1]] >= 0) + observed = contact_counts(np.maximum(labels, 0), edges[:, keep], n_labels) + degree = observed.sum(1) + expected = np.outer(degree, degree) / max(degree.sum(), 1.0) + return np.log2((observed + 1.0) / (expected + 1.0)) diff --git a/src/grass_mil/spatial/neighbourhoods.py b/src/grass_mil/spatial/neighbourhoods.py new file mode 100644 index 0000000..1fac84d --- /dev/null +++ b/src/grass_mil/spatial/neighbourhoods.py @@ -0,0 +1,33 @@ +"""Cellular neighbourhoods: the cell-type composition of every cell's k-NN window, clustered.""" + +from __future__ import annotations + +import numpy as np + +__all__ = ["window_composition", "fit_neighbourhoods"] + + +def window_composition( + coords: np.ndarray, labels: np.ndarray, n_labels: int, k: int +) -> np.ndarray: + """Fraction of each label among a cell and its ``k`` nearest neighbours.""" + from scipy.spatial import cKDTree + + _, nn = cKDTree(coords).query(coords, k=min(k + 1, len(coords))) + nn = np.asarray(nn).reshape(len(coords), -1) + counts = np.zeros((len(coords), n_labels), dtype=np.float32) + for column in nn.T: + counts[np.arange(len(coords)), labels[column]] += 1 + return counts / counts.sum(1, keepdims=True) + + +def fit_neighbourhoods(windows: list[np.ndarray], n: int, sample_size: int, seed: int): + """k-means of the windows (fitted on a sample): per-unit labels and the centres.""" + from sklearn.cluster import MiniBatchKMeans + + stacked = np.concatenate(windows) + sample = np.random.default_rng(seed).choice( + len(stacked), min(sample_size, len(stacked)), replace=False + ) + kmeans = MiniBatchKMeans(n, random_state=seed, n_init=10).fit(stacked[sample]) + return [kmeans.predict(w) for w in windows], kmeans.cluster_centers_ diff --git a/src/grass_mil/spatial/stats.py b/src/grass_mil/spatial/stats.py new file mode 100644 index 0000000..081d028 --- /dev/null +++ b/src/grass_mil/spatial/stats.py @@ -0,0 +1,73 @@ +"""Group comparisons and summaries shared by every analysis.""" + +from __future__ import annotations + +from typing import Mapping + +import numpy as np +import pandas as pd + +__all__ = ["bh", "bootstrap_mean", "clr", "patient_mean", "mann_whitney_table"] + + +def bh(p: np.ndarray) -> np.ndarray: + """Benjamini-Hochberg q-values (NaNs kept).""" + from scipy.stats import false_discovery_control + + p = np.asarray(p, dtype=float) + q = np.full(p.shape, np.nan) + ok = np.isfinite(p) + if ok.any(): + q[ok] = false_discovery_control(p[ok], method="bh") + return q + + +def bootstrap_mean(values, *, n_boot: int = 2000, seed: int = 0) -> tuple[float, float, float]: + """Mean and 95% percentile bootstrap interval of the finite ``values``.""" + v = np.asarray(values, dtype=float) + v = v[np.isfinite(v)] + if v.size == 0: + return float("nan"), float("nan"), float("nan") + draws = np.random.default_rng(seed).integers(0, v.size, size=(int(n_boot), v.size)) + means = v[draws].mean(axis=1) + return float(v.mean()), float(np.percentile(means, 2.5)), float(np.percentile(means, 97.5)) + + +def clr(shares: np.ndarray, pseudo: float = 1e-3) -> np.ndarray: + """Centred log-ratio of composition rows.""" + logs = np.log(np.asarray(shares, dtype=float) + pseudo) + return logs - logs.mean(axis=1, keepdims=True) + + +def patient_mean(per_unit: pd.DataFrame, patient_of: Mapping[str, str]) -> pd.DataFrame: + """Mean of every column over the units of each patient (index: unit id).""" + patients = [patient_of.get(str(u), str(u)) for u in per_unit.index] + return per_unit.groupby(patients).mean(numeric_only=True) + + +def mann_whitney_table( + values: pd.DataFrame, y: pd.Series, *, min_per_group: int = 3 +) -> pd.DataFrame: + """Two-sided Mann-Whitney U per column between ``y == 1`` and ``y == 0`` (BH-adjusted).""" + from scipy.stats import mannwhitneyu + + y = y.reindex(values.index) + rows = [] + for column in values.columns: + a = values.loc[y == 1, column].dropna() + b = values.loc[y == 0, column].dropna() + enough = len(a) >= min_per_group and len(b) >= min_per_group + rows.append( + { + "metric": str(column), + "n_pos": len(a), + "n_neg": len(b), + "mean_pos": float(a.mean()) if len(a) else float("nan"), + "mean_neg": float(b.mean()) if len(b) else float("nan"), + "p": float(mannwhitneyu(a, b).pvalue) if enough else float("nan"), + } + ) + table = pd.DataFrame(rows) + table["diff"] = table["mean_pos"] - table["mean_neg"] + table["q"] = bh(table["p"].to_numpy()) + return table.sort_values("p").reset_index(drop=True) diff --git a/src/grass_mil/spatial_stats/__init__.py b/src/grass_mil/spatial_stats/__init__.py deleted file mode 100644 index b846a76..0000000 --- a/src/grass_mil/spatial_stats/__init__.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Classical spatial statistics on the cell graph, per region, compared between groups. - -Neighbourhood enrichment z-scores (label permutations on the Delaunay graph), -cellular neighbourhoods (k-NN window composition clustered with k-means), -cell-cell contact fractions and Ripley cross-K, each summarised per region, -averaged per patient and compared between classes with Mann-Whitney + BH. -""" - -from .cellular_neighbourhoods import cellular_neighbourhoods, window_composition -from .contacts import contact_fractions -from .cross_k import cross_k -from .nhood_enrichment import nhood_enrichment_zscores -from .tests import bh_adjust, mann_whitney_table, patient_mean - -__all__ = [ - "bh_adjust", - "cellular_neighbourhoods", - "contact_fractions", - "cross_k", - "mann_whitney_table", - "nhood_enrichment_zscores", - "patient_mean", - "window_composition", -] diff --git a/src/grass_mil/spatial_stats/cellular_neighbourhoods.py b/src/grass_mil/spatial_stats/cellular_neighbourhoods.py deleted file mode 100644 index 7f8cb7c..0000000 --- a/src/grass_mil/spatial_stats/cellular_neighbourhoods.py +++ /dev/null @@ -1,51 +0,0 @@ -"""Cellular neighbourhoods: k-NN window composition clustered with k-means.""" - -from __future__ import annotations - -from typing import List, Sequence, Tuple - -import numpy as np - -__all__ = ["window_composition", "cellular_neighbourhoods"] - - -def window_composition( - coords: np.ndarray, labels: np.ndarray, *, n_groups: int, k: int = 10 -) -> np.ndarray: - """Fraction of each group among a cell and its ``k`` nearest neighbours.""" - from scipy.spatial import cKDTree - - coords = np.asarray(coords, dtype=float) - labels = np.asarray(labels).astype(int) - n = coords.shape[0] - kk = min(int(k) + 1, n) - _, nn = cKDTree(coords).query(coords, k=kk) - nn = np.asarray(nn).reshape(n, -1) - W = np.zeros((n, n_groups), dtype=np.float32) - for j in range(nn.shape[1]): - W[np.arange(n), labels[nn[:, j]]] += 1 - return W / W.sum(axis=1, keepdims=True) - - -def cellular_neighbourhoods( - windows: Sequence[np.ndarray], - *, - n_cn: int = 10, - sample_size: int = 200_000, - seed: int = 0, - n_init: int = 10, - batch_size: int = 4096, -) -> Tuple[List[np.ndarray], np.ndarray]: - """Fit ``MiniBatchKMeans(n_cn)`` on a sample of all windows; return per-region CN labels and centres.""" - from sklearn.cluster import MiniBatchKMeans - - W = np.concatenate(list(windows)) - rng = np.random.default_rng(int(seed)) - sub = rng.choice(len(W), min(int(sample_size), len(W)), replace=False) - km = MiniBatchKMeans( - int(n_cn), random_state=int(seed), n_init=int(n_init), batch_size=int(batch_size) - ).fit(W[sub]) - out = [] - for w in windows: - out.append(km.predict(np.asarray(w, dtype=np.float32))) - return out, km.cluster_centers_ diff --git a/src/grass_mil/spatial_stats/cli.py b/src/grass_mil/spatial_stats/cli.py deleted file mode 100644 index d7ffe58..0000000 --- a/src/grass_mil/spatial_stats/cli.py +++ /dev/null @@ -1,61 +0,0 @@ -"""``grass-mil-spatial-stats``: classical spatial statistics on a cell cache.""" - -from __future__ import annotations - -from typing import Any, Dict - -import hydra -from omegaconf import DictConfig, OmegaConf - -from grass_mil._paths import setup_project_root - -setup_project_root() - -from grass_mil.inference.cell_cache import CellCache # noqa: E402 -from grass_mil.interpretability.metaclusters_cli import load_labels_and_patients # noqa: E402 -from grass_mil.interpretability.niches_cli import resolve_regions # noqa: E402 -from grass_mil.spatial_stats.run import run_spatial_stats # noqa: E402 -from grass_mil.utils import RankedLogger # noqa: E402 - -log = RankedLogger(__name__, rank_zero_only=True) - -__all__ = ["run", "main"] - - -def run(cfg: DictConfig) -> Dict[str, Any]: - cache = CellCache.open(str(cfg.cells_dir)) - labels_of, patient_of = load_labels_and_patients(cfg.labels) - regions = resolve_regions(cfg.get("regions") or "all", cache) - return run_spatial_stats( - cache, - labels_of=labels_of, - patient_of=patient_of, - out_dir=str(cfg.output_dir), - regions=regions, - group_map=OmegaConf.to_container(cfg.group_map, resolve=True) - if cfg.get("group_map") - else None, - reference=list(cfg.reference) if cfg.get("reference") else None, - partners=list(cfg.partners) if cfg.get("partners") else None, - infiltrating=list(cfg.infiltrating) if cfg.get("infiltrating") else (), - n_perms=int(cfg.nhood.n_perms), - nhood_backend=str(cfg.nhood.backend), - window_k=int(cfg.cn.window_k), - n_cn=int(cfg.cn.n_cn), - cn_sample_size=int(cfg.cn.sample_size), - radii=[float(r) for r in cfg.crossk.radii], - seed=int(cfg.seed), - min_per_group=int(cfg.min_per_group), - positive_name=str(cfg.positive_name), - analyses=list(cfg.analyses), - log=log, - ) - - -@hydra.main(version_base="1.3", config_path="../configs", config_name="spatial_stats.yaml") -def main(cfg: DictConfig) -> None: - run(cfg) - - -if __name__ == "__main__": - main() diff --git a/src/grass_mil/spatial_stats/contacts.py b/src/grass_mil/spatial_stats/contacts.py deleted file mode 100644 index 3ed0201..0000000 --- a/src/grass_mil/spatial_stats/contacts.py +++ /dev/null @@ -1,56 +0,0 @@ -"""Contact fractions on the cell graph.""" - -from __future__ import annotations - -from typing import Dict, Sequence - -import numpy as np - -__all__ = ["contact_fractions"] - - -def contact_fractions( - labels: np.ndarray, - edge_index: np.ndarray, - *, - group_names: Sequence[str], - reference: Sequence[str], - partners: Sequence[str], - infiltrating: Sequence[str] = (), - min_cells: int = 5, -) -> Dict[str, float]: - """Fraction of reference cells touching >= 1 partner cell, homotypic contact of - the reference group, and for each infiltrating group the fraction of its cells - whose neighbours are mostly reference cells.""" - labels = np.asarray(labels).astype(int) - ei = np.asarray(edge_index).astype(int) - names = list(group_names) - gl = np.asarray(names, dtype=object)[labels] - n = labels.size - src = np.concatenate([ei[0], ei[1]]) - dst = np.concatenate([ei[1], ei[0]]) - ref = np.isin(gl, list(reference)) - out: Dict[str, float] = {} - for x in partners: - has = np.zeros(n, dtype=bool) - has[src[gl[dst] == x]] = True - out[f"contact_{x}"] = float(has[ref].mean()) if ref.sum() else float("nan") - has = np.zeros(n, dtype=bool) - has[src[ref[dst]]] = True - out["homotypic"] = float(has[ref].mean()) if ref.sum() else float("nan") - n_ref_nb = np.zeros(n) - deg = np.zeros(n) - np.add.at(n_ref_nb, src, ref[dst].astype(float)) - np.add.at(deg, src, 1.0) - frac_ref = n_ref_nb / np.maximum(deg, 1.0) - out["reference_fraction"] = float(ref.mean()) - for x in infiltrating: - mk = gl == x - out[f"{x}_intra_reference"] = ( - float((frac_ref[mk] > 0.5).mean()) if mk.sum() >= int(min_cells) else float("nan") - ) - if infiltrating: - imm = np.isin(gl, list(infiltrating)) - out["infiltration"] = float((frac_ref[imm] > 0.5).mean()) if imm.sum() else float("nan") - out["infiltrating_fraction"] = float(imm.mean()) - return out diff --git a/src/grass_mil/spatial_stats/cross_k.py b/src/grass_mil/spatial_stats/cross_k.py deleted file mode 100644 index 168216c..0000000 --- a/src/grass_mil/spatial_stats/cross_k.py +++ /dev/null @@ -1,48 +0,0 @@ -"""Ripley cross-K between a reference group and partner groups, normalised by pi r^2.""" - -from __future__ import annotations - -from typing import Dict, Sequence - -import numpy as np - -__all__ = ["cross_k"] - - -def cross_k( - coords: np.ndarray, - labels: np.ndarray, - *, - group_names: Sequence[str], - reference: Sequence[str], - partners: Sequence[str], - radii: Sequence[float] = (20.0, 50.0, 100.0), - min_cells: int = 5, -) -> Dict[str, float]: - """``K(r) = area * pairs_within_r / (n_ref * n_partner) / (pi r^2)``; 1 = CSR. - - Self pairs are removed when the partner group is the reference itself. - """ - from scipy.spatial import cKDTree - - coords = np.asarray(coords, dtype=float) - labels = np.asarray(labels).astype(int) - gl = np.asarray(list(group_names), dtype=object)[labels] - ref = np.isin(gl, list(reference)) - out: Dict[str, float] = {} - if ref.sum() < int(min_cells): - return out - area = float(np.prod(coords.max(axis=0) - coords.min(axis=0))) - tree_ref = cKDTree(coords[ref]) - for x in partners: - mk = ref if x == "reference" else (gl == x) - if mk.sum() < int(min_cells): - continue - tree_x = cKDTree(coords[mk]) - self_pairs = int(mk.sum()) if x == "reference" else 0 - for r in radii: - cnt = tree_ref.count_neighbors(tree_x, float(r)) - self_pairs - out[f"K_{x}_r{int(r) if float(r).is_integer() else r}"] = float( - area * cnt / (ref.sum() * mk.sum()) / (np.pi * float(r) ** 2) - ) - return out diff --git a/src/grass_mil/spatial_stats/nhood_enrichment.py b/src/grass_mil/spatial_stats/nhood_enrichment.py deleted file mode 100644 index 9066019..0000000 --- a/src/grass_mil/spatial_stats/nhood_enrichment.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Neighbourhood enrichment z-scores on a cell graph (label permutation null).""" - -from __future__ import annotations - -from typing import Optional - -import numpy as np - -__all__ = ["nhood_enrichment_zscores", "pair_counts"] - - -def pair_counts(labels: np.ndarray, edge_index: np.ndarray, n_groups: int) -> np.ndarray: - """Symmetric (undirected) count of edges between group pairs.""" - la = labels[edge_index[0]] - lb = labels[edge_index[1]] - M = np.zeros((n_groups, n_groups), dtype=np.float64) - np.add.at(M, (la, lb), 1.0) - return M + M.T - - -def nhood_enrichment_zscores( - labels: np.ndarray, - edge_index: np.ndarray, - *, - n_groups: int, - n_perms: int = 200, - seed: int = 0, - backend: str = "numpy", - coords: Optional[np.ndarray] = None, -) -> np.ndarray: - """z-score of observed vs permuted group-pair edge counts (squidpy's statistic). - - ``edge_index`` holds each undirected edge once or twice; counts are - symmetrised either way. ``backend='squidpy'`` delegates to - ``squidpy.gr.nhood_enrichment`` on the same graph (optional dependency). - Groups absent from the region get NaN rows/columns. - """ - labels = np.asarray(labels).astype(int) - ei = np.asarray(edge_index).astype(int) - if ei.ndim != 2 or ei.shape[0] != 2: - raise ValueError("edge_index must have shape (2, n_edges)") - # one direction only, so symmetrisation counts every edge exactly once per orientation - a, b = ei - keep = a < b - if keep.sum() == 0 and ei.shape[1] > 0: - keep = np.ones(ei.shape[1], dtype=bool) - ei = np.stack([a[keep], b[keep]]) - if backend == "squidpy": - return _squidpy_zscores( - labels, ei, n_groups=n_groups, n_perms=n_perms, seed=seed, coords=coords - ) - present = np.unique(labels) - rng = np.random.default_rng(int(seed)) - observed = pair_counts(labels, ei, n_groups) - perms = np.zeros((int(n_perms), n_groups, n_groups)) - for i in range(int(n_perms)): - perms[i] = pair_counts(rng.permutation(labels), ei, n_groups) - mean = perms.mean(axis=0) - std = perms.std(axis=0) - with np.errstate(divide="ignore", invalid="ignore"): - z = (observed - mean) / std - z[~np.isfinite(z)] = 0.0 - absent = np.setdiff1d(np.arange(n_groups), present) - z[absent, :] = np.nan - z[:, absent] = np.nan - return z - - -def _squidpy_zscores(labels, ei, *, n_groups, n_perms, seed, coords): - import anndata as ad - import pandas as pd - import scipy.sparse as sp - import squidpy as sq - - n = labels.size - A = sp.csr_matrix((np.ones(ei.shape[1]), (ei[0], ei[1])), shape=(n, n)) - A = ((A + A.T) > 0).astype(np.float32) - adata = ad.AnnData(X=np.zeros((n, 1), np.float32)) - adata.obs["g"] = pd.Categorical(labels.astype(str)) - adata.obsm["spatial"] = ( - np.asarray(coords, dtype=float) if coords is not None else np.zeros((n, 2)) - ) - adata.obsp["spatial_connectivities"] = A - adata.uns["spatial_neighbors"] = { - "connectivities_key": "spatial_connectivities", - "distances_key": "spatial_distances", - "params": {"n_neighbors": 6, "radius": None, "transform": None}, - } - sq.gr.nhood_enrichment( - adata, - cluster_key="g", - n_perms=int(n_perms), - seed=int(seed), - show_progress_bar=False, - n_jobs=1, - ) - z = adata.uns["g_nhood_enrichment"]["zscore"] - cats = [int(c) for c in adata.obs["g"].cat.categories] - out = np.full((n_groups, n_groups), np.nan) - for x, cx in enumerate(cats): - for y, cy in enumerate(cats): - out[cx, cy] = z[x, y] - return out diff --git a/src/grass_mil/spatial_stats/run.py b/src/grass_mil/spatial_stats/run.py deleted file mode 100644 index 90b7ba3..0000000 --- a/src/grass_mil/spatial_stats/run.py +++ /dev/null @@ -1,224 +0,0 @@ -"""Run the classical spatial statistics over a cell cache and compare groups.""" - -from __future__ import annotations - -import json -from pathlib import Path -from typing import Any, Dict, List, Optional, Sequence - -import numpy as np -import pandas as pd - -from grass_mil.inference.cell_cache import CellCache - -from .cellular_neighbourhoods import cellular_neighbourhoods, window_composition -from .contacts import contact_fractions -from .cross_k import cross_k -from .nhood_enrichment import nhood_enrichment_zscores -from .tests import clr, mann_whitney_table, patient_mean - -__all__ = ["run_spatial_stats"] - - -def _resolve_groups(cell_type_names: Sequence[str], group_map: Optional[Dict[str, str]]): - """Fine cell types -> coarse groups (identity by default).""" - names = list(cell_type_names) - if not group_map: - groups = sorted(set(names)) - else: - groups = sorted({group_map.get(n, n) for n in names}) - gi = {g: i for i, g in enumerate(groups)} - fine_to_group = np.array([gi[(group_map or {}).get(n, n)] for n in names], dtype=int) - return groups, fine_to_group - - -def run_spatial_stats( - cache: CellCache, - *, - labels_of: Dict[str, int], - patient_of: Optional[Dict[str, str]] = None, - out_dir: str | Path, - regions: Optional[Sequence[str]] = None, - group_map: Optional[Dict[str, str]] = None, - reference: Optional[Sequence[str]] = None, - partners: Optional[Sequence[str]] = None, - infiltrating: Sequence[str] = (), - n_perms: int = 200, - nhood_backend: str = "numpy", - window_k: int = 10, - n_cn: int = 10, - cn_sample_size: int = 200_000, - radii: Sequence[float] = (20.0, 50.0, 100.0), - seed: int = 0, - min_per_group: int = 10, - positive_name: str = "positive", - analyses: Sequence[str] = ("nhood", "cn", "contacts", "crossk"), - log=None, -) -> Dict[str, Any]: - out = Path(out_dir) - out.mkdir(parents=True, exist_ok=True) - regions = sorted(str(r) for r in (regions if regions is not None else cache.regions)) - regions = [r for r in regions if r in labels_of] - names = cache.cell_type_names or [] - groups, fine_to_group = _resolve_groups(names, group_map) - G = len(groups) - pat_of = patient_of or {r: r for r in regions} - y_pat = pd.Series({pat_of.get(r, r): int(labels_of[r]) for r in regions}) - reference = list(reference) if reference else [] - partners = list(partners) if partners else [g for g in groups if g not in reference] - analyses = [a.lower() for a in analyses] - - Z: Dict[str, np.ndarray] = {} - windows: List[np.ndarray] = [] - contacts: List[Dict[str, Any]] = [] - crossk: List[Dict[str, Any]] = [] - for i, r in enumerate(regions): - d = cache.load(r, keys=["ct", "pos", "edge_index"]) - g = fine_to_group[d["ct"].astype(int)] - pos = d["pos"].astype(float) - ei = d["edge_index"].astype(int) - if "nhood" in analyses: - Z[r] = nhood_enrichment_zscores( - g, ei, n_groups=G, n_perms=n_perms, seed=seed, backend=nhood_backend, coords=pos - ) - if "cn" in analyses: - windows.append(window_composition(pos, g, n_groups=G, k=window_k)) - if "contacts" in analyses and reference: - rec = {"region_id": r} - rec.update( - contact_fractions( - g, - ei, - group_names=groups, - reference=reference, - partners=partners, - infiltrating=infiltrating, - ) - ) - contacts.append(rec) - if "crossk" in analyses and reference: - rec = {"region_id": r} - rec.update( - cross_k( - pos, - g, - group_names=groups, - reference=reference, - partners=partners + ["reference"], - radii=radii, - ) - ) - crossk.append(rec) - if log is not None and (i + 1) % 50 == 0: - log.info(f" {i + 1}/{len(regions)} regions") - - md = [ - f"# Classical spatial statistics ({len(regions)} regions, {len(y_pat)} patients; per region -> patient mean; Mann-Whitney {positive_name} vs rest, BH within analysis)", - "", - ] - results: Dict[str, Any] = {"regions": len(regions), "groups": groups} - if "nhood" in analyses and Z: - ZT = np.stack([Z[r] for r in regions]) - rows = {} - for a in range(G): - for b in range(a, G): - rows[f"{groups[a]}|{groups[b]}"] = ZT[:, a, b] - per_region = pd.DataFrame(rows, index=regions) - per_region.to_csv(out / "nhood_zscores_per_region.csv") - T = mann_whitney_table( - patient_mean(per_region, pat_of), y_pat, min_per_group=min_per_group - ) - T[["a", "b"]] = T["metric"].str.split("|", expand=True) - T.to_csv(out / "nhood_tests.csv", index=False) - results["nhood_tests"] = str(out / "nhood_tests.csv") - md += [ - f"## Neighbourhood enrichment (z-scores on the cell graph, {n_perms} label permutations, {nhood_backend}): pairs with q < 0.1 (top 20 by p otherwise)", - "", - "| Pair | z pos | z neg | diff | p | q |", - "|---|---|---|---|---|---|", - ] - sel = T[T.q < 0.1] if (T.q < 0.1).sum() else T.head(20) - for _, r in sel.iterrows(): - md.append( - f"| {r.a}-{r.b} | {r.mean_pos:+.2f} | {r.mean_neg:+.2f} | {r['diff']:+.2f} | {r.p:.3g} | {r.q:.2g} |" - ) - if "cn" in analyses and windows: - cn_labels, centres = cellular_neighbourhoods( - windows, n_cn=n_cn, sample_size=cn_sample_size, seed=seed - ) - CN = pd.DataFrame( - {f"CN{k}": [float(np.mean(lab == k)) for lab in cn_labels] for k in range(n_cn)}, - index=regions, - ) - CN.to_csv(out / "cn_shares_per_region.csv") - CNp = patient_mean(CN, pat_of) - L = pd.DataFrame(clr(CNp.to_numpy()), index=CNp.index, columns=CNp.columns) - T = mann_whitney_table(L, y_pat, min_per_group=min_per_group).rename( - columns={"mean_pos": "clr_pos", "mean_neg": "clr_neg", "diff": "clr_diff"} - ) - share_pos = CNp[y_pat.reindex(CNp.index) == 1].mean() - share_neg = CNp[y_pat.reindex(CNp.index) == 0].mean() - T["share_pos"] = T["metric"].map(share_pos) - T["share_neg"] = T["metric"].map(share_neg) - T["composition"] = T["metric"].map( - { - f"CN{k}": ", ".join( - f"{groups[j]} {100 * centres[k][j]:.0f}%" for j in np.argsort(-centres[k])[:3] - ) - for k in range(n_cn) - } - ) - T.to_csv(out / "cn_tests.csv", index=False) - pd.DataFrame(centres, columns=groups, index=[f"CN{k}" for k in range(n_cn)]).to_csv( - out / "cn_centres.csv" - ) - results["cn_tests"] = str(out / "cn_tests.csv") - md += [ - "", - f"## Cellular neighbourhoods ({window_k}-NN windows, k-means k={n_cn}), patient-level CLR shares", - "", - "| CN | composition | share pos | share neg | CLR diff | p | q |", - "|---|---|---|---|---|---|---|", - ] - for _, r in T.iterrows(): - md.append( - f"| {r.metric} | {r.composition} | {100 * r.share_pos:.1f}% | {100 * r.share_neg:.1f}% | {r.clr_diff:+.2f} | {r.p:.3g} | {r.q:.2g} |" - ) - if contacts: - C = pd.DataFrame(contacts).set_index("region_id") - C.to_csv(out / "contacts_per_region.csv") - T = mann_whitney_table(patient_mean(C, pat_of), y_pat, min_per_group=min_per_group) - T.to_csv(out / "contact_tests.csv", index=False) - results["contact_tests"] = str(out / "contact_tests.csv") - md += [ - "", - f"## Contact fractions (reference = {', '.join(reference)})", - "", - "| Metric | pos | neg | p | q |", - "|---|---|---|---|---|", - ] - for _, r in T.iterrows(): - md.append( - f"| {r.metric} | {r.mean_pos:.3f} | {r.mean_neg:.3f} | {r.p:.3g} | {r.q:.2g} |" - ) - if crossk: - X = pd.DataFrame(crossk).set_index("region_id") - X.to_csv(out / "crossk_per_region.csv") - T = mann_whitney_table(patient_mean(X, pat_of), y_pat, min_per_group=min_per_group) - T.to_csv(out / "crossk_tests.csv", index=False) - results["crossk_tests"] = str(out / "crossk_tests.csv") - md += [ - "", - f"## Cross-K (Ripley) reference -> partner, normalised by pi r^2 (1 = complete spatial randomness); radii {list(radii)}", - "", - "| Metric | pos | neg | p | q |", - "|---|---|---|---|---|", - ] - for _, r in T.iterrows(): - md.append( - f"| {r.metric} | {r.mean_pos:.2f} | {r.mean_neg:.2f} | {r.p:.3g} | {r.q:.2g} |" - ) - (out / "report_spatial.md").write_text("\n".join(md) + "\n") - results["report"] = str(out / "report_spatial.md") - (out / "spatial_stats_summary.json").write_text(json.dumps(results, indent=1, default=str)) - return results diff --git a/src/grass_mil/spatial_stats/tests.py b/src/grass_mil/spatial_stats/tests.py deleted file mode 100644 index c327dcd..0000000 --- a/src/grass_mil/spatial_stats/tests.py +++ /dev/null @@ -1,64 +0,0 @@ -"""Per-region -> per-patient summaries and two-group tests with BH correction.""" - -from __future__ import annotations - -from typing import Dict - -import numpy as np -import pandas as pd - -__all__ = ["bh_adjust", "mann_whitney_table", "patient_mean", "clr"] - - -def bh_adjust(p) -> np.ndarray: - from scipy.stats import false_discovery_control - - p = np.asarray(p, dtype=float) - q = np.full(p.size, np.nan) - ok = ~np.isnan(p) - if ok.sum(): - q[ok] = false_discovery_control(p[ok], method="bh") - return q - - -def patient_mean(per_region: pd.DataFrame, patient_of: Dict[str, str]) -> pd.DataFrame: - """Mean of every column per patient (index = region id).""" - df = per_region.copy() - df["__patient"] = [patient_of.get(str(r), str(r)) for r in df.index] - return df.groupby("__patient").mean(numeric_only=True) - - -def clr(M: np.ndarray, pseudo: float = 1e-4) -> np.ndarray: - L = np.log(np.asarray(M, dtype=float) + pseudo) - return L - L.mean(axis=1, keepdims=True) - - -def mann_whitney_table( - values: pd.DataFrame, y: pd.Series, *, min_per_group: int = 3 -) -> pd.DataFrame: - """Mann-Whitney U (two-sided) per column between y == 1 and y == 0, BH-adjusted.""" - from scipy import stats - - yy = y.reindex(values.index) - rows = [] - for col in values.columns: - v = values[col] - a = v[yy == 1].dropna() - b = v[yy == 0].dropna() - ok = len(a) >= int(min_per_group) and len(b) >= int(min_per_group) - rows.append( - dict( - metric=str(col), - n_pos=int(len(a)), - n_neg=int(len(b)), - mean_pos=float(a.mean()) if len(a) else float("nan"), - mean_neg=float(b.mean()) if len(b) else float("nan"), - diff=float(a.mean() - b.mean()) if len(a) and len(b) else float("nan"), - p=float(stats.mannwhitneyu(a, b, alternative="two-sided").pvalue) - if ok - else float("nan"), - ) - ) - T = pd.DataFrame(rows) - T["q"] = bh_adjust(T["p"].to_numpy()) - return T.sort_values("p").reset_index(drop=True) diff --git a/src/grass_mil/train.py b/src/grass_mil/train.py index 2be226d..f96a17c 100644 --- a/src/grass_mil/train.py +++ b/src/grass_mil/train.py @@ -8,7 +8,7 @@ import hydra import lightning as L -from omegaconf import DictConfig, OmegaConf, open_dict +from omegaconf import DictConfig, open_dict from grass_mil.inference.predictions import prediction_frame from grass_mil.runtime import EpochMetricsWriter, instantiate_group, start_run From 1ab3973b82f68bf70ea36d1f317c085fe32c5763 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 15 Oct 2026 08:59:38 +0200 Subject: [PATCH 3/9] Replace the test suite with behavioural tests per subject Data (inputs, graphs, units, labels, splits, PCA, stale-graph guard), sampling (exact k-hop ego-graphs, every cell scored once, fixed roots, interior roots, radius cap), models (pooling identity, Cox loss, probe folding, augmentations, clipping hook, freezing), training (every regime, regression, survival, rebuild from checkpoint, no-validation runs), nested CV (selection fallbacks, refit without validation), prediction, extraction, niches, spatial statistics, replication, configs and entry points, and the whole workflow through the installed commands. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NqExFvaq8RhDYBj6RShCMk --- src/grass_mil/data/datamodule.py | 5 +- src/grass_mil/niches/cells.py | 5 - tests/configs/test_component_instantiation.py | 98 -- tests/configs/test_task_preset_composition.py | 103 -- tests/conftest.py | 260 +--- .../test_model_component_contracts.py | 10 - .../data/test_datamodule_sampler_exposure.py | 127 -- tests/data/test_loader_flags.py | 97 -- .../data/test_precompute_manifest_regions.py | 39 - tests/data/test_precompute_reducer_modes.py | 149 --- tests/data/test_samplers.py | 782 ----------- tests/data/test_seed_sampling.py | 223 ---- tests/eval/test_aggregation.py | 328 ----- tests/eval/test_collectors.py | 293 ----- tests/eval/test_embeddings.py | 21 - tests/eval/test_interpretability_export.py | 392 ------ tests/eval/test_io.py | 49 - tests/eval/test_metrics.py | 231 ---- tests/eval/test_predict_cli.py | 748 ----------- tests/helpers/cohort.py | 149 +++ tests/helpers/config_paths.py | 12 - tests/helpers/package_available.py | 32 - tests/helpers/run_if.py | 142 -- tests/helpers/run_sh_command.py | 22 - tests/helpers/synthetic_cohort.py | 176 --- tests/helpers/synthetic_datamodule.py | 125 -- tests/interpretability/test_attribution.py | 325 ----- tests/interpretability/test_cell_level.py | 104 -- tests/interpretability/test_core.py | 120 -- .../interpretability/test_fidelity_phase6.py | 292 ---- .../interpretability/test_fidelity_phase7.py | 176 --- tests/interpretability/test_niches.py | 236 ---- .../test_pipeline_and_report.py | 390 ------ .../test_plugins_and_tier2.py | 861 ------------ .../test_spatial_statistics.py | 242 ---- tests/interpretability/test_tissue_graph.py | 109 -- tests/interpretability/test_transfer.py | 210 --- tests/models/test_attention.py | 68 - tests/models/test_augmentations.py | 243 ---- tests/models/test_backbones.py | 80 -- tests/models/test_bagging.py | 106 -- tests/models/test_fidelity_phase1.py | 114 -- tests/models/test_fidelity_phase2.py | 284 ---- tests/models/test_fidelity_phase3.py | 295 ----- tests/models/test_fidelity_phase4.py | 236 ---- tests/models/test_fidelity_phase5.py | 273 ---- tests/models/test_losses.py | 63 - tests/models/test_node_embeddings.py | 181 --- tests/models/test_optimization.py | 153 --- tests/models/test_runtime.py | 615 --------- tests/models/test_ssl.py | 79 -- tests/spatial_stats/__init__.py | 0 tests/spatial_stats/test_fidelity_phase8.py | 171 --- tests/test_configs.py | 111 +- tests/test_data.py | 245 ++++ tests/test_e2e_synthetic.py | 227 ---- tests/test_end_to_end.py | 77 ++ tests/test_eval.py | 47 - tests/test_inference.py | 54 + tests/test_loocv.py | 223 ---- tests/test_models.py | 174 +++ tests/test_nested_cv.py | 66 + tests/test_niches.py | 85 ++ tests/test_packaging.py | 68 - tests/test_samplers.py | 154 +++ tests/test_spatial.py | 106 ++ tests/test_spatial_omics_datamodule.py | 1170 ----------------- tests/test_train.py | 202 --- tests/test_train_regimes.py | 132 -- tests/test_training.py | 113 ++ 70 files changed, 1377 insertions(+), 12521 deletions(-) delete mode 100644 tests/configs/test_component_instantiation.py delete mode 100644 tests/configs/test_task_preset_composition.py delete mode 100644 tests/contracts/test_model_component_contracts.py delete mode 100644 tests/data/test_datamodule_sampler_exposure.py delete mode 100644 tests/data/test_loader_flags.py delete mode 100644 tests/data/test_precompute_manifest_regions.py delete mode 100644 tests/data/test_precompute_reducer_modes.py delete mode 100644 tests/data/test_samplers.py delete mode 100644 tests/data/test_seed_sampling.py delete mode 100644 tests/eval/test_aggregation.py delete mode 100644 tests/eval/test_collectors.py delete mode 100644 tests/eval/test_embeddings.py delete mode 100644 tests/eval/test_interpretability_export.py delete mode 100644 tests/eval/test_io.py delete mode 100644 tests/eval/test_metrics.py delete mode 100644 tests/eval/test_predict_cli.py create mode 100644 tests/helpers/cohort.py delete mode 100644 tests/helpers/config_paths.py delete mode 100644 tests/helpers/package_available.py delete mode 100644 tests/helpers/run_if.py delete mode 100644 tests/helpers/run_sh_command.py delete mode 100644 tests/helpers/synthetic_cohort.py delete mode 100644 tests/helpers/synthetic_datamodule.py delete mode 100644 tests/interpretability/test_attribution.py delete mode 100644 tests/interpretability/test_cell_level.py delete mode 100644 tests/interpretability/test_core.py delete mode 100644 tests/interpretability/test_fidelity_phase6.py delete mode 100644 tests/interpretability/test_fidelity_phase7.py delete mode 100644 tests/interpretability/test_niches.py delete mode 100644 tests/interpretability/test_pipeline_and_report.py delete mode 100644 tests/interpretability/test_plugins_and_tier2.py delete mode 100644 tests/interpretability/test_spatial_statistics.py delete mode 100644 tests/interpretability/test_tissue_graph.py delete mode 100644 tests/interpretability/test_transfer.py delete mode 100644 tests/models/test_attention.py delete mode 100644 tests/models/test_augmentations.py delete mode 100644 tests/models/test_backbones.py delete mode 100644 tests/models/test_bagging.py delete mode 100644 tests/models/test_fidelity_phase1.py delete mode 100644 tests/models/test_fidelity_phase2.py delete mode 100644 tests/models/test_fidelity_phase3.py delete mode 100644 tests/models/test_fidelity_phase4.py delete mode 100644 tests/models/test_fidelity_phase5.py delete mode 100644 tests/models/test_losses.py delete mode 100644 tests/models/test_node_embeddings.py delete mode 100644 tests/models/test_optimization.py delete mode 100644 tests/models/test_runtime.py delete mode 100644 tests/models/test_ssl.py delete mode 100644 tests/spatial_stats/__init__.py delete mode 100644 tests/spatial_stats/test_fidelity_phase8.py create mode 100644 tests/test_data.py delete mode 100644 tests/test_e2e_synthetic.py create mode 100644 tests/test_end_to_end.py delete mode 100644 tests/test_eval.py create mode 100644 tests/test_inference.py delete mode 100644 tests/test_loocv.py create mode 100644 tests/test_models.py create mode 100644 tests/test_nested_cv.py create mode 100644 tests/test_niches.py delete mode 100644 tests/test_packaging.py create mode 100644 tests/test_samplers.py create mode 100644 tests/test_spatial.py delete mode 100644 tests/test_spatial_omics_datamodule.py delete mode 100644 tests/test_train.py delete mode 100644 tests/test_train_regimes.py create mode 100644 tests/test_training.py diff --git a/src/grass_mil/data/datamodule.py b/src/grass_mil/data/datamodule.py index 20d4e8d..c9c667f 100644 --- a/src/grass_mil/data/datamodule.py +++ b/src/grass_mil/data/datamodule.py @@ -180,7 +180,7 @@ def _group_of(self, entry: UnitEntry) -> str: def _random_assignment(self, units: Sequence[UnitEntry]) -> dict[str, str]: groups = {u.unit_id: self._group_of(u) for u in units} - # Stratify by the first target when the units carry discrete labels. + # Stratify by the first target when it holds a few integer classes. strata = None labels = self.read_labels() if labels is not None: @@ -189,7 +189,8 @@ def _random_assignment(self, units: Sequence[UnitEntry]) -> dict[str, str]: label_id = labels.resolve(u.unit_id, u.region_id, u.sample_id) if label_id is not None: values[groups[u.unit_id]] = float(labels.targets.loc[label_id].iloc[0]) - if len(values) == len(set(groups.values())) and len(set(values.values())) <= 10: + discrete = all(float(v).is_integer() for v in values.values()) + if discrete and len(values) == len(set(groups.values())) and len(set(values.values())) <= 10: strata = {g: str(v) for g, v in values.items()} ratios = self.split_cfg.get("ratios", (0.8, 0.1, 0.1)) return random_split(groups, ratios, int(self.split_cfg.get("seed", 0)), strata) diff --git a/src/grass_mil/niches/cells.py b/src/grass_mil/niches/cells.py index 4ddcb72..31602f3 100644 --- a/src/grass_mil/niches/cells.py +++ b/src/grass_mil/niches/cells.py @@ -39,11 +39,6 @@ def n(self) -> int: def unit_class(self, unit: str) -> int: return class_index(self.targets[unit][0]) - def unit_value(self, unit: str) -> float: - """The value that defines a positive unit: the class, or the event (survival).""" - t = self.targets[unit] - return float(t[1]) if self.target_type == "survival" else float(t[0]) - def positive(self, positive_class: int) -> dict[str, bool]: """Which units count as positive (class ``positive_class``; event for survival).""" if self.target_type == "classification": diff --git a/tests/configs/test_component_instantiation.py b/tests/configs/test_component_instantiation.py deleted file mode 100644 index bd85d10..0000000 --- a/tests/configs/test_component_instantiation.py +++ /dev/null @@ -1,98 +0,0 @@ -"""Compose the model/encoder and model/attention groups through Hydra and build them. - -The groups ship as plain parameter dictionaries (no ``_target_``) that the -module factories turn into components, so the only faithful test is to compose -them the way the CLI does and instantiate the resulting module. -""" - -import pytest -from hydra import compose, initialize_config_module -from hydra.core.global_hydra import GlobalHydra -from omegaconf import OmegaConf - -from tests.helpers.config_paths import CONFIGS_DIR - -CONFIG_MODULE = "grass_mil.configs" -ENCODERS = ["gin", "gcn", "gat", "graphsage", "gine", "gine_length"] -ATTENTIONS = ["gated", "gated_projected"] - - -def _compose(overrides: list[str]): - GlobalHydra.instance().clear() - with initialize_config_module(version_base="1.3", config_module=CONFIG_MODULE): - cfg = compose(config_name="train", overrides=overrides) - GlobalHydra.instance().clear() - return cfg - - -@pytest.mark.parametrize("encoder", ENCODERS) -def test_encoder_group_builds(encoder: str) -> None: - from grass_mil.models.components.backbones import EncoderConfig, GNNEncoder - - cfg = _compose([f"model/encoder={encoder}", "model.encoder.input_dim=8"]) - enc_cfg = OmegaConf.to_container(cfg.model.encoder, resolve=True) - if enc_cfg.get("categorical_embedding") is not None: - enc_cfg["categorical_embedding"] = { - **enc_cfg["categorical_embedding"], - "num_embeddings": 5, - "column_index": 0, - } - encoder_module = GNNEncoder(EncoderConfig(**enc_cfg)) - assert encoder_module.output_dim == int(cfg.model.encoder.out_dim) - - -@pytest.mark.parametrize("attention", ATTENTIONS) -def test_attention_group_builds(attention: str) -> None: - from grass_mil.models.components.factory import build_attention - - cfg = _compose([f"model/attention={attention}", "task=finetune_mil"]) - attn_cfg = OmegaConf.to_container(cfg.model.attention, resolve=True) - module = build_attention(use_attention=True, **attn_cfg) - assert module is not None - assert module.attention_c.out_features == int(cfg.model.graph_head.output_dim) - - -@pytest.mark.parametrize( - ("model_name", "task_name"), - [ - ("bgrl_module", "pretrain_bgrl"), - ("supervised_module", "finetune_mean"), - ("supervised_module", "finetune_mil"), - ("grass_mil", "finetune_mil"), - ], -) -def test_training_module_presets_instantiate(model_name: str, task_name: str) -> None: - import hydra - - cfg = _compose([f"model={model_name}", f"task={task_name}"]) - module = hydra.utils.instantiate(cfg.model) - assert module is not None - - -def test_grass_mil_preset_defaults_match_preprint() -> None: - cfg = _compose(["model=grass_mil", "task=finetune_mil"]) - enc = cfg.model.encoder - assert enc.conv_type == "gine" and int(enc.num_layers) == 4 and int(enc.hidden_dim) == 256 - assert enc.norm == "layernorm" and enc.pooling == "max" and bool(enc.train_eps) - assert int(enc.edge_feature_index) == 0 and int(enc.edge_attr_dim) == 1 - att = cfg.model.attention - assert int(att.projection_dim) == 32 and int(att.hidden_dim) == 8 and bool(att.dropout) - assert int(att.n_classes) == 2 and int(cfg.model.graph_head.output_dim) == 2 - assert int(cfg.model.graph_head.num_layers) == 2 - - -def test_sampler_strategy_resolution(): - from grass_mil.data.components.samplers import get_sampler_strategy - - base = CONFIGS_DIR / "data" / "sampler" - for name in ["identity", "shadow_native", "shadow_custom"]: - cfg = OmegaConf.load(base / f"{name}.yaml") - strategy = get_sampler_strategy(OmegaConf.to_container(cfg, resolve=True)) - assert strategy is not None - - -def test_default_data_sampler_config(): - cfg = _compose([]) - assert cfg.data.sampler.name == "shadow_custom" - cfg_mil = _compose(["data/sampler=ego_graph_mil"]) - assert cfg_mil.data.sampler.runtime.unit_as_bag is True diff --git a/tests/configs/test_task_preset_composition.py b/tests/configs/test_task_preset_composition.py deleted file mode 100644 index 6915686..0000000 --- a/tests/configs/test_task_preset_composition.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Compose task presets through Hydra the way the CLI does. - -The rest of the training tests build their config by loading YAML files directly -and assigning the nodes onto a base config. That is convenient, but it bypasses -Hydra's defaults-list resolution entirely, so a task preset can declare an -override that never takes effect and every test still passes. - -Two real bugs reached the docs that way: ``task=finetune_mil`` silently composed -a single-logit head with a single attention channel, and ``task=pretrain_bgrl`` -failed to compose at all. Both are invisible unless the config is composed the -way the entrypoints compose it. -""" - -import pytest -from hydra import compose, initialize_config_module -from hydra.core.global_hydra import GlobalHydra - -CONFIG_MODULE = "grass_mil.configs" -# Every root config that builds a SupervisedModule. `inference/predict` was -# missing here, and it had the same defaults-ordering bug: it composed a -# single-logit head, then failed to load a two-class fine-tuned checkpoint with -# a size mismatch. -ROOT_CONFIGS = ["train", "eval", "inference/predict"] -SUPERVISED_TASKS = ["finetune_mean", "finetune_mil", "finetune_survival"] - - -def _compose(config_name: str, overrides: list[str]): - GlobalHydra.instance().clear() - if config_name == "inference/predict": - # Mandatory field; irrelevant to what these tests assert. - overrides = [*overrides, "ckpt_path=/tmp/none.ckpt"] - with initialize_config_module(version_base="1.3", config_module=CONFIG_MODULE): - cfg = compose(config_name=config_name, overrides=overrides) - GlobalHydra.instance().clear() - return cfg - - -@pytest.mark.parametrize("config_name", ROOT_CONFIGS) -@pytest.mark.parametrize("task_name", SUPERVISED_TASKS) -def test_supervised_task_presets_compose(config_name: str, task_name: str) -> None: - cfg = _compose(config_name, [f"task={task_name}"]) - assert cfg.task.name == task_name - - -@pytest.mark.parametrize("config_name", ROOT_CONFIGS) -@pytest.mark.parametrize("task_name", SUPERVISED_TASKS) -def test_attention_width_matches_head_width_as_composed(config_name: str, task_name: str) -> None: - """The contract must hold on the *composed* config, not just in code. - - ``validate_attention_width`` runs at build time, but it can only see what - composition produced. A preset whose widths are silently discarded still - passes it, because one channel for one class is internally consistent. - """ - # Imported lazily on purpose. This module is collected early, and importing - # grass_mil at module scope pulls in numpy/MKL before torch spawns the - # ddp_sim subprocesses, which then die with - # "MKL_THREADING_LAYER=INTEL is incompatible with libgomp.so.1". - from grass_mil.contracts import validate_attention_width - - cfg = _compose(config_name, [f"task={task_name}"]) - validate_attention_width( - attention_width=int(cfg.model.attention.n_classes), - num_classes=int(cfg.model.graph_head.output_dim), - ) - - -@pytest.mark.parametrize("config_name", ROOT_CONFIGS) -def test_finetune_mil_composes_per_class_widths(config_name: str) -> None: - """Regression: the task preset's ``model`` block must survive composition. - - ``task/finetune_mil.yaml`` declares a two-class head and two attention - channels. If ``model`` is merged after ``task`` in the defaults list, the - module preset overwrites both and the published MIL head is not what runs. - """ - cfg = _compose(config_name, ["task=finetune_mil"]) - assert int(cfg.model.graph_head.output_dim) == 2 - assert int(cfg.model.attention.n_classes) == 2 - assert cfg.task.target_type == "categorical" - assert cfg.task.loss == "categorical_ce" - - -@pytest.mark.parametrize("config_name", ROOT_CONFIGS) -def test_pretrain_bgrl_composes_and_overrides_the_scheduler(config_name: str) -> None: - """Regression: ``defaults: - override /scheduler`` must actually resolve. - - This raised ``Could not override 'scheduler'`` from the CLI while every - training test passed, because the tests never composed this preset. - """ - cfg = _compose(config_name, ["task=pretrain_bgrl", "model=bgrl_module"]) - assert cfg.task.name == "pretrain_bgrl" - assert cfg.model._target_.endswith("BGRLModule") - # cosine_step, not the cosine_epoch default. - assert int(cfg.scheduler.T_max) == 10000 - - -@pytest.mark.parametrize("config_name", ROOT_CONFIGS) -def test_default_task_is_unchanged_by_defaults_ordering(config_name: str) -> None: - """Guards the reorder itself: the no-override baseline must not move.""" - cfg = _compose(config_name, []) - assert cfg.task.name == "finetune_mean" - assert int(cfg.model.graph_head.output_dim) == 1 - assert int(cfg.model.attention.n_classes) == 1 - assert int(cfg.scheduler.T_max) == 100 diff --git a/tests/conftest.py b/tests/conftest.py index 634d545..f8d36ff 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,198 +1,72 @@ -"""This file prepares config fixtures for other tests.""" +"""Shared fixtures: one synthetic cohort and, built once per session, a pretrained +encoder, a cross-validation and the per-cell outputs of its models.""" + +from __future__ import annotations from pathlib import Path import pytest -import rootutils -from hydra import compose, initialize_config_module -from hydra.core.global_hydra import GlobalHydra -from omegaconf import DictConfig, OmegaConf, open_dict - - -@pytest.fixture(scope="package") -def cfg_train_global() -> DictConfig: - """A pytest fixture for setting up a default Hydra DictConfig for training. - - :return: A DictConfig object containing a default Hydra configuration for training. - """ - with initialize_config_module(version_base="1.3", config_module="grass_mil.configs"): - cfg = compose(config_name="train.yaml", return_hydra_config=True, overrides=[]) - - # set defaults for all tests - with open_dict(cfg): - cfg.paths.root_dir = str(rootutils.find_root(indicator=".project-root")) - cfg.trainer.max_epochs = 1 - cfg.trainer.limit_train_batches = 1.0 - cfg.trainer.limit_val_batches = 1.0 - cfg.trainer.limit_test_batches = 1.0 - cfg.trainer.accelerator = "cpu" - cfg.trainer.devices = 1 - # Use self-contained synthetic data instead of real manifest - cfg.data = OmegaConf.create( - { - "_target_": "tests.helpers.synthetic_datamodule.SyntheticBagDataModule", - "batch_size": 4, - "num_workers": 0, - "pin_memory": False, - "input_dim": 8, - } - ) - cfg.model.encoder.input_dim = 8 - cfg.extras.print_config = False - cfg.extras.enforce_tags = False - cfg.logger = None - - return cfg - - -@pytest.fixture(scope="package") -def cfg_eval_global() -> DictConfig: - """A pytest fixture for setting up a default Hydra DictConfig for evaluation. - - :return: A DictConfig containing a default Hydra configuration for evaluation. - """ - with initialize_config_module(version_base="1.3", config_module="grass_mil.configs"): - cfg = compose(config_name="eval.yaml", return_hydra_config=True, overrides=["ckpt_path=."]) - - # set defaults for all tests - with open_dict(cfg): - cfg.paths.root_dir = str(rootutils.find_root(indicator=".project-root")) - cfg.trainer.max_epochs = 1 - cfg.trainer.limit_test_batches = 1.0 - cfg.trainer.accelerator = "cpu" - cfg.trainer.devices = 1 - # Use self-contained synthetic data instead of real manifest - cfg.data = OmegaConf.create( - { - "_target_": "tests.helpers.synthetic_datamodule.SyntheticBagDataModule", - "batch_size": 4, - "num_workers": 0, - "pin_memory": False, - "input_dim": 8, - } - ) - cfg.model.encoder.input_dim = 8 - cfg.extras.print_config = False - cfg.extras.enforce_tags = False - cfg.logger = None - - return cfg - - -@pytest.fixture(scope="package") -def cfg_predict_global() -> DictConfig: - """A pytest fixture for inference prediction configuration.""" - with initialize_config_module(version_base="1.3", config_module="grass_mil.configs"): - cfg = compose( - config_name="inference/predict.yaml", - return_hydra_config=True, - overrides=["ckpt_path=."], - ) - - with open_dict(cfg): - cfg.paths.root_dir = str(rootutils.find_root(indicator=".project-root")) - cfg.trainer.max_epochs = 1 - cfg.trainer.limit_test_batches = 1.0 - cfg.trainer.accelerator = "cpu" - cfg.trainer.devices = 1 - cfg.data = OmegaConf.create( - { - "_target_": "tests.helpers.synthetic_datamodule.SyntheticBagDataModule", - "batch_size": 4, - "num_workers": 0, - "pin_memory": False, - "input_dim": 8, - } - ) - cfg.model.encoder.input_dim = 8 - cfg.extras.print_config = False - cfg.extras.enforce_tags = False - cfg.logger = None - - return cfg - - -@pytest.fixture(scope="package") -def cfg_interpret_global() -> DictConfig: - with initialize_config_module(version_base="1.3", config_module="grass_mil.configs"): - cfg = compose( - config_name="interpretability/report.yaml", - return_hydra_config=True, - overrides=[], - ) - with open_dict(cfg): - cfg.paths.root_dir = str(rootutils.find_root(indicator=".project-root")) - cfg.extras.print_config = False - cfg.extras.enforce_tags = False - cfg.report.pdf.enabled = False - return cfg - - -@pytest.fixture(scope="function") -def cfg_train(cfg_train_global: DictConfig, tmp_path: Path) -> DictConfig: - """A pytest fixture built on top of the `cfg_train_global()` fixture, which accepts a temporary - logging path `tmp_path` for generating a temporary logging path. - - This is called by each test which uses the `cfg_train` arg. Each test generates its own temporary logging path. - - :param cfg_train_global: The input DictConfig object to be modified. - :param tmp_path: The temporary logging path. - - :return: A DictConfig with updated output and log directories corresponding to `tmp_path`. - """ - cfg = cfg_train_global.copy() - - with open_dict(cfg): - cfg.paths.output_dir = str(tmp_path) - cfg.paths.log_dir = str(tmp_path) - - yield cfg - - GlobalHydra.instance().clear() - - -@pytest.fixture(scope="function") -def cfg_interpret(cfg_interpret_global: DictConfig, tmp_path: Path) -> DictConfig: - cfg = cfg_interpret_global.copy() - with open_dict(cfg): - cfg.paths.output_dir = str(tmp_path) - cfg.paths.log_dir = str(tmp_path) - cfg.output_dir = str(tmp_path / "interpretability_report") - yield cfg - GlobalHydra.instance().clear() - - -@pytest.fixture(scope="function") -def cfg_predict(cfg_predict_global: DictConfig, tmp_path: Path) -> DictConfig: - cfg = cfg_predict_global.copy() - - with open_dict(cfg): - cfg.paths.output_dir = str(tmp_path) - cfg.paths.log_dir = str(tmp_path) - - yield cfg - - GlobalHydra.instance().clear() - - -@pytest.fixture(scope="function") -def cfg_eval(cfg_eval_global: DictConfig, tmp_path: Path) -> DictConfig: - """A pytest fixture built on top of the `cfg_eval_global()` fixture, which accepts a temporary - logging path `tmp_path` for generating a temporary logging path. - - This is called by each test which uses the `cfg_eval` arg. Each test generates its own temporary logging path. - - :param cfg_train_global: The input DictConfig object to be modified. - :param tmp_path: The temporary logging path. - - :return: A DictConfig with updated output and log directories corresponding to `tmp_path`. - """ - cfg = cfg_eval_global.copy() - - with open_dict(cfg): - cfg.paths.output_dir = str(tmp_path) - cfg.paths.log_dir = str(tmp_path) - - yield cfg - GlobalHydra.instance().clear() +from tests.helpers.cohort import Cohort, cohort_overrides, compose_cfg, make_cohort + + +@pytest.fixture(scope="session") +def cohort(tmp_path_factory) -> Cohort: + return make_cohort(tmp_path_factory.mktemp("cohort")) + + +@pytest.fixture(scope="session") +def pretrained(cohort: Cohort, tmp_path_factory) -> Path: + """``last.ckpt`` of a few BGRL steps.""" + from grass_mil.train import train + + out = tmp_path_factory.mktemp("pretrain") + cfg = compose_cfg( + "train", + cohort_overrides(cohort, finetune=False) + + [ + "regime=pretrain", + "pretrain.total_steps=4", + "pretrain.warmup_steps=1", + "pretrain.augmentation.noise_features=[area]", + ], + out, + ) + train(cfg) + return out / "checkpoints" / "last.ckpt" + + +@pytest.fixture(scope="session") +def nested_cv(cohort: Cohort, pretrained: Path, tmp_path_factory) -> Path: + """Output directory of a three-fold nested cross-validation.""" + from grass_mil.nested_cv import run_nested_cv + + out = tmp_path_factory.mktemp("nested_cv") + cfg = compose_cfg( + "nested_cv", + cohort_overrides(cohort) + + [ + f"model.init_from_ckpt={pretrained}", + f"nested_cv.folds_file={cohort.folds}", + "nested_cv.inner_k=2", + ], + out, + ) + run_nested_cv(cfg) + return out + + +@pytest.fixture(scope="session") +def cells_dir(cohort: Cohort, nested_cv: Path, tmp_path_factory) -> Path: + """Per-cell outputs of every fold model (``/fold``).""" + from grass_mil.extract import run_extract + + out = tmp_path_factory.mktemp("extract") + cfg = compose_cfg( + "extract", + cohort_overrides(cohort) + + [f"extract.nested_cv_summary={nested_cv / 'nested_cv_summary.json'}"], + out, + ) + run_extract(cfg) + return out / "cells" diff --git a/tests/contracts/test_model_component_contracts.py b/tests/contracts/test_model_component_contracts.py deleted file mode 100644 index 329ce93..0000000 --- a/tests/contracts/test_model_component_contracts.py +++ /dev/null @@ -1,10 +0,0 @@ -from pathlib import Path - - -def test_contract_markdown_exists_and_has_required_sections(): - path = Path("docs/model_component_contracts.md") - assert path.exists() - text = path.read_text() - assert "Graph Unit Contract" in text - assert "Loss Contracts" in text - assert "Sampler Contracts" in text diff --git a/tests/data/test_datamodule_sampler_exposure.py b/tests/data/test_datamodule_sampler_exposure.py deleted file mode 100644 index d642eca..0000000 --- a/tests/data/test_datamodule_sampler_exposure.py +++ /dev/null @@ -1,127 +0,0 @@ -from pathlib import Path - -from grass_mil.data.components.samplers import SamplerConfig -from grass_mil.data.spatial_omics_datamodule import SpatialOmicsDataModule - - -def _minimal_datamodule(tmp_path: Path, sampler: dict | None) -> SpatialOmicsDataModule: - data_dir = tmp_path / "data" - data_dir.mkdir(parents=True, exist_ok=True) - return SpatialOmicsDataModule( - data_dir=str(data_dir), - raw_manifest_path=str(data_dir / "raw" / "manifest.csv"), - processed_dir=str(data_dir / "processed"), - coord_scale_um=1.0, - sample_unit="tile", - batch_size=1, - num_workers=0, - pin_memory=False, - reducer_scope="sample", - keep_raw_molecular=False, - force_precompute=False, - min_cells=1, - use_molecular_features=False, - categorical_features={"include_labels": []}, - split={"train_val_test_split": [1.0, 0.0, 0.0], "split_by": "patch"}, - csv={ - "sep": ",", - "coord_columns": ["x", "y"], - "cell_id_column": "cell_id", - "categorical_label_columns": [], - "molecular_columns": None, - }, - h5ad={ - "coord_columns": ["x", "y"], - "cell_id_column": None, - "categorical_label_columns": [], - "molecular_layer": None, - }, - sce={ - "assay_name": None, - "coord_source": "colData", - "coord_key": None, - "coord_columns": ["x", "y"], - "cell_id_column": None, - "categorical_label_columns": [], - "transpose_assay": True, - }, - graph_builder={ - "name": "delaunay", - "kwargs": {"edge_features": ["distance"], "neighbor_cutoff_um": 50.0}, - }, - feature_reducer={"name": "identity", "kwargs": {}}, - tiling={"tile_size_um": 50.0, "stride_um": 50.0, "min_cells": 1}, - sampler=sampler, - val_sampler=None, - manifest={ - "sample_id": "sample_id", - "input_path": "input_path", - "input_type": "input_type", - "region_id": "region_id", - "polygons_path": "polygons_path", - }, - graph_labels=None, - transforms=[], - ) - - -def test_datamodule_sampler_defaults_to_identity(tmp_path: Path): - dm = _minimal_datamodule(tmp_path, sampler=None) - strategy = dm._ensure_sampler_strategy() - assert strategy.__class__.__name__ == "IdentityBatchStrategy" - assert dm.sampler_strategy is strategy - - -def test_datamodule_exposes_configured_sampler_strategy(tmp_path: Path): - dm = _minimal_datamodule( - tmp_path, - sampler={ - "name": "shadow_custom", - "kwargs": {}, - "runtime": { - "enabled": True, - "depth": 2, - "num_neighbors": 8, - "subgraph_batch_size": 4, - "replace": False, - "shuffle_subgraphs": True, - }, - }, - ) - strategy = dm._ensure_sampler_strategy() - assert strategy.__class__.__name__ == "ShadowCustomStrategy" - assert dm.sampler_strategy is strategy - assert strategy.runtime.enabled is True - assert strategy.runtime.depth == 2 - assert strategy.runtime.num_neighbors == 8 - - -def test_datamodule_uses_val_sampler_override(tmp_path: Path): - dm = _minimal_datamodule( - tmp_path, - sampler={"name": "identity", "kwargs": {}}, - ) - dm.val_sampler_config = SamplerConfig.from_dict( - { - "name": "shadow_custom", - "kwargs": {}, - "runtime": { - "enabled": True, - "depth": 2, - "num_neighbors": 8, - "subgraph_batch_size": 4, - "proportional_root_sampling": False, - }, - } - ) - val_strategy = dm._ensure_val_sampler_strategy() - train_strategy = dm._ensure_sampler_strategy() - assert train_strategy.__class__.__name__ == "IdentityBatchStrategy" - assert val_strategy.__class__.__name__ == "ShadowCustomStrategy" - - -def test_datamodule_val_sampler_falls_back_to_train_sampler(tmp_path: Path): - dm = _minimal_datamodule(tmp_path, sampler={"name": "identity", "kwargs": {}}) - train_strategy = dm._ensure_sampler_strategy() - val_strategy = dm._ensure_val_sampler_strategy() - assert val_strategy is train_strategy diff --git a/tests/data/test_loader_flags.py b/tests/data/test_loader_flags.py deleted file mode 100644 index b2f6262..0000000 --- a/tests/data/test_loader_flags.py +++ /dev/null @@ -1,97 +0,0 @@ -from pathlib import Path - -import numpy as np -import pytest - -pytest.importorskip("lightning") - -from grass_mil.data.components.feature_reducers import FeatureReducerConfig -from grass_mil.data.components.graph_builders import GraphBuilderConfig -from grass_mil.data.components.loaders import CsvConfig, H5adConfig, SceConfig -from grass_mil.data.components.patching import TileConfig -from grass_mil.data.components.precompute import ( - CategoricalFeatureConfig, - ManifestConfig, - PrecomputeConfig, - SpatialOmicsPreprocessor, -) - - -def test_tsv_loader_preserves_use_molecular_features_flag(tmp_path: Path) -> None: - tsv_path = tmp_path / "cells.tsv" - np.savetxt( - tsv_path, - np.array( - [ - ["c0", "0.0", "1.0"], - ["c1", "2.0", "3.0"], - ], - dtype=str, - ), - delimiter="\t", - fmt="%s", - header="cell_id\tx\ty", - comments="", - ) - - preprocessor = SpatialOmicsPreprocessor( - manifest_config=ManifestConfig(), - precompute_config=PrecomputeConfig.from_args( - raw_manifest_path=str(tmp_path / "manifest.csv"), - processed_dir=str(tmp_path / "processed"), - coord_scale_um=1.0, - sample_unit="tile", - reducer_scope="sample", - reducer_fit_mode="train_only", - keep_raw_molecular=False, - force=False, - min_cells=1, - use_molecular_features=False, - split_by="patch", - split_ratios=(1.0, 0.0, 0.0), - split_seed=42, - ), - csv_config=CsvConfig( - coord_columns=("x", "y"), - cell_id_column="cell_id", - categorical_label_columns=[], - molecular_columns=None, - use_molecular_features=False, - sep=",", - ), - h5ad_config=H5adConfig( - coord_source="obsm", - coord_key="spatial", - coord_columns=("x", "y"), - cell_id_column=None, - categorical_label_columns=[], - molecular_layer=None, - molecular_features=None, - use_molecular_features=False, - ), - sce_config=SceConfig( - assay_name=None, - coord_source="colData", - coord_key=None, - coord_columns=("x", "y"), - cell_id_column=None, - categorical_label_columns=[], - molecular_features=None, - use_molecular_features=False, - ), - graph_builder_config=GraphBuilderConfig(name="delaunay", kwargs={}), - feature_reducer_config=FeatureReducerConfig(name="identity", kwargs={}), - categorical_feature_config=CategoricalFeatureConfig(include_labels=[]), - tile_config=TileConfig(tile_size_um=10.0, stride_um=10.0, min_cells=1), - ) - - row = { - "sample_id": "s1", - "input_path": str(tsv_path), - "input_type": "tsv", - "region_id": "r1", - "polygons_path": "", - } - table = preprocessor._load_table(row) - assert table.molecular_features is None - assert table.coords.shape == (2, 2) diff --git a/tests/data/test_precompute_manifest_regions.py b/tests/data/test_precompute_manifest_regions.py deleted file mode 100644 index 2c426d7..0000000 --- a/tests/data/test_precompute_manifest_regions.py +++ /dev/null @@ -1,39 +0,0 @@ -from types import SimpleNamespace - -import pandas as pd -import pytest - -pytest.importorskip("lightning") - -from grass_mil.data.components.precompute import ManifestConfig, SpatialOmicsPreprocessor - - -def _validator() -> SimpleNamespace: - return SimpleNamespace(manifest_config=ManifestConfig()) - - -def test_validate_manifest_regions_accepts_globally_unique_region_ids() -> None: - df = pd.DataFrame( - { - "sample_id": ["s0", "s1", "s2"], - "input_path": ["a.csv", "b.csv", "c.csv"], - "input_type": ["csv", "csv", "csv"], - "region_id": ["r0", "r1", "r2"], - } - ) - - SpatialOmicsPreprocessor._validate_manifest_regions(_validator(), df) - - -def test_validate_manifest_regions_rejects_cross_sample_region_id_reuse() -> None: - df = pd.DataFrame( - { - "sample_id": ["s0", "s1"], - "input_path": ["a.csv", "b.csv"], - "input_type": ["csv", "csv"], - "region_id": ["shared_region", "shared_region"], - } - ) - - with pytest.raises(ValueError, match="region_id values must be globally unique"): - SpatialOmicsPreprocessor._validate_manifest_regions(_validator(), df) diff --git a/tests/data/test_precompute_reducer_modes.py b/tests/data/test_precompute_reducer_modes.py deleted file mode 100644 index 443ee77..0000000 --- a/tests/data/test_precompute_reducer_modes.py +++ /dev/null @@ -1,149 +0,0 @@ -import numpy as np -import pytest - -pytest.importorskip("lightning") - -from grass_mil.data.components.feature_reducers import FeatureReducerConfig -from grass_mil.data.components.graph_builders import GraphBuilderConfig -from grass_mil.data.components.loaders import CsvConfig, H5adConfig, SceConfig -from grass_mil.data.components.patching import TileConfig -from grass_mil.data.components.precompute import ( - CategoricalFeatureConfig, - GraphUnitSpec, - ManifestConfig, - PrecomputeConfig, - PreparedSample, - SpatialOmicsPreprocessor, -) -from grass_mil.data.components.spatial_types import SpatialOmicsTable - - -def _make_preprocessor(fit_mode: str) -> SpatialOmicsPreprocessor: - return SpatialOmicsPreprocessor( - manifest_config=ManifestConfig(), - precompute_config=PrecomputeConfig.from_args( - raw_manifest_path="/tmp/unused_manifest.csv", - processed_dir="/tmp/unused_processed", - coord_scale_um=1.0, - sample_unit="tile", - reducer_scope="dataset", - reducer_fit_mode=fit_mode, - keep_raw_molecular=False, - force=False, - min_cells=1, - use_molecular_features=True, - split_by="patch", - split_ratios=(0.8, 0.1, 0.1), - split_seed=42, - ), - csv_config=CsvConfig( - coord_columns=("x", "y"), - cell_id_column="cell_id", - categorical_label_columns=[], - molecular_columns=None, - use_molecular_features=True, - ), - h5ad_config=H5adConfig( - coord_source="obsm", - coord_key="spatial", - coord_columns=("x", "y"), - cell_id_column=None, - categorical_label_columns=[], - molecular_layer=None, - molecular_features=None, - use_molecular_features=True, - ), - sce_config=SceConfig( - assay_name=None, - coord_source="colData", - coord_key=None, - coord_columns=("x", "y"), - cell_id_column=None, - categorical_label_columns=[], - use_molecular_features=True, - ), - graph_builder_config=GraphBuilderConfig(name="delaunay", kwargs={}), - feature_reducer_config=FeatureReducerConfig(name="pca", kwargs={"n_components": 1}), - categorical_feature_config=CategoricalFeatureConfig(include_labels=[]), - tile_config=TileConfig(tile_size_um=10.0, stride_um=10.0, min_cells=1), - ) - - -def _make_prepared_and_units(): - table_train = SpatialOmicsTable( - coords=np.zeros((3, 2), dtype=float), - molecular_features=np.array([[0.0, 1.0], [0.0, 1.0], [0.0, 1.0]], dtype=float), - molecular_feature_names=["a", "b"], - categorical_labels={}, - cell_ids=np.array(["a0", "a1", "a2"]), - sample_id="s0", - region_id="r0", - ) - table_test = SpatialOmicsTable( - coords=np.zeros((3, 2), dtype=float), - molecular_features=np.array([[10.0, 11.0], [10.0, 11.0], [10.0, 11.0]], dtype=float), - molecular_feature_names=["a", "b"], - categorical_labels={}, - cell_ids=np.array(["b0", "b1", "b2"]), - sample_id="s1", - region_id="r1", - ) - prepared = [ - PreparedSample( - table_idx=0, - table=table_train, - sample_id="s0", - region_id="r0", - patch_indices=[np.array([0, 1, 2])], - ), - PreparedSample( - table_idx=1, - table=table_test, - sample_id="s1", - region_id="r1", - patch_indices=[np.array([0, 1, 2])], - ), - ] - units = [ - GraphUnitSpec( - table_idx=0, - sample_id="s0", - region_id="r0", - patch_idx=0, - patch_id="s0_r0_patch_0", - indices=np.array([0, 1, 2]), - split="train", - ), - GraphUnitSpec( - table_idx=1, - sample_id="s1", - region_id="r1", - patch_idx=0, - patch_id="s1_r1_patch_0", - indices=np.array([0, 1, 2]), - split="test", - ), - ] - return prepared, units - - -def test_prepare_reducer_global_fits_on_all_data(): - preprocessor = _make_preprocessor("global") - prepared, units = _make_prepared_and_units() - state, reducer = preprocessor._prepare_reducer(prepared, units) - - assert state["fit_mode"] == "global" - assert state["fitted_on"] == "global" - assert "global" in state - assert np.isclose(float(reducer.mean_[0, 0]), 5.0) - - -def test_prepare_reducer_train_only_fits_on_train_split(): - preprocessor = _make_preprocessor("train_only") - prepared, units = _make_prepared_and_units() - state, reducer = preprocessor._prepare_reducer(prepared, units) - - assert state["fit_mode"] == "train_only" - assert state["fitted_on"] == "train_only" - assert "global" in state - assert np.isclose(float(reducer.mean_[0, 0]), 0.0) diff --git a/tests/data/test_samplers.py b/tests/data/test_samplers.py deleted file mode 100644 index e421308..0000000 --- a/tests/data/test_samplers.py +++ /dev/null @@ -1,782 +0,0 @@ -import pytest -import torch - - -def _toy_dataset(n_graphs: int = 3, n_nodes: int = 20): - from torch_geometric.data import Data - - graphs = [] - for _ in range(n_graphs): - edge_index = torch.randint(0, n_nodes, (2, 64), dtype=torch.long) - x = torch.randn(n_nodes, 8) - edge_attr = torch.rand(edge_index.size(1), 2) - graphs.append( - Data(x=x, edge_index=edge_index, edge_attr=edge_attr, num_nodes=n_nodes) - ) - return graphs - - -def _toy_dataset_with_graph_attrs(n_graphs: int = 2, n_nodes: int = 20): - from torch_geometric.data import Data - - graphs = [] - for i in range(n_graphs): - edge_index = torch.randint(0, n_nodes, (2, 64), dtype=torch.long) - x = torch.randn(n_nodes, 8) - edge_attr = torch.rand(edge_index.size(1), 2) - graph_y = torch.tensor([[float(i % 2)]], dtype=torch.float32) - graph_w = torch.ones_like(graph_y) - data = Data( - x=x, - edge_index=edge_index, - edge_attr=edge_attr, - graph_y=graph_y, - graph_w=graph_w, - num_nodes=n_nodes, - ) - data.sample_id = f"sample_{i}" - data.region_id = f"region_{i}" - data.patch_id = f"patch_{i}" - graphs.append(data) - return graphs - - -def _attach_cell_type_metadata(data, labels): - data.categorical_codes = torch.tensor(labels, dtype=torch.long).view(-1, 1) - data.categorical_slices = {"cell_type": 0} - return data - - -def test_sequence_assignment_preserves_schema_name_lists_without_nesting(): - from grass_mil.data.components.samplers import _select_sequence_value_for_subgraph - - # Typical Batch collation of list metadata from 2 subgraphs. - value = [["edge_type", "distance"], ["edge_type", "distance"]] - out = _select_sequence_value_for_subgraph( - key="edge_attr_names", - value=value, - subgraph_index=0, - num_subgraphs=2, - ) - assert out == ["edge_type", "distance"] - - -def test_sequence_assignment_does_not_slice_flat_schema_name_lists_by_length(): - from grass_mil.data.components.samplers import _select_sequence_value_for_subgraph - - # Coincidental length match with num_subgraphs should not imply per-subgraph slicing. - value = ["edge_type", "distance"] - out = _select_sequence_value_for_subgraph( - key="edge_attr_names", - value=value, - subgraph_index=0, - num_subgraphs=2, - ) - assert out == ["edge_type", "distance"] - - -def test_sequence_assignment_still_slices_per_subgraph_non_schema_lists(): - from grass_mil.data.components.samplers import _select_sequence_value_for_subgraph - - value = ["sample_0", "sample_1"] - out = _select_sequence_value_for_subgraph( - key="sample_id", - value=value, - subgraph_index=1, - num_subgraphs=2, - ) - assert out == ["sample_1"] - - -def test_identity_sampler_strategy(): - pytest.importorskip("torch_geometric") - from grass_mil.data.components.samplers import get_sampler_strategy - - dataset = _toy_dataset() - strategy = get_sampler_strategy({"name": "identity", "kwargs": {}}) - loader = strategy.build_dataset_loader( - dataset=dataset, - batch_size=2, - num_workers=0, - pin_memory=False, - shuffle=False, - ) - batch = next(iter(loader)) - assert hasattr(batch, "x") - assert hasattr(batch, "edge_index") - - -def test_shadow_native_requires_no_post_transform(): - pytest.importorskip("torch_geometric") - try: - import torch_sparse # noqa: F401 - except Exception: - pytest.skip("torch_sparse is not available") - - from grass_mil.data.components.samplers import get_sampler_strategy - - data = _toy_dataset(n_graphs=1)[0] - strategy = get_sampler_strategy({"name": "shadow_native", "kwargs": {}}) - with pytest.raises(ValueError): - strategy.build_unit_loader( - data=data, - depth=2, - num_neighbors=8, - batch_size=4, - transform=lambda x: x, - ) - - -def test_shadow_custom_unit_loader(): - pytest.importorskip("torch_geometric") - try: - import torch_sparse # noqa: F401 - except Exception: - pytest.skip("torch_sparse is not available") - - from grass_mil.data.components.samplers import get_sampler_strategy - - data = _toy_dataset(n_graphs=1)[0] - strategy = get_sampler_strategy({"name": "shadow_custom", "kwargs": {}}) - loader = strategy.build_unit_loader( - data=data, - depth=2, - num_neighbors=8, - batch_size=4, - transform=lambda d: d, - ) - batch = next(iter(loader)) - assert hasattr(batch, "x") - assert hasattr(batch, "edge_index") - - -def test_shadow_custom_without_torch_sparse_rejects_transform(monkeypatch): - pytest.importorskip("torch_geometric") - from grass_mil.data.components import samplers as samplers_mod - from grass_mil.data.components.samplers import ShadowCustomStrategy - - data = _toy_dataset(n_graphs=1)[0] - monkeypatch.setattr(samplers_mod, "WITH_TORCH_SPARSE", False) - strategy = ShadowCustomStrategy(runtime={"enabled": False}) - with pytest.raises(ValueError, match="requires torch-sparse"): - strategy.build_unit_loader( - data=data, - depth=2, - num_neighbors=8, - batch_size=4, - transform=lambda d: d, - ) - - -def test_shadow_custom_runtime_dataset_loader_uses_unit_loader(monkeypatch): - pytest.importorskip("torch_geometric") - from grass_mil.data.components.samplers import get_sampler_strategy - - dataset = _toy_dataset_with_graph_attrs(n_graphs=2) - strategy = get_sampler_strategy( - { - "name": "shadow_custom", - "kwargs": {}, - "runtime": { - "enabled": True, - "depth": 2, - "num_neighbors": 8, - "subgraph_batch_size": 4, - "replace": False, - "shuffle_subgraphs": True, - }, - } - ) - - call_counter = {"n": 0} - - def _fake_unit_loader(data, **kwargs): - call_counter["n"] += 1 - return [data] - - monkeypatch.setattr(strategy, "build_unit_loader", _fake_unit_loader) - - loader = strategy.build_dataset_loader( - dataset=dataset, - batch_size=2, - num_workers=0, - pin_memory=False, - shuffle=False, - ) - batches = list(loader) - assert call_counter["n"] == len(dataset) - assert len(batches) == len(dataset) - - -def test_shadow_runtime_normalizes_graph_level_metadata(monkeypatch): - pytest.importorskip("torch_geometric") - from torch_geometric.data import Batch - - from grass_mil.data.components.samplers import get_sampler_strategy - - dataset = _toy_dataset_with_graph_attrs(n_graphs=1) - strategy = get_sampler_strategy( - { - "name": "shadow_custom", - "kwargs": {}, - "runtime": {"enabled": True, "depth": 2, "num_neighbors": 8}, - } - ) - - def _fake_unit_loader(data, **kwargs): - b = Batch.from_data_list([data, data]) - # Simulate graph-level attrs carried as singletons before normalization. - b.sample_id = data.sample_id - b.region_id = data.region_id - b.patch_id = data.patch_id - b.graph_y = data.graph_y - b.graph_w = data.graph_w - return [b] - - monkeypatch.setattr(strategy, "build_unit_loader", _fake_unit_loader) - - loader = strategy.build_dataset_loader( - dataset=dataset, - batch_size=1, - num_workers=0, - pin_memory=False, - shuffle=False, - ) - batch = next(iter(loader)) - assert isinstance(batch.sample_id, list) and len(batch.sample_id) == 2 - assert isinstance(batch.region_id, list) and len(batch.region_id) == 2 - assert batch.graph_y.size(0) == 2 - assert batch.graph_w.size(0) == 2 - - -def test_runtime_loader_len_matches_yielded_batches(monkeypatch): - pytest.importorskip("torch_geometric") - from grass_mil.data.components.samplers import get_sampler_strategy - - dataset = _toy_dataset_with_graph_attrs(n_graphs=2, n_nodes=6) - strategy = get_sampler_strategy( - { - "name": "shadow_custom", - "kwargs": {}, - "runtime": { - "enabled": True, - "depth": 2, - "num_neighbors": 8, - "subgraph_batch_size": 4, - "proportional_root_sampling": False, - }, - } - ) - - def _fake_unit_loader(data, **kwargs): - roots = kwargs.get("node_idx") - if roots is None: - n_roots = int(data.num_nodes) - else: - n_roots = int(roots.numel()) - bs = int(kwargs["batch_size"]) - n_batches = (n_roots + bs - 1) // bs - return [data] * n_batches - - monkeypatch.setattr(strategy, "build_unit_loader", _fake_unit_loader) - loader = strategy.build_dataset_loader( - dataset=dataset, - batch_size=1, - num_workers=0, - pin_memory=False, - shuffle=False, - ) - yielded = sum(1 for _ in loader) - assert len(loader) == yielded - - -def test_runtime_shadow_config_proportional_defaults(): - from grass_mil.data.components.samplers import RuntimeShadowConfig - - cfg = RuntimeShadowConfig.from_dict({}) - assert cfg.proportional_root_sampling is True - assert cfg.property_name == "cell_type" - assert cfg.weight_mode == "inverse" - assert cfg.min_weight > 0 - - -def test_shadow_runtime_passes_weighted_node_idx(monkeypatch): - pytest.importorskip("torch_geometric") - from grass_mil.data.components.samplers import get_sampler_strategy - - unit = _toy_dataset_with_graph_attrs(n_graphs=1, n_nodes=6)[0] - unit = _attach_cell_type_metadata(unit, [0, 0, 0, 0, 1, 1]) - dataset = [unit] - strategy = get_sampler_strategy( - { - "name": "shadow_custom", - "kwargs": {}, - "runtime": { - "enabled": True, - "depth": 2, - "num_neighbors": 8, - "subgraph_batch_size": 4, - "proportional_root_sampling": True, - "property_name": "cell_type", - "weight_mode": "inverse", - }, - } - ) - - observed = {"node_idx": None} - - def _fake_unit_loader(data, **kwargs): - observed["node_idx"] = kwargs.get("node_idx") - return [data] - - monkeypatch.setattr(strategy, "build_unit_loader", _fake_unit_loader) - loader = strategy.build_dataset_loader( - dataset=dataset, - batch_size=1, - num_workers=0, - pin_memory=False, - shuffle=False, - ) - _ = list(loader) - node_idx = observed["node_idx"] - assert isinstance(node_idx, torch.Tensor) - assert node_idx.dtype == torch.long - assert node_idx.numel() == unit.num_nodes - - -def test_shadow_runtime_falls_back_to_none_when_property_missing(monkeypatch): - pytest.importorskip("torch_geometric") - from grass_mil.data.components.samplers import get_sampler_strategy - - unit = _toy_dataset_with_graph_attrs(n_graphs=1, n_nodes=6)[0] - dataset = [unit] - strategy = get_sampler_strategy( - { - "name": "shadow_custom", - "kwargs": {}, - "runtime": { - "enabled": True, - "depth": 2, - "num_neighbors": 8, - "subgraph_batch_size": 4, - "proportional_root_sampling": True, - "property_name": "cell_type", - "weight_mode": "inverse", - "node_idx": None, - }, - } - ) - - observed = {"node_idx": "unset", "unit_num_nodes": int(unit.num_nodes)} - - def _fake_unit_loader(data, **kwargs): - observed["node_idx"] = kwargs.get("node_idx") - return [data] - - monkeypatch.setattr(strategy, "build_unit_loader", _fake_unit_loader) - loader = strategy.build_dataset_loader( - dataset=dataset, - batch_size=1, - num_workers=0, - pin_memory=False, - shuffle=False, - ) - _ = list(loader) - node_idx = observed["node_idx"] - assert isinstance(node_idx, torch.Tensor) - assert node_idx.dtype == torch.long - assert node_idx.numel() == observed["unit_num_nodes"] - assert torch.equal( - node_idx, torch.arange(observed["unit_num_nodes"], dtype=torch.long) - ) - - -def test_shadow_runtime_can_disable_proportional_sampling(monkeypatch): - pytest.importorskip("torch_geometric") - from grass_mil.data.components.samplers import get_sampler_strategy - - unit = _toy_dataset_with_graph_attrs(n_graphs=1, n_nodes=6)[0] - unit = _attach_cell_type_metadata(unit, [0, 0, 0, 0, 1, 1]) - dataset = [unit] - explicit_node_idx = torch.tensor([0, 1], dtype=torch.long) - strategy = get_sampler_strategy( - { - "name": "shadow_custom", - "kwargs": {}, - "runtime": { - "enabled": True, - "depth": 2, - "num_neighbors": 8, - "subgraph_batch_size": 4, - "proportional_root_sampling": False, - "property_name": "cell_type", - "weight_mode": "inverse", - "node_idx": explicit_node_idx, - }, - } - ) - - observed = {"node_idx": None} - - def _fake_unit_loader(data, **kwargs): - observed["node_idx"] = kwargs.get("node_idx") - return [data] - - monkeypatch.setattr(strategy, "build_unit_loader", _fake_unit_loader) - loader = strategy.build_dataset_loader( - dataset=dataset, - batch_size=1, - num_workers=0, - pin_memory=False, - shuffle=False, - ) - _ = list(loader) - assert torch.equal(observed["node_idx"], explicit_node_idx) - - -@pytest.mark.parametrize("weight_mode", ["inverse", "sqrt_inverse", "proportional"]) -def test_shadow_runtime_accepts_weight_modes(weight_mode): - from grass_mil.data.components.samplers import RuntimeShadowConfig - - cfg = RuntimeShadowConfig.from_dict({"weight_mode": weight_mode}) - assert cfg.weight_mode == weight_mode - - -def test_runtime_shadow_config_accepts_subsample_defaults(): - from grass_mil.data.components.samplers import RuntimeShadowConfig - - cfg = RuntimeShadowConfig.from_dict({}) - assert cfg.subsample_fraction == 1.0 - assert cfg.subsample_seed is None - - -@pytest.mark.parametrize("fraction", [0.0, -0.1, 1.1]) -def test_runtime_shadow_config_rejects_invalid_subsample_fraction(fraction): - from grass_mil.data.components.samplers import RuntimeShadowConfig - - with pytest.raises(ValueError, match="subsample_fraction"): - RuntimeShadowConfig.from_dict({"subsample_fraction": fraction}) - - -def test_shadow_runtime_subsamples_roots_with_seed(monkeypatch): - pytest.importorskip("torch_geometric") - from grass_mil.data.components.samplers import get_sampler_strategy - - unit = _toy_dataset_with_graph_attrs(n_graphs=1, n_nodes=10)[0] - dataset = [unit] - strategy = get_sampler_strategy( - { - "name": "shadow_custom", - "kwargs": {}, - "runtime": { - "enabled": True, - "depth": 2, - "num_neighbors": 8, - "subgraph_batch_size": 4, - "proportional_root_sampling": False, - "subsample_fraction": 0.4, - "subsample_seed": 17, - }, - } - ) - - observed = {"node_idx": None} - - def _fake_unit_loader(data, **kwargs): - observed["node_idx"] = kwargs.get("node_idx") - return [data] - - monkeypatch.setattr(strategy, "build_unit_loader", _fake_unit_loader) - loader = strategy.build_dataset_loader( - dataset=dataset, - batch_size=1, - num_workers=0, - pin_memory=False, - shuffle=False, - ) - _ = list(loader) - first = observed["node_idx"] - assert isinstance(first, torch.Tensor) - assert first.dtype == torch.long - assert int(first.numel()) == 4 - - observed["node_idx"] = None - _ = list(loader) - second = observed["node_idx"] - assert isinstance(second, torch.Tensor) - assert torch.equal(first, second) - - -def test_shadow_runtime_subsampled_len_matches_yielded_batches(monkeypatch): - pytest.importorskip("torch_geometric") - from grass_mil.data.components.samplers import get_sampler_strategy - - dataset = _toy_dataset_with_graph_attrs(n_graphs=2, n_nodes=10) - strategy = get_sampler_strategy( - { - "name": "shadow_custom", - "kwargs": {}, - "runtime": { - "enabled": True, - "depth": 2, - "num_neighbors": 8, - "subgraph_batch_size": 3, - "proportional_root_sampling": False, - "subsample_fraction": 0.5, - "subsample_seed": 11, - }, - } - ) - - def _fake_unit_loader(data, **kwargs): - roots = kwargs.get("node_idx") - n_roots = int(roots.numel()) if roots is not None else int(data.num_nodes) - bs = int(kwargs["batch_size"]) - n_batches = (n_roots + bs - 1) // bs - return [data] * n_batches - - monkeypatch.setattr(strategy, "build_unit_loader", _fake_unit_loader) - loader = strategy.build_dataset_loader( - dataset=dataset, - batch_size=1, - num_workers=0, - pin_memory=False, - shuffle=False, - ) - yielded = sum(1 for _ in loader) - assert len(loader) == yielded - - -def test_shadow_custom_per_subgraph_transform_preserves_node_level_tensors(): - """Node-level tensors must survive per-subgraph transform application. - - ``categorical_codes`` carries the cell-type codes the encoder embeds and - ``pos`` carries coordinates; neither is recoverable from ``x``. They were - previously dropped as "already represented by x", which silently removed - cell type from every sampled subgraph. - """ - pytest.importorskip("torch_geometric") - try: - import torch_sparse # noqa: F401 - except Exception: - pytest.skip("torch_sparse is not available") - - from grass_mil.data.components.samplers import get_sampler_strategy - - n_nodes = 20 - data = _toy_dataset(n_graphs=1, n_nodes=n_nodes)[0] - _attach_cell_type_metadata(data, list(range(n_nodes))) - data.pos = torch.arange(n_nodes, dtype=torch.float32).view(-1, 1).repeat(1, 2) - - strategy = get_sampler_strategy({"name": "shadow_custom", "kwargs": {}}) - loader = strategy.build_unit_loader( - data=data, - depth=2, - num_neighbors=8, - batch_size=4, - transform=lambda d: d, - ) - batch = next(iter(loader)) - - assert hasattr(batch, "categorical_codes"), "cell-type codes were dropped by sampling" - assert hasattr(batch, "pos"), "coordinates were dropped by sampling" - assert batch.categorical_codes.size(0) == batch.x.size(0) - assert batch.pos.size(0) == batch.x.size(0) - - # The codes must be the ones belonging to the sampled nodes, not a stale - # or truncated slice: node i of the source graph carries code i. - expected = data.categorical_codes[batch.n_id].view(-1) - assert torch.equal(batch.categorical_codes.view(-1), expected) - - -def test_node_level_codes_batch_along_the_node_dimension(): - """Guard the attribute name against PyTorch Geometric's index heuristic. - - ``Data.__cat_dim__`` returns -1 for any key containing ``index``, so an - attribute named ``categorical_index`` is concatenated along the feature - dimension when graphs are batched -- silently corrupting per-node codes (or - raising, when the graphs differ in size). The codes must batch along dim 0 - like any other node attribute. - """ - pytest.importorskip("torch_geometric") - from torch_geometric.data import Batch, Data - - graphs = [] - for n_nodes in (4, 7): - data = Data(x=torch.randn(n_nodes, 3), edge_index=torch.empty(2, 0, dtype=torch.long)) - data.num_nodes = n_nodes - _attach_cell_type_metadata(data, list(range(n_nodes))) - graphs.append(data) - - batch = Batch.from_data_list(graphs) - assert batch.categorical_codes.shape == (11, 1) - assert torch.equal( - batch.categorical_codes.view(-1), - torch.tensor([0, 1, 2, 3, 0, 1, 2, 3, 4, 5, 6]), - ) - - -def _dense_grid_graph(side: int = 9, spacing: float = 10.0): - """Grid graph with coordinates, for interior/radius sampler wiring.""" - import numpy as np - from torch_geometric.data import Data - - xs, ys = np.meshgrid(np.arange(side), np.arange(side), indexing="ij") - coords = np.column_stack([xs.ravel(), ys.ravel()]).astype(float) * spacing - - edges = [] - for i in range(side): - for j in range(side): - node = i * side + j - if i + 1 < side: - edges.append((node, (i + 1) * side + j)) - if j + 1 < side: - edges.append((node, i * side + (j + 1))) - edge_index = torch.tensor(np.array(edges).T, dtype=torch.long) - n_nodes = side * side - data = Data( - x=torch.randn(n_nodes, 4), - edge_index=edge_index, - edge_attr=torch.rand(edge_index.size(1), 2), - ) - data.pos = torch.tensor(coords, dtype=torch.float32) - data.num_nodes = n_nodes - data.sample_id = "s0" - data.region_id = "r0" - data.patch_id = "p0" - return data - - -def test_runtime_interior_seeds_restrict_roots_to_the_tissue_interior(): - pytest.importorskip("torch_geometric") - from grass_mil.data.components.samplers import ( - RuntimeShadowConfig, - _resolve_interior_indices, - ) - - data = _dense_grid_graph(side=9) - runtime = RuntimeShadowConfig.from_dict( - {"enabled": True, "depth": 1, "interior_seeds": {"enabled": True}} - ) - interior = _resolve_interior_indices(data, runtime) - assert interior is not None - assert int(interior.numel()) < data.num_nodes - - # No seed may sit on the outer ring of the grid. - coords = data.pos.numpy() - selected = coords[interior.numpy()] - assert selected[:, 0].min() > coords[:, 0].min() - assert selected[:, 0].max() < coords[:, 0].max() - - disabled = RuntimeShadowConfig.from_dict({"enabled": True, "depth": 1}) - assert _resolve_interior_indices(data, disabled) is None - - -def test_runtime_ego_radius_shrinks_sampled_subgraphs(): - pytest.importorskip("torch_geometric") - try: - import torch_sparse # noqa: F401 - except Exception: - pytest.skip("torch_sparse is not available") - - from grass_mil.data.components.samplers import get_sampler_strategy - - data = _dense_grid_graph(side=9, spacing=10.0) - - def _subgraph_sizes(runtime): - strategy = get_sampler_strategy( - {"name": "shadow_custom", "kwargs": {}, "runtime": runtime} - ) - loader = strategy.build_dataset_loader( - dataset=[data], - batch_size=1, - num_workers=0, - pin_memory=False, - shuffle=False, - transform=lambda d: d, - ) - batch = next(iter(loader)) - return int(batch.x.size(0)) - - base_runtime = { - "enabled": True, - "depth": 3, - "num_neighbors": 8, - "subgraph_batch_size": 4, - "shuffle_subgraphs": False, - "proportional_root_sampling": False, - } - unbounded = _subgraph_sizes(base_runtime) - bounded = _subgraph_sizes( - { - **base_runtime, - # r = 5*3 + 0 = 15, i.e. 1.5 grid steps: anything two or more - # steps from the root is cut even though it is within 3 hops. - "ego_radius": { - "enabled": True, - "radius_per_hop": 5.0, - "radius_offset": 0.0, - }, - } - ) - assert bounded < unbounded - - -def test_shadow_batches_split_into_subgraphs_without_a_transform(): - """ShaDow batches must be splittable even when no transform is configured. - - ``Batch.to_data_list()`` only works for batches assembled by - ``Batch.from_data_list()``. The ShaDow samplers construct their ``Batch`` - directly and delimit subgraphs with ``ptr``, so ``to_data_list()`` raises on - them. - - That mattered because ``predict_step`` used ``to_data_list()`` behind a bare - ``except`` to recover per-instance subgraphs: it silently produced ``None``, - ``instance_graphs`` were never collected, and the spatial-table export failed - several steps later complaining about a missing payload field. Every spatial - plugin (Moran's I, Ripley, neighbourhood enrichment) was unreachable from the - CLI as a result. - - The existing coverage missed it because the only test exercising this path - passed a transform, and the transform branch re-collates via - ``Batch.from_data_list()`` -- the one case where ``to_data_list()`` works. - """ - pytest.importorskip("torch_geometric") - try: - import torch_sparse # noqa: F401 - except Exception: - pytest.skip("torch_sparse is not available") - - from grass_mil.batching import extract_subgraphs - from grass_mil.data.components.samplers import get_sampler_strategy - - n_nodes = 20 - data = _toy_dataset(n_graphs=1, n_nodes=n_nodes)[0] - _attach_cell_type_metadata(data, list(range(n_nodes))) - data.pos = torch.arange(n_nodes, dtype=torch.float32).view(-1, 1).repeat(1, 2) - - strategy = get_sampler_strategy({"name": "shadow_custom", "kwargs": {}}) - loader = strategy.build_unit_loader( - data=data, depth=2, num_neighbors=8, batch_size=4, transform=None - ) - batch = next(iter(loader)) - - # The precondition: this batch is not one PyG can un-batch itself. - with pytest.raises(Exception): - batch.to_data_list() - - subgraphs = extract_subgraphs(batch) - assert subgraphs is not None, "ShaDow batch could not be split into subgraphs" - assert len(subgraphs) == int(batch.ptr.numel() - 1) - - for sub in subgraphs: - # Node-level tensors must come along, sliced to this subgraph. Spatial - # export needs `pos`; the encoder needs `categorical_codes`. - assert sub.categorical_codes.size(0) == sub.x.size(0) - assert sub.pos.size(0) == sub.x.size(0) - # Edges must be relabelled into the subgraph's own index space. - if sub.edge_index.numel(): - assert int(sub.edge_index.max()) < sub.x.size(0) - - # Node counts must partition the batch, with nothing dropped or duplicated. - assert sum(int(s.x.size(0)) for s in subgraphs) == int(batch.x.size(0)) diff --git a/tests/data/test_seed_sampling.py b/tests/data/test_seed_sampling.py deleted file mode 100644 index 8ae76b0..0000000 --- a/tests/data/test_seed_sampling.py +++ /dev/null @@ -1,223 +0,0 @@ -"""Contracts for interior seed restriction and the ego-radius cutoff.""" - -from __future__ import annotations - -import numpy as np -import pytest -import torch - -from grass_mil.data.components.ego_radius import ( - EgoRadiusConfig, - apply_ego_radius_cutoff, - resolve_ego_radius, -) -from grass_mil.data.components.seed_sampling import ( - InteriorSeedConfig, - InteriorSeedIndexCache, - compute_interior_seed_indices, - convex_hull_interior_mask, -) - - -def _grid(side: int = 9, spacing: float = 1.0): - """A regular square grid with 4-neighbour connectivity.""" - xs, ys = np.meshgrid(np.arange(side), np.arange(side), indexing="ij") - coords = np.column_stack([xs.ravel(), ys.ravel()]).astype(float) * spacing - - edges = [] - for i in range(side): - for j in range(side): - node = i * side + j - if i + 1 < side: - edges.append((node, (i + 1) * side + j)) - if j + 1 < side: - edges.append((node, i * side + (j + 1))) - edge_index = np.array(edges, dtype=np.int64).T - return coords, edge_index - - -def test_interior_mask_excludes_the_hull_rim() -> None: - side = 9 - coords, edge_index = _grid(side) - - # A tolerance below the grid spacing makes only the cells lying exactly on - # the hull polygon seeds, so one hop removes exactly the outer ring. - mask = convex_hull_interior_mask( - coords, edge_index, n_hops=1, seed_tolerance_multiplier=0.5 - ) - assert mask.sum() == (side - 2) ** 2 - - grid_mask = mask.reshape(side, side) - assert not grid_mask[0].any() and not grid_mask[-1].any() - assert not grid_mask[:, 0].any() and not grid_mask[:, -1].any() - assert grid_mask[1:-1, 1:-1].all() - - -def test_seed_tolerance_widens_the_rim_by_proximity_to_the_hull() -> None: - """Seeds are cells *near* the hull, not only those exactly on it. - - The tolerance is a multiple of the median nearest-neighbour distance, so at - the default multiplier a unit-spaced grid seeds its two outermost rings. - """ - side = 9 - coords, edge_index = _grid(side) - - tight = convex_hull_interior_mask( - coords, edge_index, n_hops=1, seed_tolerance_multiplier=0.5 - ) - default = convex_hull_interior_mask(coords, edge_index, n_hops=1) - - assert tight.sum() == (side - 2) ** 2 - assert default.sum() == (side - 4) ** 2 - # The wider tolerance removes a strict superset of cells. - assert not (default & ~tight).any() - - -def test_interior_shrinks_monotonically_with_hop_count() -> None: - coords, edge_index = _grid(11) - sizes = [ - int(convex_hull_interior_mask(coords, edge_index, n_hops=hops).sum()) - for hops in (1, 2, 3) - ] - assert sizes[0] > sizes[1] > sizes[2] - - -def test_interior_mask_falls_back_to_all_nodes_on_degenerate_input() -> None: - # Collinear points have no area hull, and tiny graphs have no meaningful rim. - collinear = np.array([[0.0, 0.0], [1.0, 0.0], [2.0, 0.0], [3.0, 0.0]]) - edge_index = np.array([[0, 1, 2], [1, 2, 3]], dtype=np.int64) - assert convex_hull_interior_mask(collinear, edge_index, n_hops=2).all() - - tiny = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]]) - assert convex_hull_interior_mask(tiny, np.empty((2, 0), dtype=np.int64), n_hops=2).all() - - # A non-positive rim width disables the restriction entirely. - coords, edges = _grid(6) - assert convex_hull_interior_mask(coords, edges, n_hops=0).all() - - -def test_interior_mask_can_drop_disconnected_islands() -> None: - """An interior island cut off from the main body is folded into the rim.""" - coords, edge_index = _grid(9) - # Sever the centre cell from its neighbours: it survives the rim test but is - # no longer connected to the interior body. - centre = 4 * 9 + 4 - keep = ~((edge_index[0] == centre) | (edge_index[1] == centre)) - severed = edge_index[:, keep] - - with_islands = convex_hull_interior_mask( - coords, severed, n_hops=1, keep_largest_component=False - ) - without_islands = convex_hull_interior_mask( - coords, severed, n_hops=1, keep_largest_component=True - ) - assert with_islands[centre] - assert not without_islands[centre] - assert without_islands.sum() == with_islands.sum() - 1 - - -def test_compute_interior_seed_indices_falls_back_without_coordinates() -> None: - from torch_geometric.data import Data - - data = Data(x=torch.randn(5, 2), edge_index=torch.empty(2, 0, dtype=torch.long)) - data.num_nodes = 5 - indices = compute_interior_seed_indices( - data, InteriorSeedConfig(enabled=True), n_hops=2 - ) - assert indices.tolist() == list(range(5)) - - -def test_interior_seed_cache_round_trip_and_corruption_recovery(tmp_path) -> None: - cache = InteriorSeedIndexCache(str(tmp_path)) - indices = np.array([1, 2, 3], dtype=np.int64) - cache.store("region_a__hops2", indices) - np.testing.assert_array_equal(cache.load("region_a__hops2"), indices) - - assert cache.load("missing") is None - - # A truncated cache entry must trigger recomputation, not an exception. - corrupt = tmp_path / "region_a__hops2.npy" - corrupt.write_bytes(b"not a numpy file") - assert cache.load("region_a__hops2") is None - - -def test_disabled_cache_is_a_no_op() -> None: - cache = InteriorSeedIndexCache(None) - cache.store("anything", np.array([1])) - assert cache.load("anything") is None - - -def test_resolve_ego_radius_is_linear_in_depth() -> None: - # The published NSCLC setting: r_k = 75k + 55, so r_4 = 355. - assert resolve_ego_radius(4, radius_per_hop=75.0, radius_offset=55.0) == 355.0 - assert resolve_ego_radius(0, radius_per_hop=75.0, radius_offset=55.0) == 55.0 - with pytest.raises(ValueError): - resolve_ego_radius(-1, radius_per_hop=75.0, radius_offset=55.0) - - -def _line_graph(n: int = 6, spacing: float = 10.0): - from torch_geometric.data import Data - - pos = torch.arange(n, dtype=torch.float32).unsqueeze(-1) * spacing - pos = torch.cat([pos, torch.zeros_like(pos)], dim=1) - src = torch.arange(n - 1, dtype=torch.long) - edge_index = torch.stack([src, src + 1], dim=0) - data = Data( - x=torch.randn(n, 3), - edge_index=edge_index, - edge_attr=torch.rand(edge_index.size(1), 2), - ) - data.pos = pos - data.num_nodes = n - data.root_n_id = torch.tensor([0], dtype=torch.long) - data.categorical_codes = torch.arange(n, dtype=torch.long).view(-1, 1) - return data - - -def test_ego_radius_cutoff_drops_distant_nodes_and_keeps_the_root() -> None: - data = _line_graph(n=6, spacing=10.0) - out = apply_ego_radius_cutoff(data, radius=25.0) - - # Nodes at 0, 10, 20 are within the radius; 30, 40, 50 are not. - assert int(out.num_nodes) == 3 - assert int(out.root_n_id.item()) == 0 - torch.testing.assert_close(out.pos[:, 0], torch.tensor([0.0, 10.0, 20.0])) - # Node-level tensors are sliced consistently with x. - assert out.categorical_codes.view(-1).tolist() == [0, 1, 2] - assert out.edge_index.max() < out.num_nodes - - -def test_ego_radius_cutoff_is_a_no_op_when_everything_fits() -> None: - data = _line_graph(n=4, spacing=1.0) - out = apply_ego_radius_cutoff(data, radius=1000.0) - assert int(out.num_nodes) == 4 - assert out is data - - -def test_ego_radius_cutoff_can_drop_nodes_disconnected_from_the_root() -> None: - """Within the radius but unreachable: an island, not a neighbourhood.""" - from torch_geometric.data import Data - - pos = torch.tensor([[0.0, 0.0], [1.0, 0.0], [2.0, 0.0]]) - # Node 2 is close to the root but only connected to node 1... which we omit. - edge_index = torch.tensor([[0], [1]], dtype=torch.long) - data = Data(x=torch.randn(3, 2), edge_index=edge_index) - data.pos = pos - data.num_nodes = 3 - data.root_n_id = torch.tensor([0], dtype=torch.long) - - kept = apply_ego_radius_cutoff(data, radius=10.0, keep_root_component=False) - assert int(kept.num_nodes) == 3 - - pruned = apply_ego_radius_cutoff(data, radius=10.0, keep_root_component=True) - assert int(pruned.num_nodes) == 2 - - -def test_ego_radius_config_rejects_unknown_keys() -> None: - with pytest.raises(ValueError, match="Unknown ego radius keys"): - EgoRadiusConfig.from_dict({"enabled": True, "radius": 10.0}) - - -def test_interior_seed_config_rejects_unknown_keys() -> None: - with pytest.raises(ValueError, match="Unknown interior seed keys"): - InteriorSeedConfig.from_dict({"enabled": True, "hops": 2}) diff --git a/tests/eval/test_aggregation.py b/tests/eval/test_aggregation.py deleted file mode 100644 index 366ed5d..0000000 --- a/tests/eval/test_aggregation.py +++ /dev/null @@ -1,328 +0,0 @@ -from __future__ import annotations - -import torch -import pytest - -from grass_mil.inference.aggregation import aggregate_group_logits -from grass_mil.inference.schemas import BatchPredictionPayload - - -def test_aggregate_group_logits_mean() -> None: - payload = BatchPredictionPayload( - bag_ids=["r1", "r1", "r2"], - bag_logits=torch.tensor([[1.0], [3.0], [2.0]]), - bag_targets=torch.tensor([[1.0], [1.0], [0.0]]), - bag_attention=[torch.tensor([0.1, 0.9]), torch.tensor([0.9, 0.1]), None], - row_region_ids=["region_a", "region_a", "region_b"], - row_sample_ids=["sample_x", "sample_x", "sample_y"], - ) - out = aggregate_group_logits(payload, mode="mean") - assert out.bag_ids == ["r1", "r2"] - assert torch.allclose(out.bag_logits[0], torch.tensor([2.0])) - assert out.bag_attention is None - - -def test_aggregate_group_logits_max() -> None: - payload = BatchPredictionPayload( - bag_ids=["r1", "r1", "r2"], - bag_logits=torch.tensor([[1.0], [3.0], [2.0]]), - bag_targets=None, - bag_attention=None, - row_region_ids=["region_a", "region_a", "region_b"], - row_sample_ids=["sample_x", "sample_x", "sample_y"], - ) - out = aggregate_group_logits(payload, mode="max") - assert out.bag_ids == ["r1", "r2"] - assert torch.allclose(out.bag_logits[0], torch.tensor([3.0])) - - -def test_aggregate_group_logits_mean_drops_non_aggregable_attention_rows() -> None: - payload = BatchPredictionPayload( - bag_ids=["r1", "r1"], - bag_logits=torch.tensor([[1.0], [2.0]]), - bag_targets=None, - bag_attention=[torch.tensor([0.1, 0.9]), torch.tensor([0.8, 0.2])], - row_region_ids=["region_a", "region_a"], - row_sample_ids=["sample_x", "sample_x"], - ) - out = aggregate_group_logits(payload, mode="mean") - assert out.bag_ids == ["r1"] - assert out.bag_attention is None - - -def test_aggregate_group_logits_attention_weighted_uses_instance_softmax() -> None: - """Per-class attention: identical columns reproduce a shared weighting.""" - payload = BatchPredictionPayload( - bag_ids=["r1", "r1", "r2"], - bag_logits=torch.tensor([[1.0, 10.0], [3.0, 30.0], [2.0, 20.0]]), - bag_targets=None, - bag_attention=None, - row_region_ids=["region_a", "region_a", "region_b"], - row_sample_ids=["sample_x", "sample_x", "sample_y"], - instance_logits=torch.tensor([[1.0, 10.0], [3.0, 30.0], [2.0, 20.0]]), - instance_attention_logits=torch.tensor([[1.0, 1.0], [9.0, 9.0], [1.0, 1.0]]), - instance_patch_ids=["r1", "r1", "r2"], - instance_region_ids=["region_a", "region_a", "region_b"], - instance_sample_ids=["sample_x", "sample_x", "sample_y"], - ) - out = aggregate_group_logits(payload, mode="attention_weighted") - assert out.bag_ids == ["r1", "r2"] - assert torch.allclose(out.bag_logits[0], torch.tensor([2.9993, 29.9933]), atol=1e-3) - assert torch.allclose(out.bag_logits[1], torch.tensor([2.0, 20.0]), atol=1e-6) - assert out.bag_attention is not None - assert len(out.bag_attention) == 2 - assert out.bag_attention[0] is not None - # Attention is [n_instances, n_classes], one column per class. - assert out.bag_attention[0].shape == (2, 2) - assert torch.allclose( - out.bag_attention[0], - torch.tensor([[0.000335, 0.000335], [0.999665, 0.999665]]), - atol=1e-4, - ) - - -def test_aggregate_group_logits_attention_weighted_supports_per_class_attention() -> None: - """Each class column is normalised over instances independently.""" - payload = BatchPredictionPayload( - bag_ids=["r1", "r1"], - bag_logits=torch.tensor([[1.0, 10.0], [3.0, 30.0]]), - bag_targets=None, - bag_attention=None, - row_region_ids=["region_a", "region_a"], - row_sample_ids=["sample_x", "sample_x"], - instance_logits=torch.tensor([[1.0, 10.0], [3.0, 30.0]]), - # Class 0 attends the first instance, class 1 the second. - instance_attention_logits=torch.tensor([[9.0, 1.0], [1.0, 9.0]]), - instance_patch_ids=["r1", "r1"], - instance_region_ids=["region_a", "region_a"], - instance_sample_ids=["sample_x", "sample_x"], - ) - out = aggregate_group_logits(payload, mode="attention_weighted") - - attention = out.bag_attention[0] - assert attention.shape == (2, 2) - assert torch.allclose(attention.sum(dim=0), torch.ones(2), atol=1e-6) - - # The identity: bag logit for class c is sum_i A[i,c] * l[i,c]. - expected = (payload.instance_logits * attention).sum(dim=0) - assert torch.allclose(out.bag_logits[0], expected, atol=1e-6) - # Class 0 is pulled toward instance 0's logit, class 1 toward instance 1's. - assert out.bag_logits[0][0] < 1.5 - assert out.bag_logits[0][1] > 29.0 - - -def test_aggregate_group_logits_rejects_mismatched_attention_width() -> None: - payload = BatchPredictionPayload( - bag_ids=["r1", "r1"], - bag_logits=torch.tensor([[1.0, 10.0], [3.0, 30.0]]), - bag_targets=None, - bag_attention=None, - row_region_ids=["region_a", "region_a"], - row_sample_ids=["sample_x", "sample_x"], - instance_logits=torch.tensor([[1.0, 10.0], [3.0, 30.0]]), - instance_attention_logits=torch.tensor([[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]]), - instance_patch_ids=["r1", "r1"], - instance_region_ids=["region_a", "region_a"], - instance_sample_ids=["sample_x", "sample_x"], - ) - with pytest.raises(ValueError, match="Attention emits 3 channel"): - aggregate_group_logits(payload, mode="attention_weighted") - - -def test_aggregate_group_logits_mean_supports_region_scope() -> None: - payload = BatchPredictionPayload( - bag_ids=["p0", "p1", "p2"], - bag_logits=torch.tensor([[1.0], [3.0], [8.0]]), - bag_targets=None, - bag_attention=None, - row_region_ids=["rA", "rA", "rB"], - row_sample_ids=["s0", "s0", "s1"], - ) - out = aggregate_group_logits(payload, mode="mean", bag_scope="region") - assert out.bag_ids == ["rA", "rB"] - assert torch.allclose(out.bag_logits[0], torch.tensor([2.0])) - assert torch.allclose(out.bag_logits[1], torch.tensor([8.0])) - - -def test_aggregate_group_logits_region_mean_uses_instance_logits() -> None: - payload = BatchPredictionPayload( - bag_ids=["chunk_a", "chunk_b", "chunk_c"], - bag_logits=torch.tensor([[5.0], [2.0], [3.0]]), - bag_targets=torch.tensor([[1.0], [1.0], [0.0]]), - bag_attention=None, - row_region_ids=["rA", "rA", "rB"], - row_sample_ids=["sA", "sA", "sB"], - instance_logits=torch.tensor([[1.0], [9.0], [2.0], [3.0]]), - instance_attention_logits=None, - instance_patch_ids=["chunk_a", "chunk_a", "chunk_b", "chunk_c"], - instance_region_ids=["rA", "rA", "rA", "rB"], - instance_sample_ids=["sA", "sA", "sA", "sB"], - ) - out = aggregate_group_logits(payload, mode="mean", bag_scope="region") - assert out.bag_ids == ["rA", "rB"] - assert torch.allclose(out.bag_logits[0], torch.tensor([4.0])) - assert torch.allclose(out.bag_logits[1], torch.tensor([3.0])) - - -def test_aggregate_group_logits_region_max_uses_instance_logits() -> None: - payload = BatchPredictionPayload( - bag_ids=["chunk_a", "chunk_b"], - bag_logits=torch.tensor([[5.0], [2.0]]), - bag_targets=None, - bag_attention=None, - row_region_ids=["rA", "rA"], - row_sample_ids=["sA", "sA"], - instance_logits=torch.tensor([[1.0], [9.0], [2.0]]), - instance_attention_logits=None, - instance_patch_ids=["chunk_a", "chunk_a", "chunk_b"], - instance_region_ids=["rA", "rA", "rA"], - instance_sample_ids=["sA", "sA", "sA"], - ) - out = aggregate_group_logits(payload, mode="max", bag_scope="region") - assert out.bag_ids == ["rA"] - assert torch.allclose(out.bag_logits[0], torch.tensor([9.0])) - - -def test_aggregate_group_logits_region_scope_falls_back_to_sample_ids_for_rows() -> None: - payload = BatchPredictionPayload( - bag_ids=["p0", "p1", "p2"], - bag_logits=torch.tensor([[1.0], [3.0], [8.0]]), - bag_targets=None, - bag_attention=None, - row_region_ids=[None, "", "r1"], - row_sample_ids=["s0", "s0", "s1"], - ) - with pytest.warns(RuntimeWarning, match="Fell back from 'row_region_ids' to 'row_sample_ids'"): - out = aggregate_group_logits(payload, mode="mean", bag_scope="region") - assert out.bag_ids == ["r1", "s0"] - assert torch.allclose(out.bag_logits[0], torch.tensor([8.0])) - assert torch.allclose(out.bag_logits[1], torch.tensor([2.0])) - - -def test_aggregate_group_logits_region_scope_falls_back_to_sample_ids_for_instances() -> None: - payload = BatchPredictionPayload( - bag_ids=["chunk_a", "chunk_b"], - bag_logits=torch.tensor([[5.0], [2.0]]), - bag_targets=None, - bag_attention=None, - row_region_ids=[None, "rB"], - row_sample_ids=["sA", "sB"], - instance_logits=torch.tensor([[1.0], [9.0], [3.0]]), - instance_attention_logits=None, - instance_patch_ids=["chunk_a", "chunk_a", "chunk_b"], - instance_region_ids=[None, "", "rB"], - instance_sample_ids=["sA", "sA", "sB"], - ) - with pytest.warns( - RuntimeWarning, - match="Fell back from 'instance_region_ids' to 'instance_sample_ids'", - ): - out = aggregate_group_logits(payload, mode="mean", bag_scope="region") - assert out.bag_ids == ["rB", "sA"] - assert torch.allclose(out.bag_logits[0], torch.tensor([3.0])) - assert torch.allclose(out.bag_logits[1], torch.tensor([5.0])) - - -def test_aggregate_group_logits_region_scope_raises_when_region_and_sample_missing() -> None: - payload = BatchPredictionPayload( - bag_ids=["p0"], - bag_logits=torch.tensor([[1.0]]), - bag_targets=None, - bag_attention=None, - row_region_ids=[None], - row_sample_ids=[None], - ) - with pytest.raises( - ValueError, match="Missing both 'row_region_ids' and fallback 'row_sample_ids'" - ): - aggregate_group_logits(payload, mode="mean", bag_scope="region") - - -def test_aggregate_group_logits_sample_scope_uses_instance_logits() -> None: - payload = BatchPredictionPayload( - bag_ids=["chunk_a", "chunk_b", "chunk_c"], - bag_logits=torch.tensor([[1.0], [3.0], [5.0]]), - bag_targets=None, - bag_attention=None, - row_region_ids=["r1", "r2", "r3"], - row_sample_ids=["sA", "sA", "sB"], - instance_logits=torch.tensor([[1.0], [5.0], [10.0], [7.0]]), - instance_attention_logits=None, - instance_patch_ids=["chunk_a", "chunk_b", "chunk_b", "chunk_c"], - instance_region_ids=["r1", "r2", "r2", "r3"], - instance_sample_ids=["sA", "sA", "sA", "sB"], - ) - out_mean = aggregate_group_logits(payload, mode="mean", bag_scope="sample") - out_max = aggregate_group_logits(payload, mode="max", bag_scope="sample") - assert out_mean.bag_ids == ["sA", "sB"] - assert torch.allclose(out_mean.bag_logits[0], torch.tensor([16.0 / 3.0])) - assert torch.allclose(out_mean.bag_logits[1], torch.tensor([7.0])) - assert torch.allclose(out_max.bag_logits[0], torch.tensor([10.0])) - assert torch.allclose(out_max.bag_logits[1], torch.tensor([7.0])) - - -def test_attention_weighted_region_scope_with_subsample_and_seed() -> None: - payload = BatchPredictionPayload( - bag_ids=["p0", "p1", "p2", "p3"], - bag_logits=torch.tensor([[0.0], [0.0], [0.0], [0.0]]), - bag_targets=None, - bag_attention=None, - row_region_ids=["rA", "rA", "rB", "rB"], - row_sample_ids=["s0", "s0", "s1", "s1"], - instance_logits=torch.tensor([[1.0], [2.0], [10.0], [20.0]]), - instance_attention_logits=torch.tensor([[1.0], [2.0], [1.0], [2.0]]), - instance_patch_ids=["p0", "p1", "p2", "p3"], - instance_region_ids=["rA", "rA", "rB", "rB"], - instance_sample_ids=["s0", "s0", "s1", "s1"], - ) - out = aggregate_group_logits( - payload, - mode="attention_weighted", - bag_scope="region", - subsample_fraction=0.5, - subsample_seed=7, - ) - assert out.bag_ids == ["rA", "rB"] - assert out.bag_logits.shape == (2, 1) - - -def test_attention_weighted_subsample_fraction_zero_raises() -> None: - payload = BatchPredictionPayload( - bag_ids=["r1"], - bag_logits=torch.tensor([[1.0]]), - bag_targets=None, - bag_attention=None, - row_region_ids=["region_a"], - row_sample_ids=["sample_x"], - instance_logits=torch.tensor([[1.0]]), - instance_attention_logits=torch.tensor([[1.0]]), - instance_patch_ids=["r1"], - instance_region_ids=["region_a"], - instance_sample_ids=["sample_x"], - ) - with pytest.raises(ValueError, match="subsample_fraction=0.0"): - aggregate_group_logits( - payload, - mode="attention_weighted", - subsample_fraction=0.0, - ) - - -def test_instance_aggregation_raises_when_row_group_metadata_missing() -> None: - payload = BatchPredictionPayload( - bag_ids=["chunk_a", "chunk_b"], - bag_logits=torch.tensor([[5.0], [2.0]]), - bag_targets=torch.tensor([[1.0], [0.0]]), - bag_attention=None, - row_region_ids=[None, "rB"], - row_sample_ids=[None, "sB"], - instance_logits=torch.tensor([[1.0], [9.0], [3.0]]), - instance_attention_logits=None, - instance_patch_ids=["chunk_a", "chunk_a", "chunk_b"], - instance_region_ids=["rA", "rA", "rB"], - instance_sample_ids=["sA", "sA", "sB"], - ) - with pytest.raises( - ValueError, match="Missing both 'row_region_ids' and fallback 'row_sample_ids'" - ): - aggregate_group_logits(payload, mode="mean", bag_scope="region") diff --git a/tests/eval/test_collectors.py b/tests/eval/test_collectors.py deleted file mode 100644 index f1911b3..0000000 --- a/tests/eval/test_collectors.py +++ /dev/null @@ -1,293 +0,0 @@ -from __future__ import annotations - -import pytest -import torch -from types import SimpleNamespace - -from grass_mil.inference.collectors import collect_inference_payload - - -class _FakeTrainer: - def __init__(self, outputs): - self._outputs = outputs - - def predict(self, model=None, datamodule=None, ckpt_path=None): - return self._outputs - - -def test_collect_inference_payload_normalizes_predictions() -> None: - outputs = [ - { - "bag_ids": ["region_1", "region_2"], - "bag_logits": torch.tensor([[1.0], [0.0]]), - "bag_targets": torch.tensor([[1.0], [0.0]]), - "row_region_ids": ["region_1", "region_2"], - "row_sample_ids": ["sample_a", "sample_b"], - }, - { - "bag_ids": ["region_1"], - "bag_logits": torch.tensor([[2.0]]), - "bag_targets": torch.tensor([[1.0]]), - "bag_attention": {"region_1": torch.tensor([0.2, 0.8])}, - "row_region_ids": ["region_1"], - "row_sample_ids": ["sample_a"], - "instance_logits": torch.tensor([[1.5], [2.5]]), - "instance_attention_logits": torch.tensor([[0.1], [0.2]]), - "instance_patch_ids": ["region_1", "region_1"], - "instance_region_ids": ["region_1", "region_1"], - "instance_sample_ids": ["sample_a", "sample_a"], - }, - ] - trainer = _FakeTrainer(outputs) - collected = collect_inference_payload( - trainer=trainer, # type: ignore[arg-type] - model=None, # type: ignore[arg-type] - datamodule=None, # type: ignore[arg-type] - ckpt_path=None, - ) - payload = collected.prediction_payload - assert payload.bag_logits.shape[0] == 3 - assert payload.bag_targets is not None - assert payload.bag_attention is not None - assert payload.bag_ids[0] == "region_1" - assert len(payload.bag_attention) == 3 - assert payload.bag_attention[0] is None - assert payload.bag_attention[1] is None - assert torch.allclose(payload.bag_attention[2], torch.tensor([0.2, 0.8])) - assert payload.row_region_ids == ["region_1", "region_2", "region_1"] - assert payload.row_sample_ids == ["sample_a", "sample_b", "sample_a"] - assert payload.instance_logits is not None - assert payload.instance_logits.shape == (2, 1) - assert payload.instance_patch_ids == ["region_1", "region_1"] - assert collected.embedding_payload is None - - -def test_collect_inference_payload_can_skip_instance_payload_collection() -> None: - outputs = [ - { - "bag_ids": ["region_1"], - "bag_logits": torch.tensor([[1.0]]), - "row_region_ids": ["region_1"], - "row_sample_ids": ["sample_a"], - # Intentionally malformed instance payload to ensure collector ignores it. - "instance_logits": torch.tensor([[1.0], [2.0]]), - "instance_patch_ids": ["patch_1"], - "instance_region_ids": ["region_1"], - "instance_sample_ids": ["sample_a"], - } - ] - trainer = _FakeTrainer(outputs) - collected = collect_inference_payload( - trainer=trainer, # type: ignore[arg-type] - model=None, # type: ignore[arg-type] - datamodule=None, # type: ignore[arg-type] - ckpt_path=None, - include_instance_payload=False, - ) - payload = collected.prediction_payload - assert payload.bag_logits.shape == (1, 1) - assert payload.instance_logits is None - assert payload.instance_patch_ids is None - assert payload.instance_region_ids is None - assert payload.instance_sample_ids is None - - -def test_collect_inference_payload_aggregates_duplicate_embedding_bag_ids_with_counts() -> None: - outputs = [ - { - "bag_ids": ["region_1"], - "bag_logits": torch.tensor([[1.0]]), - "row_region_ids": ["region_1"], - "row_sample_ids": ["sample_a"], - "embedding_bag_ids": ["region_1", "region_2"], - "graph_embeddings": torch.tensor([[1.0, 1.0], [10.0, 10.0]]), - "embedding_bag_counts": [2, 1], - "node_embeddings": torch.tensor([[0.1, 0.2], [0.3, 0.4]]), - "node_bag_ids": ["region_1", "region_2"], - }, - { - "bag_ids": ["region_2"], - "bag_logits": torch.tensor([[2.0]]), - "row_region_ids": ["region_2"], - "row_sample_ids": ["sample_b"], - "embedding_bag_ids": ["region_1"], - "graph_embeddings": torch.tensor([[3.0, 3.0]]), - "embedding_bag_counts": [1], - "node_embeddings": torch.tensor([[0.5, 0.6]]), - "node_bag_ids": ["region_1"], - }, - ] - trainer = _FakeTrainer(outputs) - collected = collect_inference_payload( - trainer=trainer, # type: ignore[arg-type] - model=None, # type: ignore[arg-type] - datamodule=None, # type: ignore[arg-type] - ckpt_path=None, - include_embeddings=True, - include_node_embeddings=True, - ) - payload = collected.embedding_payload - assert payload is not None - assert payload.bag_ids == ["region_1", "region_2"] - # region_1 weighted mean: (1.0*2 + 3.0*1) / 3 - assert torch.allclose(payload.graph_embeddings[0], torch.tensor([5.0 / 3.0, 5.0 / 3.0])) - assert torch.allclose(payload.graph_embeddings[1], torch.tensor([10.0, 10.0])) - assert payload.node_embeddings is not None - assert payload.node_embeddings.shape == (3, 2) - assert payload.node_bag_ids == ["region_1", "region_2", "region_1"] - - -def test_collect_inference_payload_raises_when_embedding_fields_missing() -> None: - outputs = [ - { - "bag_ids": ["region_1"], - "bag_logits": torch.tensor([[1.0]]), - "row_region_ids": ["region_1"], - "row_sample_ids": ["sample_a"], - "embedding_bag_ids": ["region_1"], - # Missing graph_embeddings and embedding_bag_counts. - } - ] - trainer = _FakeTrainer(outputs) - with pytest.raises(ValueError, match="missing required embedding fields"): - collect_inference_payload( - trainer=trainer, # type: ignore[arg-type] - model=None, # type: ignore[arg-type] - datamodule=None, # type: ignore[arg-type] - ckpt_path=None, - include_embeddings=True, - ) - - -def test_collect_inference_payload_raises_when_node_ids_missing() -> None: - outputs = [ - { - "bag_ids": ["region_1"], - "bag_logits": torch.tensor([[1.0]]), - "row_region_ids": ["region_1"], - "row_sample_ids": ["sample_a"], - "embedding_bag_ids": ["region_1"], - "graph_embeddings": torch.tensor([[1.0, 1.0]]), - "embedding_bag_counts": [1], - "node_embeddings": torch.tensor([[0.1, 0.2]]), - } - ] - trainer = _FakeTrainer(outputs) - with pytest.raises(ValueError, match="must include both node_embeddings and node_bag_ids"): - collect_inference_payload( - trainer=trainer, # type: ignore[arg-type] - model=None, # type: ignore[arg-type] - datamodule=None, # type: ignore[arg-type] - ckpt_path=None, - include_embeddings=True, - include_node_embeddings=True, - ) - - -def test_collect_inference_payload_collects_instance_export_fields() -> None: - outputs = [ - { - "bag_ids": ["region_1"], - "bag_logits": torch.tensor([[1.0]]), - "row_region_ids": ["region_1"], - "row_sample_ids": ["sample_a"], - "instance_logits": torch.tensor([[1.5], [2.5]]), - "instance_attention_logits": torch.tensor([[0.1], [0.2]]), - "instance_patch_ids": ["patch_0", "patch_1"], - "instance_bag_ids": ["sample_a::region_1", "sample_a::region_1"], - "instance_region_ids": ["region_1", "region_1"], - "instance_sample_ids": ["sample_a", "sample_a"], - "instance_embeddings": torch.tensor([[0.11, 0.12], [0.21, 0.22]]), - "instance_composition": torch.tensor([[1.0, 0.0], [0.2, 0.8]]), - "instance_centroids": torch.tensor([[10.0, 11.0], [12.0, 13.0]]), - "instance_graphs": [ - SimpleNamespace(name="g0"), - SimpleNamespace(name="g1"), - ], - } - ] - trainer = _FakeTrainer(outputs) - collected = collect_inference_payload( - trainer=trainer, # type: ignore[arg-type] - model=None, # type: ignore[arg-type] - datamodule=None, # type: ignore[arg-type] - ckpt_path=None, - include_instance_embeddings=True, - ) - payload = collected.prediction_payload - assert payload.instance_bag_ids == ["sample_a::region_1", "sample_a::region_1"] - assert payload.instance_embeddings is not None - assert payload.instance_embeddings.shape == (2, 2) - assert payload.instance_composition is not None - assert payload.instance_composition.shape == (2, 2) - assert payload.instance_centroids is not None - assert payload.instance_centroids.shape == (2, 2) - assert payload.instance_graphs is not None - assert len(payload.instance_graphs) == 2 - - -def test_collect_inference_payload_requires_instance_embeddings_when_enabled() -> None: - outputs = [ - { - "bag_ids": ["region_1"], - "bag_logits": torch.tensor([[1.0]]), - "row_region_ids": ["region_1"], - "row_sample_ids": ["sample_a"], - "instance_logits": torch.tensor([[1.5]]), - "instance_patch_ids": ["patch_0"], - "instance_bag_ids": ["sample_a::region_1"], - "instance_region_ids": ["region_1"], - "instance_sample_ids": ["sample_a"], - } - ] - trainer = _FakeTrainer(outputs) - with pytest.raises(KeyError, match="instance_embeddings"): - collect_inference_payload( - trainer=trainer, # type: ignore[arg-type] - model=None, # type: ignore[arg-type] - datamodule=None, # type: ignore[arg-type] - ckpt_path=None, - include_instance_embeddings=True, - ) - - -def test_collect_inference_payload_collects_instance_graphs_across_chunks() -> None: - outputs = [ - { - "bag_ids": ["region_1"], - "bag_logits": torch.tensor([[1.0]]), - "row_region_ids": ["region_1"], - "row_sample_ids": ["sample_a"], - "instance_logits": torch.tensor([[1.5], [2.5]]), - "instance_patch_ids": ["patch_0", "patch_1"], - "instance_bag_ids": ["sample_a::region_1", "sample_a::region_1"], - "instance_region_ids": ["region_1", "region_1"], - "instance_sample_ids": ["sample_a", "sample_a"], - "instance_embeddings": torch.tensor([[0.1, 0.2], [0.3, 0.4]]), - "instance_graphs": [SimpleNamespace(name="g0"), SimpleNamespace(name="g1")], - }, - { - "bag_ids": ["region_1"], - "bag_logits": torch.tensor([[1.0]]), - "row_region_ids": ["region_1"], - "row_sample_ids": ["sample_a"], - "instance_logits": torch.tensor([[3.5]]), - "instance_patch_ids": ["patch_2"], - "instance_bag_ids": ["sample_a::region_1"], - "instance_region_ids": ["region_1"], - "instance_sample_ids": ["sample_a"], - "instance_embeddings": torch.tensor([[0.5, 0.6]]), - "instance_graphs": [SimpleNamespace(name="g2")], - }, - ] - trainer = _FakeTrainer(outputs) - collected = collect_inference_payload( - trainer=trainer, # type: ignore[arg-type] - model=None, # type: ignore[arg-type] - datamodule=None, # type: ignore[arg-type] - ckpt_path=None, - include_instance_embeddings=True, - ) - payload = collected.prediction_payload - assert payload.instance_graphs is not None - assert [graph.name for graph in payload.instance_graphs] == ["g0", "g1", "g2"] diff --git a/tests/eval/test_embeddings.py b/tests/eval/test_embeddings.py deleted file mode 100644 index a4382ec..0000000 --- a/tests/eval/test_embeddings.py +++ /dev/null @@ -1,21 +0,0 @@ -from __future__ import annotations - -import torch - -from grass_mil.inference.io import embeddings_to_dataframe, node_embeddings_to_dataframe -from grass_mil.inference.schemas import EmbeddingPayload - - -def test_embedding_payload_to_dataframes() -> None: - payload = EmbeddingPayload( - bag_ids=["r2", "r1"], - graph_embeddings=torch.tensor([[2.0, 3.0], [1.0, 1.5]]), - node_embeddings=torch.tensor([[0.1, 0.2], [0.3, 0.4]]), - node_bag_ids=["r1", "r2"], - ) - graph_df = embeddings_to_dataframe(payload) - node_df = node_embeddings_to_dataframe(payload) - - assert list(graph_df.columns) == ["bag_id", "graph_emb_0", "graph_emb_1"] - assert graph_df.iloc[0]["bag_id"] == "r1" - assert list(node_df.columns) == ["bag_id", "node_emb_0", "node_emb_1"] diff --git a/tests/eval/test_interpretability_export.py b/tests/eval/test_interpretability_export.py deleted file mode 100644 index 0a58556..0000000 --- a/tests/eval/test_interpretability_export.py +++ /dev/null @@ -1,392 +0,0 @@ -from __future__ import annotations - -import math -from types import SimpleNamespace - -import pytest -import torch - -from grass_mil.inference.interpretability_export import build_instance_table, build_spatial_table -from grass_mil.inference.schemas import BatchPredictionPayload - - -def test_build_instance_table_adds_uniform_attention_and_unique_ids() -> None: - payload = BatchPredictionPayload( - bag_ids=["b0"], - bag_logits=torch.tensor([[1.0]]), - bag_targets=None, - bag_attention=None, - instance_logits=torch.tensor([[0.2], [0.5], [1.2]]), - instance_patch_ids=["p0", "p0", "p1"], - instance_bag_ids=["b0", "b0", "b1"], - instance_region_ids=["r0", "r0", "r1"], - instance_sample_ids=["s0", "s0", "s1"], - instance_embeddings=torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]), - instance_composition=torch.tensor([[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]]), - ) - frame = build_instance_table(payload, require_composition=True) - assert list(frame["instance_id"]) == ["p0", "p0::1", "p1"] - assert "attention" in frame.columns - # Bag b0 has two instances with uniform fallback attention. - assert frame.loc[frame["bag_id"] == "b0", "attention"].tolist() == [0.5, 0.5] - assert frame.loc[frame["bag_id"] == "b1", "attention"].tolist() == [1.0] - expected_scores = [1.0 / (1.0 + math.exp(-v)) for v in [0.2, 0.5, 1.2]] - assert frame["score"].tolist() == pytest.approx(expected_scores) - assert "inst_emb_0" in frame.columns - assert "comp_0" in frame.columns - - -def test_build_instance_table_supports_identity_score_mode() -> None: - payload = BatchPredictionPayload( - bag_ids=["b0"], - bag_logits=torch.tensor([[1.0]]), - bag_targets=None, - bag_attention=None, - instance_logits=torch.tensor([[0.2], [0.5], [1.2]]), - instance_patch_ids=["p0", "p1", "p2"], - instance_bag_ids=["b0", "b0", "b1"], - instance_region_ids=["r0", "r0", "r1"], - instance_sample_ids=["s0", "s0", "s1"], - instance_embeddings=torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]), - instance_composition=torch.tensor([[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]]), - ) - frame = build_instance_table( - payload, - require_composition=True, - score_mode="identity", - score_logit_index=0, - ) - assert frame["score"].tolist() == pytest.approx([0.2, 0.5, 1.2]) - - -def test_build_instance_table_supports_softmax_score_mode() -> None: - logits = torch.tensor([[2.0, 1.0], [0.5, 3.0], [1.5, 1.5]]) - payload = BatchPredictionPayload( - bag_ids=["b0"], - bag_logits=torch.tensor([[1.0, 0.0]]), - bag_targets=None, - bag_attention=None, - instance_logits=logits, - instance_patch_ids=["p0", "p1", "p2"], - instance_bag_ids=["b0", "b0", "b1"], - instance_region_ids=["r0", "r0", "r1"], - instance_sample_ids=["s0", "s0", "s1"], - instance_embeddings=torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]), - instance_composition=torch.tensor([[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]]), - ) - frame = build_instance_table( - payload, - require_composition=True, - score_mode="softmax", - score_logit_index=1, - ) - expected_scores = torch.softmax(logits, dim=1)[:, 1].tolist() - assert frame["score"].tolist() == pytest.approx(expected_scores) - - -def test_build_instance_table_rejects_unknown_score_mode() -> None: - payload = BatchPredictionPayload( - bag_ids=["b0"], - bag_logits=torch.tensor([[1.0]]), - bag_targets=None, - bag_attention=None, - instance_logits=torch.tensor([[0.2], [0.5]]), - instance_patch_ids=["p0", "p1"], - instance_bag_ids=["b0", "b0"], - instance_region_ids=["r0", "r0"], - instance_sample_ids=["s0", "s0"], - instance_embeddings=torch.tensor([[0.1, 0.2], [0.3, 0.4]]), - instance_composition=torch.tensor([[1.0, 0.0], [0.5, 0.5]]), - ) - with pytest.raises(ValueError, match="Unsupported score_mode"): - build_instance_table(payload, require_composition=True, score_mode="unknown") - - -def test_build_instance_table_rejects_out_of_bounds_score_index() -> None: - payload = BatchPredictionPayload( - bag_ids=["b0"], - bag_logits=torch.tensor([[1.0]]), - bag_targets=None, - bag_attention=None, - instance_logits=torch.tensor([[0.2], [0.5]]), - instance_patch_ids=["p0", "p1"], - instance_bag_ids=["b0", "b0"], - instance_region_ids=["r0", "r0"], - instance_sample_ids=["s0", "s0"], - instance_embeddings=torch.tensor([[0.1, 0.2], [0.3, 0.4]]), - instance_composition=torch.tensor([[1.0, 0.0], [0.5, 0.5]]), - ) - with pytest.raises(ValueError, match="score_logit_index is out of bounds"): - build_instance_table(payload, require_composition=True, score_logit_index=3) - - -def test_build_spatial_table_uses_payload_connectivity_edges() -> None: - payload = BatchPredictionPayload( - bag_ids=["b0"], - bag_logits=torch.tensor([[1.0]]), - bag_targets=None, - bag_attention=None, - instance_logits=torch.tensor([[0.2], [0.5], [1.2]]), - instance_patch_ids=["p0", "p1", "p2"], - instance_bag_ids=["b0", "b0", "b0"], - instance_region_ids=["r0", "r0", "r0"], - instance_sample_ids=["s0", "s0", "s0"], - instance_embeddings=torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]), - instance_composition=torch.tensor([[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]]), - instance_graphs=[ - SimpleNamespace( - root_n_id=torch.tensor([0], dtype=torch.long), - n_id=torch.tensor([100, 101], dtype=torch.long), - edge_index=torch.tensor([[0], [1]], dtype=torch.long), - edge_attr=torch.tensor([[1.0]], dtype=torch.float32), - edge_attr_names=["distance"], - ), - SimpleNamespace( - root_n_id=torch.tensor([0], dtype=torch.long), - n_id=torch.tensor([101, 100, 102], dtype=torch.long), - edge_index=torch.tensor([[0, 0], [1, 2]], dtype=torch.long), - edge_attr=torch.tensor([[1.0], [2.0]], dtype=torch.float32), - edge_attr_names=["distance"], - ), - SimpleNamespace( - root_n_id=torch.tensor([0], dtype=torch.long), - n_id=torch.tensor([102, 101], dtype=torch.long), - edge_index=torch.tensor([[0], [1]], dtype=torch.long), - edge_attr=torch.tensor([[2.0]], dtype=torch.float32), - edge_attr_names=["distance"], - ), - ], - ) - instance_frame = build_instance_table(payload, require_composition=True) - spatial = build_spatial_table(instance_frame, payload, undirected=True) - assert set(spatial.columns) == {"source_id", "target_id", "distance"} - # Undirected expansion guarantees reverse edges. - assert ("p0", "p1") in set(zip(spatial["source_id"], spatial["target_id"])) - assert ("p1", "p0") in set(zip(spatial["source_id"], spatial["target_id"])) - - -def test_build_spatial_table_requires_payload_connectivity() -> None: - payload = BatchPredictionPayload( - bag_ids=["b0"], - bag_logits=torch.tensor([[1.0]]), - bag_targets=None, - bag_attention=None, - instance_logits=torch.tensor([[0.2], [0.5]]), - instance_patch_ids=["p0", "p1"], - instance_bag_ids=["b0", "b0"], - instance_region_ids=["r0", "r0"], - instance_sample_ids=["s0", "s0"], - instance_embeddings=torch.tensor([[0.1, 0.2], [0.3, 0.4]]), - instance_composition=torch.tensor([[1.0, 0.0], [0.5, 0.5]]), - ) - instance_frame = build_instance_table(payload, require_composition=True) - with pytest.raises(ValueError, match="missing instance_graphs"): - build_spatial_table(instance_frame, payload) - - -def test_build_instance_table_exports_per_class_attention_columns() -> None: - """Per-class attention is exported, not rejected. - - The attribution identity needs one attention column per class, plus the raw - pre-softmax logits (which, unlike the within-bag normalised scores, are - comparable across bags). - """ - payload = BatchPredictionPayload( - bag_ids=["b0"], - bag_logits=torch.tensor([[1.0, 0.0]]), - bag_targets=None, - bag_attention=None, - instance_logits=torch.tensor([[0.2, 0.7], [0.5, 0.1]]), - instance_attention_logits=torch.tensor([[0.1, 0.2], [0.3, 0.4]]), - instance_patch_ids=["p0", "p1"], - instance_bag_ids=["b0", "b0"], - instance_region_ids=["r0", "r0"], - instance_sample_ids=["s0", "s0"], - instance_embeddings=torch.tensor([[0.1, 0.2], [0.3, 0.4]]), - instance_composition=torch.tensor([[1.0, 0.0], [0.5, 0.5]]), - ) - table = build_instance_table(payload, require_composition=True) - - for column in ("attention", "attention_c0", "attention_c1"): - assert column in table.columns - assert "attention_logit_c0" in table.columns - - # Each class column is a distribution over the bag's instances. - assert table["attention_c0"].sum() == pytest.approx(1.0) - assert table["attention_c1"].sum() == pytest.approx(1.0) - # The canonical `attention` column defaults to the last class. - assert table["attention"].tolist() == pytest.approx(table["attention_c1"].tolist()) - # Raw logits are passed through unnormalised. - assert table["attention_logit_c0"].tolist() == pytest.approx([0.1, 0.3]) - - -def test_build_instance_table_attention_class_index_is_selectable() -> None: - payload = BatchPredictionPayload( - bag_ids=["b0"], - bag_logits=torch.tensor([[1.0, 0.0]]), - bag_targets=None, - bag_attention=None, - instance_logits=torch.tensor([[0.2, 0.7], [0.5, 0.1]]), - instance_attention_logits=torch.tensor([[0.1, 5.0], [0.3, 0.4]]), - instance_patch_ids=["p0", "p1"], - instance_bag_ids=["b0", "b0"], - instance_region_ids=["r0", "r0"], - instance_sample_ids=["s0", "s0"], - instance_embeddings=torch.tensor([[0.1, 0.2], [0.3, 0.4]]), - instance_composition=torch.tensor([[1.0, 0.0], [0.5, 0.5]]), - ) - table = build_instance_table(payload, require_composition=True, attention_class_index=0) - assert table["attention"].tolist() == pytest.approx(table["attention_c0"].tolist()) - - with pytest.raises(ValueError, match="out of range"): - build_instance_table(payload, require_composition=True, attention_class_index=5) - - -def _instance_graph(codes, n_id, root_global): - """Minimal stand-in for a sampled ego-graph.""" - from torch_geometric.data import Data - - g = Data(x=torch.zeros(len(n_id), 1)) - g.categorical_codes = torch.tensor(codes).view(-1, 1) - g.n_id = torch.tensor(n_id) - g.root_n_id = torch.tensor([root_global]) - g.root_n_id_is_global = torch.tensor([True]) - g.categorical_slices = {"cell_type": 0} - return g - - -def test_build_instance_table_emits_the_root_cell_type() -> None: - """One label per instance: the type of the cell it is rooted at. - - There is exactly one instance per cell -- an instance is that cell's k-hop - ego-graph -- so the instance table is a cell table. `filtration_curves` and - `per_niche_cell_type_enrichment` both need this column and raise without it. - - It must be the root cell's own type. Summarising the neighbourhood instead - (its dominant type, say) would quietly turn cell-level statistics into - statistics over neighbourhood labels. The neighbourhood composition is - exported separately as `comp_*`. - """ - # Instance 0 is rooted at global node 7, whose code is 1 ("Tcell"), even - # though its neighbourhood is mostly code 0 ("Epithelial"). - graphs = [ - _instance_graph(codes=[0, 0, 1], n_id=[5, 6, 7], root_global=7), - _instance_graph(codes=[0, 1, 1], n_id=[8, 9, 10], root_global=8), - ] - payload = BatchPredictionPayload( - bag_ids=["b0"], - bag_logits=torch.tensor([[1.0]]), - bag_targets=None, - bag_attention=None, - instance_logits=torch.tensor([[0.2], [0.5]]), - instance_patch_ids=["p0", "p1"], - instance_bag_ids=["b0", "b0"], - instance_region_ids=["r0", "r0"], - instance_sample_ids=["s0", "s0"], - instance_embeddings=torch.tensor([[0.1, 0.2], [0.3, 0.4]]), - instance_composition=torch.tensor([[0.67, 0.33], [0.33, 0.67]]), - instance_graphs=graphs, - ) - frame = build_instance_table( - payload, - require_composition=True, - composition_column_names=["comp_Epithelial", "comp_Tcell"], - ) - assert "cell_type" in frame.columns - # Root types, not the neighbourhood majority: instance 0 is majority - # Epithelial but rooted at a Tcell, and instance 1 the other way round. - assert frame["cell_type"].tolist() == ["Tcell", "Epithelial"] - - -def test_instance_cell_type_column_can_be_disabled() -> None: - graphs = [_instance_graph(codes=[0, 1], n_id=[3, 4], root_global=3)] - payload = BatchPredictionPayload( - bag_ids=["b0"], - bag_logits=torch.tensor([[1.0]]), - bag_targets=None, - bag_attention=None, - instance_logits=torch.tensor([[0.2]]), - instance_patch_ids=["p0"], - instance_bag_ids=["b0"], - instance_region_ids=["r0"], - instance_sample_ids=["s0"], - instance_embeddings=torch.tensor([[0.1, 0.2]]), - instance_composition=torch.tensor([[0.7, 0.3]]), - instance_graphs=graphs, - ) - frame = build_instance_table(payload, require_composition=True, cell_type_column=None) - assert "cell_type" not in frame.columns - - -def test_build_cell_tables_pools_every_node_with_its_own_type() -> None: - """Cell-level tables: one row per node of every sampled ego-graph. - - The instance table describes neighbourhoods, reducing their cell content to - composition fractions. Cell-type enrichment within a niche, per-cell-type - autocorrelation and distance filtration all need the cells themselves, which - that summary cannot reconstruct. - - Cells are pooled with duplicates -- overlapping ego-graphs revisit the same - cell -- so (patch_id, node_id) identifies a real cell, not cell_id. - """ - from grass_mil.inference.interpretability_export import build_cell_tables - - g0 = _instance_graph(codes=[0, 1, 1], n_id=[5, 6, 7], root_global=5) - g0.pos = torch.tensor([[0.0, 0.0], [1.0, 0.0], [2.0, 0.0]]) - g0.edge_index = torch.tensor([[0, 1], [1, 2]]) - g0.edge_attr = torch.tensor([[1.0, 1.0], [2.0, 1.0]]) - g0.edge_attr_names = ["distance", "neighbor"] - g0.patch_id = ["patchA"] - - payload = BatchPredictionPayload( - bag_ids=["b0"], - bag_logits=torch.tensor([[1.0]]), - bag_targets=None, - bag_attention=None, - instance_logits=torch.tensor([[0.2]]), - instance_patch_ids=["p0"], - instance_bag_ids=["b0"], - instance_region_ids=["r0"], - instance_sample_ids=["s0"], - instance_embeddings=torch.tensor([[0.1, 0.2]]), - instance_composition=torch.tensor([[0.34, 0.66]]), - instance_graphs=[g0], - ) - cells, edges = build_cell_tables( - payload, - ["inst0"], - composition_column_names=["comp_Epithelial", "comp_Tcell"], - ) - - # Every node becomes a row, carrying its own type -- not the neighbourhood's. - assert len(cells) == 3 - assert cells["cell_type"].tolist() == ["Epithelial", "Tcell", "Tcell"] - assert cells["instance_id"].unique().tolist() == ["inst0"] - assert cells["patch_id"].unique().tolist() == ["patchA"] - # Exactly one root, and it is the node whose n_id matches root_n_id. - assert cells["is_root"].sum() == 1 - assert cells.loc[cells["is_root"], "node_id"].tolist() == [5] - - # Edges carry their distance, so filtration over thresholds is possible. - assert len(edges) == 2 - assert edges["distance"].tolist() == [1.0, 2.0] - assert set(edges["source_id"]) <= set(cells["cell_id"]) - - -def test_build_cell_tables_returns_none_without_subgraphs() -> None: - from grass_mil.inference.interpretability_export import build_cell_tables - - payload = BatchPredictionPayload( - bag_ids=["b0"], - bag_logits=torch.tensor([[1.0]]), - bag_targets=None, - bag_attention=None, - instance_logits=torch.tensor([[0.2]]), - instance_patch_ids=["p0"], - instance_bag_ids=["b0"], - instance_region_ids=["r0"], - instance_sample_ids=["s0"], - instance_embeddings=torch.tensor([[0.1, 0.2]]), - instance_composition=torch.tensor([[0.7, 0.3]]), - ) - assert build_cell_tables(payload, ["inst0"]) is None diff --git a/tests/eval/test_io.py b/tests/eval/test_io.py deleted file mode 100644 index e9b00bc..0000000 --- a/tests/eval/test_io.py +++ /dev/null @@ -1,49 +0,0 @@ -from __future__ import annotations - -import json -import math -from pathlib import Path - -from grass_mil.inference.io import write_json - - -def test_write_json_sanitizes_non_finite_floats(tmp_path: Path) -> None: - payload = { - "nan_metric": float("nan"), - "pos_inf_metric": float("inf"), - "neg_inf_metric": float("-inf"), - "nested": {"values": [1.0, float("nan"), 3.0]}, - } - out_path = tmp_path / "metrics.json" - - write_json(payload, out_path) - - contents = out_path.read_text() - assert "NaN" not in contents - assert "Infinity" not in contents - assert "-Infinity" not in contents - - parsed = json.loads(contents) - assert parsed["nan_metric"] is None - assert parsed["pos_inf_metric"] is None - assert parsed["neg_inf_metric"] is None - assert parsed["nested"]["values"] == [1.0, None, 3.0] - - -def test_write_json_preserves_finite_values(tmp_path: Path) -> None: - payload = { - "accuracy": 0.95, - "count": 3, - "items": [0.0, -2.5], - "nested": {"ok": True}, - } - out_path = tmp_path / "summary.json" - - write_json(payload, out_path) - - parsed = json.loads(out_path.read_text()) - assert parsed["accuracy"] == 0.95 - assert parsed["count"] == 3 - assert parsed["items"] == [0.0, -2.5] - assert parsed["nested"]["ok"] is True - assert math.isfinite(parsed["accuracy"]) diff --git a/tests/eval/test_metrics.py b/tests/eval/test_metrics.py deleted file mode 100644 index f8981bd..0000000 --- a/tests/eval/test_metrics.py +++ /dev/null @@ -1,231 +0,0 @@ -from __future__ import annotations - -import math - -import pytest -import torch -from omegaconf import OmegaConf -from sklearn.metrics import roc_auc_score - -from grass_mil.inference.metrics import compute_task_metrics -from grass_mil.inference.schemas import BatchPredictionPayload - - -def test_compute_binary_metrics() -> None: - payload = BatchPredictionPayload( - bag_ids=["a", "b", "c", "d"], - bag_logits=torch.tensor([[2.0], [-1.0], [1.5], [-2.0]]), - bag_targets=torch.tensor([[1.0], [0.0], [1.0], [0.0]]), - bag_attention=None, - ) - cfg = OmegaConf.create({"categorical": {"threshold": 0.5}, "per_group": True}) - metrics = compute_task_metrics(payload=payload, target_type="binary", metrics_cfg=cfg) - assert "accuracy" in metrics["global"] - assert "balanced_accuracy" in metrics["global"] - assert "roc_auc" in metrics["global"] - assert metrics["global"]["accuracy"] >= 0.5 - assert "a" in metrics["per_group"] - - -def test_compute_regression_metrics() -> None: - payload = BatchPredictionPayload( - bag_ids=["a", "b", "c"], - bag_logits=torch.tensor([[1.0], [2.0], [3.0]]), - bag_targets=torch.tensor([[1.1], [1.9], [3.2]]), - bag_attention=None, - ) - cfg = OmegaConf.create({"categorical": {"threshold": 0.5}, "per_group": False}) - metrics = compute_task_metrics(payload=payload, target_type="regression", metrics_cfg=cfg) - assert "mae" in metrics["global"] - assert "rmse" in metrics["global"] - assert "r2" in metrics["global"] - - -def test_compute_survival_metrics() -> None: - payload = BatchPredictionPayload( - bag_ids=["a", "b", "c", "d"], - bag_logits=torch.tensor([[0.8], [0.1], [0.7], [0.2]]), - bag_targets=torch.tensor([[2.0, 1.0], [4.0, 1.0], [3.0, 0.0], [5.0, 1.0]]), - bag_attention=None, - ) - cfg = OmegaConf.create({"categorical": {"threshold": 0.5}, "per_group": False}) - metrics = compute_task_metrics(payload=payload, target_type="survival", metrics_cfg=cfg) - assert "c_index" in metrics["global"] - - -def test_compute_categorical_metrics() -> None: - payload = BatchPredictionPayload( - bag_ids=["a", "b", "c", "d"], - bag_logits=torch.tensor( - [ - [4.0, 0.1, -1.0], - [0.1, 3.0, 0.2], - [-0.2, 0.1, 2.5], - [0.5, 1.5, 0.2], - ] - ), - bag_targets=torch.tensor([0, 1, 2, 1]), - bag_attention=None, - ) - cfg = OmegaConf.create({"categorical": {"threshold": 0.5}, "per_group": True}) - metrics = compute_task_metrics(payload=payload, target_type="categorical", metrics_cfg=cfg) - assert "accuracy" in metrics["global"] - assert "balanced_accuracy" in metrics["global"] - assert "precision" in metrics["global"] - assert "recall" in metrics["global"] - assert "f1" in metrics["global"] - assert "a" in metrics["per_group"] - - -def test_compute_binary_metrics_matches_sklearn_auc_with_ties() -> None: - payload = BatchPredictionPayload( - bag_ids=["a", "b", "c", "d"], - bag_logits=torch.tensor([[0.0], [0.0], [-1.0], [2.0]]), - bag_targets=torch.tensor([[1.0], [0.0], [0.0], [1.0]]), - bag_attention=None, - ) - cfg = OmegaConf.create({"categorical": {"threshold": 0.5}, "per_group": False}) - metrics = compute_task_metrics(payload=payload, target_type="binary", metrics_cfg=cfg) - expected = roc_auc_score( - payload.bag_targets.view(-1).numpy(), - torch.sigmoid(payload.bag_logits).view(-1).numpy(), - ) - assert metrics["global"]["roc_auc"] == expected - - -def test_compute_binary_multitask_metrics_are_per_task_not_flattened() -> None: - # Task 0 is perfect, task 1 is perfectly wrong; macro accuracy should be 0.5. - payload = BatchPredictionPayload( - bag_ids=["a", "b", "c", "d"], - bag_logits=torch.tensor( - [ - [8.0, 8.0], - [-8.0, -8.0], - [8.0, 8.0], - [-8.0, -8.0], - ] - ), - bag_targets=torch.tensor( - [ - [1.0, 0.0], - [0.0, 1.0], - [1.0, 0.0], - [0.0, 1.0], - ] - ), - bag_attention=None, - ) - cfg = OmegaConf.create({"categorical": {"threshold": 0.5}, "per_group": False}) - metrics = compute_task_metrics(payload=payload, target_type="binary", metrics_cfg=cfg) - - assert metrics["global"]["accuracy"] == 0.5 - assert metrics["global"]["balanced_accuracy"] == 0.5 - assert "per_task" in metrics["global"] - assert metrics["global"]["per_task"]["task_0"]["accuracy"] == 1.0 - assert metrics["global"]["per_task"]["task_0"]["balanced_accuracy"] == 1.0 - assert metrics["global"]["per_task"]["task_1"]["accuracy"] == 0.0 - assert metrics["global"]["per_task"]["task_1"]["balanced_accuracy"] == 0.0 - - -def _reference_concordance_index(risk, time, event) -> float: - """Straightforward O(n^2) reference for the vectorised implementation.""" - concordant = ties = comparable = 0.0 - n = len(risk) - for i in range(n): - if event[i] <= 0: - continue - for j in range(n): - if i == j or not time[i] < time[j]: - continue - comparable += 1.0 - if risk[i] > risk[j]: - concordant += 1.0 - elif risk[i] == risk[j]: - ties += 1.0 - if comparable == 0.0: - return float("nan") - return (concordant + 0.5 * ties) / comparable - - -@pytest.mark.parametrize("seed", [0, 1, 2, 3]) -def test_concordance_index_matches_the_reference_double_loop(seed: int) -> None: - import numpy as np - - from grass_mil.inference.metrics import _survival_concordance_index - - rng = np.random.default_rng(seed) - n = 50 - risk = rng.normal(size=n) - time = rng.uniform(1.0, 100.0, size=n) - event = (rng.uniform(size=n) > 0.4).astype(float) - - got = _survival_concordance_index(torch.tensor(risk), torch.tensor(time), torch.tensor(event)) - expected = _reference_concordance_index(risk, time, event) - assert got == pytest.approx(expected, abs=1e-9) - - -def test_concordance_index_chunking_does_not_change_the_result() -> None: - import numpy as np - - from grass_mil.inference.metrics import _survival_concordance_index - - rng = np.random.default_rng(7) - n = 37 - risk = torch.tensor(rng.normal(size=n)) - time = torch.tensor(rng.uniform(1.0, 50.0, size=n)) - event = torch.tensor((rng.uniform(size=n) > 0.3).astype(float)) - - whole = _survival_concordance_index(risk, time, event, chunk_size=1024) - chunked = _survival_concordance_index(risk, time, event, chunk_size=4) - assert whole == pytest.approx(chunked, abs=1e-12) - - -def test_concordance_index_edge_cases() -> None: - from grass_mil.inference.metrics import _survival_concordance_index - - # No observed events: no comparable pairs at all. - assert math.isnan( - _survival_concordance_index( - torch.tensor([1.0, 2.0]), torch.tensor([1.0, 2.0]), torch.tensor([0.0, 0.0]) - ) - ) - # Perfect ordering: earlier event carries the higher risk. - assert _survival_concordance_index( - torch.tensor([3.0, 2.0, 1.0]), - torch.tensor([1.0, 2.0, 3.0]), - torch.tensor([1.0, 1.0, 1.0]), - ) == pytest.approx(1.0) - # Fully reversed ordering. - assert _survival_concordance_index( - torch.tensor([1.0, 2.0, 3.0]), - torch.tensor([1.0, 2.0, 3.0]), - torch.tensor([1.0, 1.0, 1.0]), - ) == pytest.approx(0.0) - # All risks tied: every comparable pair counts a half. - assert _survival_concordance_index( - torch.tensor([1.0, 1.0, 1.0]), - torch.tensor([1.0, 2.0, 3.0]), - torch.tensor([1.0, 1.0, 1.0]), - ) == pytest.approx(0.5) - - -def test_concordance_index_is_fast_for_large_cohorts() -> None: - """The old double loop synchronised the device once per pair.""" - import time as timing - - import numpy as np - - from grass_mil.inference.metrics import _survival_concordance_index - - rng = np.random.default_rng(0) - n = 4000 - risk = torch.tensor(rng.normal(size=n)) - time_col = torch.tensor(rng.uniform(1.0, 100.0, size=n)) - event = torch.tensor((rng.uniform(size=n) > 0.5).astype(float)) - - started = timing.perf_counter() - value = _survival_concordance_index(risk, time_col, event) - elapsed = timing.perf_counter() - started - - assert 0.0 <= value <= 1.0 - assert elapsed < 5.0 diff --git a/tests/eval/test_predict_cli.py b/tests/eval/test_predict_cli.py deleted file mode 100644 index 05f9384..0000000 --- a/tests/eval/test_predict_cli.py +++ /dev/null @@ -1,748 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from types import SimpleNamespace - -import pandas as pd -import pytest -import torch -from hydra.core.hydra_config import HydraConfig -from omegaconf import DictConfig, open_dict -from omegaconf import OmegaConf - -from grass_mil.inference.schemas import ( - AggregatedPredictionPayload, - BatchPredictionPayload, - CollectedInferencePayload, - EmbeddingPayload, -) -from grass_mil.inference.predict import ( - _needs_instance_payload, - _plan_preforward_subsampling, - predict, -) -from grass_mil.train import train - - -def test_predict_cli_smoke(cfg_train: DictConfig, cfg_predict: DictConfig, tmp_path: Path) -> None: - with open_dict(cfg_train): - cfg_train.trainer.max_epochs = 1 - cfg_train.test = True - HydraConfig().set_config(cfg_train) - _, _ = train(cfg_train) - - ckpt_path = tmp_path / "checkpoints" / "last.ckpt" - assert ckpt_path.exists() - - with open_dict(cfg_predict): - cfg_predict.ckpt_path = str(ckpt_path) - cfg_predict.embeddings.extract_node = False - cfg_predict.predict.output_subdir = "predict_artifacts" - HydraConfig().set_config(cfg_predict) - - summary, _ = predict(cfg_predict) - assert summary["row_count"] > 0 - - output_dir = tmp_path / "predict_artifacts" - assert (output_dir / "predictions.csv").exists() - assert (output_dir / "metrics.json").exists() - assert (output_dir / "graph_embeddings.csv").exists() - - -class _SamplerCfg: - def __init__(self, *, name: str, runtime: dict): - self.name = name - self.runtime = runtime - - -class _DataModule: - def __init__(self, *, sampler_name: str, runtime: dict): - self.sampler_config = _SamplerCfg(name=sampler_name, runtime=runtime) - self.sampler_strategy = object() - self.val_sampler_strategy = object() - - -def test_plan_preforward_subsampling_applies_for_enabled_runtime_sampler() -> None: - datamodule = _DataModule(sampler_name="shadow_native", runtime={"enabled": True}) - cfg = OmegaConf.create( - { - "aggregation": { - "enabled": True, - "mode": "mean", - "subsample_fraction": 0.5, - "subsample_seed": 123, - } - } - ) - decision = _plan_preforward_subsampling(datamodule, cfg) - assert decision.requested is True - assert decision.effective is True - assert decision.reason == "applied" - assert decision.strategy_name == "shadow_native" - assert decision.runtime_enabled is True - assert datamodule.sampler_config.runtime["subsample_fraction"] == 0.5 - assert datamodule.sampler_config.runtime["subsample_seed"] == 123 - assert datamodule.sampler_strategy is None - assert datamodule.val_sampler_strategy is None - - -def test_plan_preforward_subsampling_not_requested_for_full_fraction() -> None: - datamodule = _DataModule(sampler_name="shadow_native", runtime={"enabled": True}) - datamodule.sampler_strategy = "keep" - datamodule.val_sampler_strategy = "keep_val" - cfg = OmegaConf.create( - { - "aggregation": { - "enabled": True, - "mode": "max", - "subsample_fraction": 1.0, - "subsample_seed": 7, - } - } - ) - decision = _plan_preforward_subsampling(datamodule, cfg) - assert decision.requested is False - assert decision.effective is False - assert decision.reason == "not_requested" - assert "subsample_fraction" not in datamodule.sampler_config.runtime - assert datamodule.sampler_strategy == "keep" - assert datamodule.val_sampler_strategy == "keep_val" - - -def test_plan_preforward_subsampling_identity_sampler_is_ineffective() -> None: - datamodule = _DataModule(sampler_name="identity", runtime={"enabled": True}) - cfg = OmegaConf.create( - { - "aggregation": { - "enabled": True, - "mode": "mean", - "subsample_fraction": 0.5, - "subsample_seed": 7, - } - } - ) - decision = _plan_preforward_subsampling(datamodule, cfg) - assert decision.requested is True - assert decision.effective is False - assert decision.reason == "identity_sampler" - assert decision.strategy_name == "identity" - assert decision.runtime_enabled is None - - -def test_plan_preforward_subsampling_runtime_disabled_is_ineffective() -> None: - datamodule = _DataModule(sampler_name="shadow_custom", runtime={"enabled": False}) - cfg = OmegaConf.create( - { - "aggregation": { - "enabled": True, - "mode": "max", - "subsample_fraction": 0.5, - "subsample_seed": 11, - } - } - ) - decision = _plan_preforward_subsampling(datamodule, cfg) - assert decision.requested is True - assert decision.effective is False - assert decision.reason == "runtime_disabled" - assert decision.strategy_name == "shadow_custom" - assert decision.runtime_enabled is False - - -def test_plan_preforward_subsampling_unsupported_mode_is_ineffective() -> None: - datamodule = _DataModule(sampler_name="shadow_native", runtime={"enabled": True}) - cfg = OmegaConf.create( - { - "aggregation": { - "enabled": True, - "mode": "none", - "subsample_fraction": 0.5, - "subsample_seed": 5, - } - } - ) - decision = _plan_preforward_subsampling(datamodule, cfg) - assert decision.requested is True - assert decision.effective is False - assert decision.reason == "unsupported_mode" - - -@pytest.mark.parametrize( - ("cfg_dict", "expected"), - [ - ({"aggregation": {"enabled": False, "mode": "mean"}}, False), - ({"aggregation": {"enabled": True, "mode": "none"}}, False), - ({"aggregation": {"enabled": True, "mode": "mean"}}, True), - ({"aggregation": {"enabled": True, "mode": "max"}}, True), - ({"aggregation": {"enabled": True, "mode": "attention_weighted"}}, True), - ], -) -def test_needs_instance_payload(cfg_dict, expected) -> None: - cfg = OmegaConf.create(cfg_dict) - assert _needs_instance_payload(cfg) is expected - - -def test_predict_raises_when_preforward_subsampling_is_requested_but_ineffective( - monkeypatch, tmp_path: Path -) -> None: - import grass_mil.inference.predict as predict_module - - datamodule = _DataModule(sampler_name="identity", runtime={"enabled": True}) - - def _fake_instantiate(cfg, *args, **kwargs): - target = cfg.get("_target_") - if target == "fake.DataModule": - return datamodule - if target == "fake.Model": - return SimpleNamespace() - if target == "fake.Trainer": - return SimpleNamespace() - raise AssertionError(f"Unexpected instantiate target: {target}") - - monkeypatch.setattr(predict_module.hydra.utils, "instantiate", _fake_instantiate) - monkeypatch.setattr(predict_module, "instantiate_loggers", lambda cfg: []) - monkeypatch.setattr(predict_module, "log_hyperparameters", lambda obj: None) - - cfg = OmegaConf.create( - { - "ckpt_path": "dummy.ckpt", - "data": {"_target_": "fake.DataModule"}, - "model": {"_target_": "fake.Model"}, - "trainer": {"_target_": "fake.Trainer"}, - "logger": None, - "paths": {"output_dir": str(tmp_path)}, - "predict": { - "save_predictions": False, - "save_metrics": False, - "include_targets": False, - "include_attention": False, - "output_subdir": "predict_artifacts", - }, - "output": { - "predictions_filename": "predictions.csv", - "metrics_filename": "metrics.json", - "summary_filename": "inference_summary.json", - }, - "aggregation": { - "enabled": True, - "mode": "mean", - "bag_scope": "region", - "subsample_fraction": 0.5, - "subsample_seed": 17, - }, - "metrics": {"enabled": False, "per_group": False, "categorical": {"threshold": 0.5}}, - "embeddings": {"enabled": False, "save": False, "extract_node": False}, - "task": {"target_type": "binary"}, - } - ) - with pytest.raises(ValueError, match="Preforward subsampling was requested"): - predict(cfg) - - -def test_predict_disables_post_subsampling_only_when_preforward_is_effective( - monkeypatch, tmp_path: Path -) -> None: - import grass_mil.inference.predict as predict_module - - datamodule = _DataModule(sampler_name="shadow_native", runtime={"enabled": True}) - captured = {} - - def _fake_instantiate(cfg, *args, **kwargs): - target = cfg.get("_target_") - if target == "fake.DataModule": - return datamodule - if target == "fake.Model": - return SimpleNamespace() - if target == "fake.Trainer": - return SimpleNamespace() - raise AssertionError(f"Unexpected instantiate target: {target}") - - def _fake_aggregate(payload, **kwargs): - captured.update(kwargs) - return AggregatedPredictionPayload( - bag_ids=["r0"], - bag_logits=OmegaConf.create([[1.0]]), # type: ignore[arg-type] - bag_targets=None, - bag_attention=None, - metadata={"mode": "mean", "bag_scope": "region"}, - row_region_ids=["r0"], - row_sample_ids=["s0"], - ) - - monkeypatch.setattr(predict_module.hydra.utils, "instantiate", _fake_instantiate) - monkeypatch.setattr(predict_module, "instantiate_loggers", lambda cfg: []) - monkeypatch.setattr(predict_module, "log_hyperparameters", lambda obj: None) - monkeypatch.setattr( - predict_module, - "collect_inference_payload", - lambda **kwargs: CollectedInferencePayload( - prediction_payload=BatchPredictionPayload( - bag_ids=["r0"], - bag_logits=OmegaConf.create([[1.0]]), # type: ignore[arg-type] - bag_targets=None, - bag_attention=None, - row_region_ids=["r0"], - row_sample_ids=["s0"], - instance_logits=OmegaConf.create([[1.0]]), # type: ignore[arg-type] - instance_attention_logits=OmegaConf.create([[1.0]]), # type: ignore[arg-type] - instance_patch_ids=["p0"], - instance_region_ids=["r0"], - instance_sample_ids=["s0"], - ), - embedding_payload=None, - ), - ) - monkeypatch.setattr(predict_module, "aggregate_group_logits", _fake_aggregate) - monkeypatch.setattr( - predict_module, - "predictions_to_dataframe", - lambda payload, include_targets, include_attention: pd.DataFrame( - {"bag_id": ["r0"], "logit_0": [1.0]} - ), - ) - - cfg = OmegaConf.create( - { - "ckpt_path": "dummy.ckpt", - "data": {"_target_": "fake.DataModule"}, - "model": {"_target_": "fake.Model"}, - "trainer": {"_target_": "fake.Trainer"}, - "logger": None, - "paths": {"output_dir": str(tmp_path)}, - "predict": { - "save_predictions": False, - "save_metrics": False, - "include_targets": False, - "include_attention": False, - "output_subdir": "predict_artifacts", - }, - "output": { - "predictions_filename": "predictions.csv", - "metrics_filename": "metrics.json", - "summary_filename": "inference_summary.json", - }, - "aggregation": { - "enabled": True, - "mode": "mean", - "bag_scope": "region", - "subsample_fraction": 0.5, - "subsample_seed": 17, - }, - "metrics": {"enabled": False, "per_group": False, "categorical": {"threshold": 0.5}}, - "embeddings": {"enabled": False, "save": False, "extract_node": False}, - "task": {"target_type": "binary"}, - } - ) - summary, _ = predict(cfg) - - assert captured["subsample_fraction"] == 1.0 - assert captured["subsample_seed"] is None - assert summary["aggregation"]["preforward_subsampling_applied"] is True - assert summary["aggregation"]["preforward_subsampling"]["effective"] is True - - -def test_predict_writes_embeddings_from_single_pass_collected_payload( - monkeypatch, tmp_path: Path -) -> None: - import grass_mil.inference.predict as predict_module - - datamodule = _DataModule(sampler_name="shadow_native", runtime={"enabled": True}) - wrote_paths = [] - - def _fake_instantiate(cfg, *args, **kwargs): - target = cfg.get("_target_") - if target == "fake.DataModule": - return datamodule - if target == "fake.Model": - return SimpleNamespace() - if target == "fake.Trainer": - return SimpleNamespace() - raise AssertionError(f"Unexpected instantiate target: {target}") - - monkeypatch.setattr(predict_module.hydra.utils, "instantiate", _fake_instantiate) - monkeypatch.setattr(predict_module, "instantiate_loggers", lambda cfg: []) - monkeypatch.setattr(predict_module, "log_hyperparameters", lambda obj: None) - monkeypatch.setattr( - predict_module, - "collect_inference_payload", - lambda **kwargs: CollectedInferencePayload( - prediction_payload=BatchPredictionPayload( - bag_ids=["r0"], - bag_logits=OmegaConf.create([[1.0]]), # type: ignore[arg-type] - bag_targets=None, - bag_attention=None, - row_region_ids=["r0"], - row_sample_ids=["s0"], - ), - embedding_payload=EmbeddingPayload( - bag_ids=["r0"], - graph_embeddings=OmegaConf.create([[0.1, 0.2]]), # type: ignore[arg-type] - node_embeddings=OmegaConf.create([[0.3, 0.4]]), # type: ignore[arg-type] - node_bag_ids=["r0"], - ), - ), - ) - monkeypatch.setattr( - predict_module, - "aggregate_group_logits", - lambda *args, **kwargs: AggregatedPredictionPayload( - bag_ids=["r0"], - bag_logits=OmegaConf.create([[1.0]]), # type: ignore[arg-type] - bag_targets=None, - bag_attention=None, - metadata={"mode": "mean", "bag_scope": "region"}, - row_region_ids=["r0"], - row_sample_ids=["s0"], - ), - ) - monkeypatch.setattr( - predict_module, - "predictions_to_dataframe", - lambda payload, include_targets, include_attention: pd.DataFrame( - {"bag_id": ["r0"], "logit_0": [1.0]} - ), - ) - monkeypatch.setattr( - predict_module, - "embeddings_to_dataframe", - lambda payload: pd.DataFrame({"bag_id": payload.bag_ids, "graph_emb_0": [0.1]}), - ) - monkeypatch.setattr( - predict_module, - "node_embeddings_to_dataframe", - lambda payload: pd.DataFrame({"bag_id": payload.node_bag_ids, "node_emb_0": [0.3]}), - ) - - def _fake_write_dataframe(frame, path): - wrote_paths.append(str(path)) - - monkeypatch.setattr(predict_module, "write_dataframe", _fake_write_dataframe) - - cfg = OmegaConf.create( - { - "ckpt_path": "dummy.ckpt", - "data": {"_target_": "fake.DataModule"}, - "model": {"_target_": "fake.Model"}, - "trainer": {"_target_": "fake.Trainer"}, - "logger": None, - "paths": {"output_dir": str(tmp_path)}, - "predict": { - "save_predictions": False, - "save_metrics": False, - "include_targets": False, - "include_attention": False, - "output_subdir": "predict_artifacts", - }, - "output": { - "predictions_filename": "predictions.csv", - "metrics_filename": "metrics.json", - "summary_filename": "inference_summary.json", - }, - "aggregation": { - "enabled": True, - "mode": "mean", - "bag_scope": "region", - "subsample_fraction": 1.0, - "subsample_seed": None, - }, - "metrics": {"enabled": False, "per_group": False, "categorical": {"threshold": 0.5}}, - "embeddings": { - "enabled": True, - "save": True, - "extract_node": True, - "filename": "graph_embeddings.csv", - "node_filename": "node_embeddings.csv", - }, - "task": {"target_type": "binary"}, - } - ) - - summary, _ = predict(cfg) - assert summary["embeddings_path"] is not None - assert any(path.endswith("graph_embeddings.csv") for path in wrote_paths) - assert any(path.endswith("node_embeddings.csv") for path in wrote_paths) - - -def test_predict_summary_includes_aggregation_metadata(monkeypatch, tmp_path: Path) -> None: - import grass_mil.inference.predict as predict_module - - class _FakeDataModule: - pass - - class _FakeModel: - pass - - class _FakeTrainer: - pass - - fake_datamodule = _FakeDataModule() - fake_model = _FakeModel() - fake_trainer = _FakeTrainer() - - def _fake_instantiate(cfg, *args, **kwargs): - target = cfg.get("_target_") - if target == "fake.DataModule": - return fake_datamodule - if target == "fake.Model": - return fake_model - if target == "fake.Trainer": - return fake_trainer - raise AssertionError(f"Unexpected instantiate target: {target}") - - captured_summary_payload = {} - - def _fake_write_json(payload, path): - captured_summary_payload.clear() - captured_summary_payload.update(payload) - - monkeypatch.setattr(predict_module.hydra.utils, "instantiate", _fake_instantiate) - monkeypatch.setattr(predict_module, "instantiate_loggers", lambda cfg: []) - monkeypatch.setattr(predict_module, "log_hyperparameters", lambda obj: None) - monkeypatch.setattr( - predict_module, - "collect_inference_payload", - lambda **kwargs: CollectedInferencePayload( - prediction_payload=BatchPredictionPayload( - bag_ids=["r0"], - bag_logits=OmegaConf.create([[1.0]]), # type: ignore[arg-type] - bag_targets=None, - bag_attention=None, - row_region_ids=["r0"], - row_sample_ids=["s0"], - ), - embedding_payload=None, - ), - ) - monkeypatch.setattr( - predict_module, - "aggregate_group_logits", - lambda *args, **kwargs: AggregatedPredictionPayload( - bag_ids=["r0"], - bag_logits=OmegaConf.create([[1.0]]), # type: ignore[arg-type] - bag_targets=None, - bag_attention=None, - metadata={"mode": "mean", "bag_scope": "region"}, - row_region_ids=["r0"], - row_sample_ids=["s0"], - ), - ) - monkeypatch.setattr( - predict_module, - "predictions_to_dataframe", - lambda payload, include_targets, include_attention: pd.DataFrame( - {"bag_id": ["r0"], "logit_0": [1.0]} - ), - ) - monkeypatch.setattr(predict_module, "write_json", _fake_write_json) - - cfg = OmegaConf.create( - { - "ckpt_path": "dummy.ckpt", - "data": {"_target_": "fake.DataModule"}, - "model": {"_target_": "fake.Model"}, - "trainer": {"_target_": "fake.Trainer"}, - "logger": None, - "paths": {"output_dir": str(tmp_path)}, - "predict": { - "save_predictions": False, - "save_metrics": False, - "include_targets": False, - "include_attention": False, - "output_subdir": "predict_artifacts", - }, - "output": { - "predictions_filename": "predictions.csv", - "metrics_filename": "metrics.json", - "summary_filename": "inference_summary.json", - }, - "aggregation": { - "enabled": True, - "mode": "mean", - "bag_scope": "region", - "subsample_fraction": 1.0, - "subsample_seed": None, - }, - "metrics": {"enabled": False, "per_group": False, "categorical": {"threshold": 0.5}}, - "embeddings": {"enabled": False, "save": False, "extract_node": False}, - "task": {"target_type": "binary"}, - } - ) - - summary, _ = predict(cfg) - assert "aggregation" in summary - assert summary["aggregation"]["mode"] == "mean" - assert summary["aggregation"]["bag_scope"] == "region" - assert "preforward_subsampling_applied" in summary["aggregation"] - assert "preforward_subsampling" in summary["aggregation"] - assert summary["aggregation"]["preforward_subsampling"]["reason"] == "not_requested" - assert captured_summary_payload["aggregation"]["mode"] == "mean" - - -def test_predict_raises_when_ckpt_path_missing(tmp_path: Path) -> None: - cfg = OmegaConf.create( - { - "ckpt_path": "", - "data": {"_target_": "fake.DataModule"}, - "model": {"_target_": "fake.Model"}, - "trainer": {"_target_": "fake.Trainer"}, - "logger": None, - "paths": {"output_dir": str(tmp_path)}, - "predict": { - "save_predictions": False, - "save_metrics": False, - "include_targets": False, - "include_attention": False, - "output_subdir": "predict_artifacts", - }, - "output": { - "predictions_filename": "predictions.csv", - "metrics_filename": "metrics.json", - "summary_filename": "inference_summary.json", - }, - "aggregation": {"enabled": False, "mode": "none"}, - "metrics": {"enabled": False, "per_group": False, "categorical": {"threshold": 0.5}}, - "embeddings": {"enabled": False, "save": False, "extract_node": False}, - "task": {"target_type": "binary"}, - } - ) - with pytest.raises(ValueError, match="Missing required `ckpt_path` for inference."): - predict(cfg) - - -def test_predict_can_export_interpretability_tables(monkeypatch, tmp_path: Path) -> None: - import grass_mil.inference.predict as predict_module - - class _FakeDataModule: - class _Dataset: - label_maps = {"cell_type": {"A": 0, "B": 1}} - - dataset_train = _Dataset() - - class _FakeModel: - pass - - class _FakeTrainer: - pass - - fake_datamodule = _FakeDataModule() - fake_model = _FakeModel() - fake_trainer = _FakeTrainer() - captured = {"collector_kwargs": None, "frames": {}} - - def _fake_instantiate(cfg, *args, **kwargs): - target = cfg.get("_target_") - if target == "fake.DataModule": - return fake_datamodule - if target == "fake.Model": - return fake_model - if target == "fake.Trainer": - return fake_trainer - raise AssertionError(f"Unexpected instantiate target: {target}") - - def _fake_collect(**kwargs): - captured["collector_kwargs"] = kwargs - return CollectedInferencePayload( - prediction_payload=BatchPredictionPayload( - bag_ids=["r0"], - bag_logits=OmegaConf.create([[1.0]]), # type: ignore[arg-type] - bag_targets=None, - bag_attention=None, - row_region_ids=["r0"], - row_sample_ids=["s0"], - instance_logits=torch.tensor([[0.2], [0.4]]), - instance_attention_logits=torch.tensor([[0.1], [0.2]]), - instance_patch_ids=["p0", "p1"], - instance_bag_ids=["s0::r0", "s0::r0"], - instance_region_ids=["r0", "r0"], - instance_sample_ids=["s0", "s0"], - instance_embeddings=torch.tensor([[0.11, 0.12], [0.21, 0.22]]), - instance_composition=torch.tensor([[1.0, 0.0], [0.0, 1.0]]), - instance_centroids=torch.tensor([[0.0, 0.0], [1.0, 0.0]]), - instance_graphs=[ - SimpleNamespace( - root_n_id=torch.tensor([0], dtype=torch.long), - n_id=torch.tensor([100, 101], dtype=torch.long), - edge_index=torch.tensor([[0], [1]], dtype=torch.long), - edge_attr=torch.tensor([[1.0]], dtype=torch.float32), - edge_attr_names=["distance"], - ), - SimpleNamespace( - root_n_id=torch.tensor([0], dtype=torch.long), - n_id=torch.tensor([101, 100], dtype=torch.long), - edge_index=torch.tensor([[0], [1]], dtype=torch.long), - edge_attr=torch.tensor([[1.0]], dtype=torch.float32), - edge_attr_names=["distance"], - ), - ], - ), - embedding_payload=None, - ) - - def _fake_write_dataframe(frame, path): - captured["frames"][str(path)] = frame.copy() - - monkeypatch.setattr(predict_module.hydra.utils, "instantiate", _fake_instantiate) - monkeypatch.setattr(predict_module, "instantiate_loggers", lambda cfg: []) - monkeypatch.setattr(predict_module, "log_hyperparameters", lambda obj: None) - monkeypatch.setattr(predict_module, "collect_inference_payload", _fake_collect) - monkeypatch.setattr( - predict_module, - "predictions_to_dataframe", - lambda payload, include_targets, include_attention: pd.DataFrame( - {"bag_id": ["r0"], "logit_0": [1.0]} - ), - ) - monkeypatch.setattr(predict_module, "write_dataframe", _fake_write_dataframe) - - cfg = OmegaConf.create( - { - "ckpt_path": "dummy.ckpt", - "data": {"_target_": "fake.DataModule"}, - "model": {"_target_": "fake.Model"}, - "trainer": {"_target_": "fake.Trainer"}, - "logger": None, - "paths": {"output_dir": str(tmp_path)}, - "predict": { - "save_predictions": False, - "save_metrics": False, - "include_targets": False, - "include_attention": False, - "output_subdir": "predict_artifacts", - }, - "output": { - "predictions_filename": "predictions.csv", - "metrics_filename": "metrics.json", - "summary_filename": "inference_summary.json", - }, - "aggregation": {"enabled": False, "mode": "none", "subsample_fraction": 1.0}, - "metrics": {"enabled": False, "per_group": False, "categorical": {"threshold": 0.5}}, - "embeddings": {"enabled": False, "save": False, "extract_node": False}, - "interpretability": { - "enabled": True, - "instance_filename": "instance_table.csv", - "spatial_filename": "spatial_table.csv", - "id_column": "instance_id", - "bag_id_column": "bag_id", - "composition_label": "cell_type", - "composition_prefix": "comp_", - "embedding_prefix": "inst_emb", - "score_column": "score", - "attention_column": "attention", - "require_composition": True, - "spatial": { - "enabled": True, - "undirected": True, - }, - }, - "task": {"target_type": "binary"}, - } - ) - - summary, _ = predict(cfg) - assert captured["collector_kwargs"]["include_instance_payload"] is True - assert captured["collector_kwargs"]["include_instance_embeddings"] is True - assert "interpretability" in summary - assert summary["interpretability"]["instance_table_path"].endswith("instance_table.csv") - assert summary["interpretability"]["spatial_table_path"].endswith("spatial_table.csv") - assert any(path.endswith("instance_table.csv") for path in captured["frames"].keys()) - assert any(path.endswith("spatial_table.csv") for path in captured["frames"].keys()) diff --git a/tests/helpers/cohort.py b/tests/helpers/cohort.py new file mode 100644 index 0000000..23cc7fd --- /dev/null +++ b/tests/helpers/cohort.py @@ -0,0 +1,149 @@ +"""A small synthetic cohort in the public data contract, and config helpers. + +Regions are CSV cell tables (``cell_id, x, y, cell_type, area``) listed in a +manifest; labels live in a CSV with ``region_id``, ``label``, ``patient_id``, +``time`` and ``event``; outer folds are a JSON list of region-id lists grouped +by patient. Positive regions carry a planted niche: cell type ``D`` packed +around a ``C`` cell; in negative regions ``D`` is scattered. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Sequence + +import numpy as np +import pandas as pd +from hydra import compose, initialize_config_module +from hydra.core.global_hydra import GlobalHydra +from hydra.core.hydra_config import HydraConfig +from omegaconf import DictConfig, open_dict + +CELL_TYPES = ["A", "B", "C", "D"] + + +@dataclass +class Cohort: + root: Path + manifest: Path + labels: Path + folds: Path + regions: list[str] + label_of: dict[str, int] + patient_of: dict[str, str] + + +def _cells(rng: np.random.Generator, n: int, positive: bool, extent: float) -> pd.DataFrame: + xy = rng.uniform(0.0, extent, size=(n, 2)) + types = rng.choice(CELL_TYPES[:3], size=n, p=[0.45, 0.35, 0.20]) + d = rng.choice(n, size=max(8, n // 10), replace=False) + types[d] = "D" + if positive: + centre = xy[rng.choice(np.flatnonzero(types == "C"))] + xy[d] = np.clip(centre + rng.normal(0.0, extent * 0.04, size=(d.size, 2)), 0.0, extent) + return pd.DataFrame( + { + "cell_id": [f"c{i}" for i in range(n)], + "x": xy[:, 0], + "y": xy[:, 1], + "cell_type": types, + "area": rng.lognormal(4.0, 0.3, size=n), + } + ) + + +def make_cohort( + root: Path, + *, + n_patients: int = 6, + regions_per_patient: int = 2, + n_cells: int = 120, + n_folds: int = 3, + seed: int = 0, +) -> Cohort: + raw = root / "raw" + raw.mkdir(parents=True, exist_ok=True) + rng = np.random.default_rng(seed) + manifest, labels = [], [] + label_of, patient_of = {}, {} + for p in range(n_patients): + patient, label = f"P{p:02d}", p % 2 + for r in range(regions_per_patient): + region = f"{patient}_core{r}" + _cells(rng, n_cells, bool(label), 600.0).to_csv(raw / f"{region}.csv", index=False) + manifest.append( + {"sample_id": region, "input_path": f"{region}.csv", "region_id": region} + ) + labels.append( + { + "region_id": region, + "label": label, + "patient_id": patient, + "time": float(rng.exponential(10.0 if label else 30.0)), + "event": int(rng.random() < 0.8), + "score": float(label + rng.normal(0, 0.3)), + } + ) + label_of[region], patient_of[region] = label, patient + pd.DataFrame(manifest).to_csv(raw / "manifest.csv", index=False) + pd.DataFrame(labels).to_csv(raw / "labels.csv", index=False) + patients = sorted(set(patient_of.values())) + folds: list[list[str]] = [[] for _ in range(n_folds)] + for i, patient in enumerate(patients): + folds[(i // 2) % n_folds].extend(r for r in label_of if patient_of[r] == patient) + (raw / "folds.json").write_text(json.dumps(folds)) + return Cohort( + root, + raw / "manifest.csv", + raw / "labels.csv", + raw / "folds.json", + list(label_of), + label_of, + patient_of, + ) + + +def cohort_overrides( + cohort: Cohort, processed: str = "processed", *, finetune: bool = True +) -> list[str]: + """Overrides pointing the packaged configs at the cohort, with a tiny model.""" + small = ["finetune.hyperbatch_size=3", "finetune.probe.roots_per_unit=16"] if finetune else [] + return small + [ + f"data.manifest_path={cohort.manifest}", + f"data.processed_dir={cohort.root / processed}", + "data.input.cell_id_column=cell_id", + "data.input.features=[area]", + "data.graph.neighbor_cutoff=55", + f"data.labels.file={cohort.labels}", + "data.labels.patient_column=patient_id", + "task.target_columns=[label]", + "data.sampler.depth=2", + "data.sampler.roots_per_unit=24", + "data.sampler.batch_size=16", + "data.val_roots_per_unit=32", + "model.encoder.num_layers=2", + "model.encoder.hidden_dim=16", + "model.encoder.out_dim=16", + "trainer.max_epochs=2", + "trainer.min_epochs=1", + "+trainer.enable_progress_bar=false", + "+trainer.enable_model_summary=false", + "~callbacks.rich_progress_bar", + "~callbacks.model_summary", + ] + + +def compose_cfg(name: str, overrides: Sequence[str], output_dir: Path | None = None) -> DictConfig: + """Compose a packaged config as its entry point would (Hydra runtime set).""" + GlobalHydra.instance().clear() + with initialize_config_module(version_base="1.3", config_module="grass_mil.configs"): + cfg = compose(config_name=name, overrides=list(overrides), return_hydra_config=True) + HydraConfig.instance().set_config(cfg) + if output_dir is not None: + with open_dict(cfg): + cfg.paths.output_dir = str(output_dir) + if "trainer" in cfg: + cfg.trainer.default_root_dir = str(output_dir) + return cfg diff --git a/tests/helpers/config_paths.py b/tests/helpers/config_paths.py deleted file mode 100644 index 57f9b60..0000000 --- a/tests/helpers/config_paths.py +++ /dev/null @@ -1,12 +0,0 @@ -"""Locate the packaged Hydra config tree. - -Configs ship inside the package (``grass_mil/configs``) so that they resolve -identically from a source checkout and from an installed wheel. Tests must not -reach for them via a CWD-relative path. -""" - -from pathlib import Path - -from grass_mil import configs as _configs - -CONFIGS_DIR = Path(_configs.__file__).parent diff --git a/tests/helpers/package_available.py b/tests/helpers/package_available.py deleted file mode 100644 index 0afdba8..0000000 --- a/tests/helpers/package_available.py +++ /dev/null @@ -1,32 +0,0 @@ -import platform - -import pkg_resources -from lightning.fabric.accelerators import TPUAccelerator - - -def _package_available(package_name: str) -> bool: - """Check if a package is available in your environment. - - :param package_name: The name of the package to be checked. - - :return: `True` if the package is available. `False` otherwise. - """ - try: - return pkg_resources.require(package_name) is not None - except pkg_resources.DistributionNotFound: - return False - - -_TPU_AVAILABLE = TPUAccelerator.is_available() - -_IS_WINDOWS = platform.system() == "Windows" - -_SH_AVAILABLE = not _IS_WINDOWS and _package_available("sh") - -_DEEPSPEED_AVAILABLE = not _IS_WINDOWS and _package_available("deepspeed") -_FAIRSCALE_AVAILABLE = not _IS_WINDOWS and _package_available("fairscale") - -_WANDB_AVAILABLE = _package_available("wandb") -_NEPTUNE_AVAILABLE = _package_available("neptune") -_COMET_AVAILABLE = _package_available("comet_ml") -_MLFLOW_AVAILABLE = _package_available("mlflow") diff --git a/tests/helpers/run_if.py b/tests/helpers/run_if.py deleted file mode 100644 index 9703af4..0000000 --- a/tests/helpers/run_if.py +++ /dev/null @@ -1,142 +0,0 @@ -"""Adapted from: - -https://github.com/PyTorchLightning/pytorch-lightning/blob/master/tests/helpers/runif.py -""" - -import sys -from typing import Any, Dict, Optional - -import pytest -import torch -from packaging.version import Version -from pkg_resources import get_distribution -from pytest import MarkDecorator - -from tests.helpers.package_available import ( - _COMET_AVAILABLE, - _DEEPSPEED_AVAILABLE, - _FAIRSCALE_AVAILABLE, - _IS_WINDOWS, - _MLFLOW_AVAILABLE, - _NEPTUNE_AVAILABLE, - _SH_AVAILABLE, - _TPU_AVAILABLE, - _WANDB_AVAILABLE, -) - - -class RunIf: - """RunIf wrapper for conditional skipping of tests. - - Fully compatible with `@pytest.mark`. - - Example: - - ```python - @RunIf(min_torch="1.8") - @pytest.mark.parametrize("arg1", [1.0, 2.0]) - def test_wrapper(arg1): - assert arg1 > 0 - ``` - """ - - def __new__( - cls, - min_gpus: int = 0, - min_torch: Optional[str] = None, - max_torch: Optional[str] = None, - min_python: Optional[str] = None, - skip_windows: bool = False, - sh: bool = False, - tpu: bool = False, - fairscale: bool = False, - deepspeed: bool = False, - wandb: bool = False, - neptune: bool = False, - comet: bool = False, - mlflow: bool = False, - **kwargs: Dict[Any, Any], - ) -> MarkDecorator: - """Creates a new `@RunIf` `MarkDecorator` decorator. - - :param min_gpus: Min number of GPUs required to run test. - :param min_torch: Minimum pytorch version to run test. - :param max_torch: Maximum pytorch version to run test. - :param min_python: Minimum python version required to run test. - :param skip_windows: Skip test for Windows platform. - :param tpu: If TPU is available. - :param sh: If `sh` module is required to run the test. - :param fairscale: If `fairscale` module is required to run the test. - :param deepspeed: If `deepspeed` module is required to run the test. - :param wandb: If `wandb` module is required to run the test. - :param neptune: If `neptune` module is required to run the test. - :param comet: If `comet` module is required to run the test. - :param mlflow: If `mlflow` module is required to run the test. - :param kwargs: Native `pytest.mark.skipif` keyword arguments. - """ - conditions = [] - reasons = [] - - if min_gpus: - conditions.append(torch.cuda.device_count() < min_gpus) - reasons.append(f"GPUs>={min_gpus}") - - if min_torch: - torch_version = get_distribution("torch").version - conditions.append(Version(torch_version) < Version(min_torch)) - reasons.append(f"torch>={min_torch}") - - if max_torch: - torch_version = get_distribution("torch").version - conditions.append(Version(torch_version) >= Version(max_torch)) - reasons.append(f"torch<{max_torch}") - - if min_python: - py_version = ( - f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}" - ) - conditions.append(Version(py_version) < Version(min_python)) - reasons.append(f"python>={min_python}") - - if skip_windows: - conditions.append(_IS_WINDOWS) - reasons.append("does not run on Windows") - - if tpu: - conditions.append(not _TPU_AVAILABLE) - reasons.append("TPU") - - if sh: - conditions.append(not _SH_AVAILABLE) - reasons.append("sh") - - if fairscale: - conditions.append(not _FAIRSCALE_AVAILABLE) - reasons.append("fairscale") - - if deepspeed: - conditions.append(not _DEEPSPEED_AVAILABLE) - reasons.append("deepspeed") - - if wandb: - conditions.append(not _WANDB_AVAILABLE) - reasons.append("wandb") - - if neptune: - conditions.append(not _NEPTUNE_AVAILABLE) - reasons.append("neptune") - - if comet: - conditions.append(not _COMET_AVAILABLE) - reasons.append("comet") - - if mlflow: - conditions.append(not _MLFLOW_AVAILABLE) - reasons.append("mlflow") - - reasons = [rs for cond, rs in zip(conditions, reasons) if cond] - return pytest.mark.skipif( - condition=any(conditions), - reason=f"Requires: [{' + '.join(reasons)}]", - **kwargs, - ) diff --git a/tests/helpers/run_sh_command.py b/tests/helpers/run_sh_command.py deleted file mode 100644 index fdd2ed6..0000000 --- a/tests/helpers/run_sh_command.py +++ /dev/null @@ -1,22 +0,0 @@ -from typing import List - -import pytest - -from tests.helpers.package_available import _SH_AVAILABLE - -if _SH_AVAILABLE: - import sh - - -def run_sh_command(command: List[str]) -> None: - """Default method for executing shell commands with `pytest` and `sh` package. - - :param command: A list of shell commands as strings. - """ - msg = None - try: - sh.python(command) - except sh.ErrorReturnCode as e: - msg = e.stderr.decode() - if msg: - pytest.fail(msg=msg) diff --git a/tests/helpers/synthetic_cohort.py b/tests/helpers/synthetic_cohort.py deleted file mode 100644 index 3b3c025..0000000 --- a/tests/helpers/synthetic_cohort.py +++ /dev/null @@ -1,176 +0,0 @@ -"""A small synthetic cohort on disk, in the public data contract. - -Regions are CSV cell tables (``cell_id, x, y, cell_type, area``) listed in a -manifest; labels live in a CSV with a region id, a target and a patient column; -outer folds are a JSON list of region-id lists grouped by patient. Positive -regions carry a planted niche: cell type ``D`` forms a compact cluster next to -``C`` cells, while in negatives ``D`` is scattered. -""" - -from __future__ import annotations - -import json -import os -from dataclasses import dataclass, field -from pathlib import Path -from typing import Dict, List, Sequence - -import numpy as np -import pandas as pd -from hydra import compose, initialize_config_module -from hydra.core.global_hydra import GlobalHydra -from omegaconf import DictConfig, open_dict - -CELL_TYPES = ["A", "B", "C", "D"] - - -@dataclass -class SyntheticCohort: - root: Path - manifest_path: Path - label_path: Path - folds_path: Path - processed_dir: Path - region_ids: List[str] - patient_of: Dict[str, str] - label_of: Dict[str, int] - folds: List[List[str]] = field(default_factory=list) - - -def _region_cells(rng: np.random.Generator, n_cells: int, positive: bool, extent: float): - xy = rng.uniform(0.0, extent, size=(n_cells, 2)) - types = rng.choice(CELL_TYPES[:3], size=n_cells, p=[0.45, 0.35, 0.20]) - n_d = max(8, n_cells // 10) - d_idx = rng.choice(n_cells, size=n_d, replace=False) - types[d_idx] = "D" - if positive: - # Planted niche: D cells packed around a C cell. - c_idx = np.where(types == "C")[0] - centre = xy[rng.choice(c_idx)] if c_idx.size else xy[0] - xy[d_idx] = centre + rng.normal(0.0, extent * 0.04, size=(n_d, 2)) - xy = np.clip(xy, 0.0, extent) - area = rng.lognormal(mean=4.0, sigma=0.3, size=n_cells) - return pd.DataFrame( - { - "cell_id": [f"c{i}" for i in range(n_cells)], - "x": xy[:, 0], - "y": xy[:, 1], - "cell_type": types, - "area": area, - } - ) - - -def make_synthetic_cohort( - root: str | Path, - *, - n_patients: int = 6, - regions_per_patient: int = 2, - n_cells: int = 150, - n_outer_folds: int = 3, - extent: float = 600.0, - seed: int = 0, -) -> SyntheticCohort: - root = Path(root) - raw = root / "raw" - raw.mkdir(parents=True, exist_ok=True) - rng = np.random.default_rng(seed) - - rows = [] - label_rows = [] - region_ids: List[str] = [] - patient_of: Dict[str, str] = {} - label_of: Dict[str, int] = {} - for p in range(n_patients): - patient = f"P{p:02d}" - label = int(p % 2 == 1) - for r in range(regions_per_patient): - rid = f"{patient}_core{r}" - df = _region_cells(rng, n_cells, positive=bool(label), extent=extent) - path = raw / f"{rid}.csv" - df.to_csv(path, index=False) - rows.append( - {"sample_id": rid, "input_path": str(path), "input_type": "csv", "region_id": rid} - ) - label_rows.append({"REGION_ID": rid, "STAS": label, "PATIENT": patient}) - region_ids.append(rid) - patient_of[rid] = patient - label_of[rid] = label - manifest_path = raw / "manifest.csv" - pd.DataFrame(rows).to_csv(manifest_path, index=False) - label_path = raw / "labels.csv" - pd.DataFrame(label_rows).to_csv(label_path, index=False) - - # Patient-grouped outer folds: consecutive (negative, positive) patient pairs - # are dealt round-robin, so every fold holds both classes. - patients = sorted(set(patient_of.values())) - folds: List[List[str]] = [[] for _ in range(n_outer_folds)] - for i, patient in enumerate(patients): - folds[(i // 2) % n_outer_folds].extend(r for r in region_ids if patient_of[r] == patient) - folds_path = raw / "folds.json" - folds_path.write_text(json.dumps(folds)) - - return SyntheticCohort( - root=root, - manifest_path=manifest_path, - label_path=label_path, - folds_path=folds_path, - processed_dir=root / "processed", - region_ids=region_ids, - patient_of=patient_of, - label_of=label_of, - folds=folds, - ) - - -def cohort_overrides(cohort: SyntheticCohort, output_dir: str | Path) -> List[str]: - """Hydra overrides that point the public configs at the synthetic cohort.""" - return [ - "regime=mil", - f"paths.output_dir={output_dir}", - f"paths.root_dir={cohort.root}", - f"data.raw_manifest_path={cohort.manifest_path}", - f"data.processed_dir={cohort.processed_dir}", - "data.sample_unit=full", - "data.min_cells=5", - "data.csv.categorical_label_columns=[cell_type]", - "data.csv.molecular_columns=[area]", - "data.graph_builder.kwargs.edge_features=[log_distance,neighbor]", - "data.graph_builder.kwargs.neighbor_cutoff_um=55.0", - "data.split.split_by=region", - "data.split.train_val_test_split=[0.6,0.2,0.2]", - f"data.graph_labels.label_file={cohort.label_path}", - "data.graph_labels.id_column=REGION_ID", - "data.graph_labels.tasks=[STAS]", - "data.graph_labels.patient_column=PATIENT", - "data.sampler.runtime.depth=2", - "data.sampler.runtime.subsample_count=24", - "data.val_sampler.runtime.subsample_count=32", - "model.encoder.hidden_dim=16", - "model.encoder.out_dim=16", - "model.encoder.num_layers=2", - "model.attention.projection_dim=8", - "model.attention.hidden_dim=4", - "task.region_accumulation.hyperbatch_size=4", - "trainer.max_epochs=2", - "trainer.min_epochs=1", - "trainer.accelerator=cpu", - "trainer.devices=1", - "trainer.precision=32", - "+trainer.enable_progress_bar=false", - "~callbacks.rich_progress_bar", - "extras.print_config=false", - "extras.enforce_tags=false", - f"seed={os.environ.get('GRASS_MIL_TEST_SEED', '0')}", - ] - - -def compose_cfg(config_name: str, overrides: Sequence[str]) -> DictConfig: - GlobalHydra.instance().clear() - with initialize_config_module(version_base="1.3", config_module="grass_mil.configs"): - cfg = compose(config_name=config_name, overrides=list(overrides)) - GlobalHydra.instance().clear() - if "trainer" in cfg: - with open_dict(cfg): - cfg.trainer.default_root_dir = cfg.paths.output_dir - return cfg diff --git a/tests/helpers/synthetic_datamodule.py b/tests/helpers/synthetic_datamodule.py deleted file mode 100644 index 6585981..0000000 --- a/tests/helpers/synthetic_datamodule.py +++ /dev/null @@ -1,125 +0,0 @@ -from __future__ import annotations - -import lightning as L -import torch - -try: - from torch_geometric.data import Data - from torch_geometric.loader import DataLoader -except Exception as exc: # pragma: no cover - raise ImportError("torch_geometric is required for synthetic test datamodule") from exc - - -def _make_graph( - region_id: str, - label: float, - num_nodes: int = 6, - input_dim: int = 8, - patch_idx: int = 0, - survival: bool = False, -) -> Data: - x = torch.randn(num_nodes, input_dim) - edge_index = torch.tensor( - [[0, 1, 2, 3, 4, 1, 2, 3, 4, 5], [1, 2, 3, 4, 5, 0, 1, 2, 3, 4]], - dtype=torch.long, - ) - if survival: - # [time, event]; order is load-bearing for the Cox loss. Times are - # spread and events mixed so risk sets are non-degenerate -- if every - # event sat at the largest time the partial likelihood would be - # constant and carry no gradient. - follow_up = 1.0 + float(patch_idx % 7) + 3.0 * label - observed = float((patch_idx % 3) != 0) - graph_y = torch.tensor([[follow_up, observed]], dtype=torch.float32) - else: - graph_y = torch.tensor([[label]], dtype=torch.float32) - graph_w = torch.ones((1, 1), dtype=torch.float32) - data = Data(x=x, edge_index=edge_index, graph_y=graph_y, graph_w=graph_w) - data.sample_id = "sample_a" - data.region_id = region_id - data.patch_id = f"{region_id}_patch_{patch_idx}" - data.graph_label_names = ["time", "event"] if survival else ["label"] - return data - - -class SyntheticBagDataModule(L.LightningDataModule): - def __init__( - self, - batch_size: int = 4, - num_workers: int = 0, - pin_memory: bool = False, - input_dim: int = 8, - graphs_per_region: int = 20, - survival: bool = False, - ) -> None: - super().__init__() - self.batch_size = batch_size - self.num_workers = num_workers - self.pin_memory = pin_memory - self.input_dim = input_dim - self.graphs_per_region = graphs_per_region - self.survival = survival - self.dataset_train = None - self.dataset_val = None - self.dataset_test = None - - def setup(self, stage=None): - data = [] - for i in range(self.graphs_per_region): - data.append( - _make_graph( - "region_0", - 0.0, - input_dim=self.input_dim, - patch_idx=2 * i, - survival=self.survival, - ) - ) - data.append( - _make_graph( - "region_1", - 1.0, - input_dim=self.input_dim, - patch_idx=2 * i + 1, - survival=self.survival, - ) - ) - self.dataset_train = data - self.dataset_val = data - self.dataset_test = data - - def train_dataloader(self): - return DataLoader( - self.dataset_train, - batch_size=self.batch_size, - num_workers=self.num_workers, - pin_memory=self.pin_memory, - shuffle=True, - ) - - def val_dataloader(self): - return DataLoader( - self.dataset_val, - batch_size=self.batch_size, - num_workers=self.num_workers, - pin_memory=self.pin_memory, - shuffle=False, - ) - - def test_dataloader(self): - return DataLoader( - self.dataset_test, - batch_size=self.batch_size, - num_workers=self.num_workers, - pin_memory=self.pin_memory, - shuffle=False, - ) - - def predict_dataloader(self): - return DataLoader( - self.dataset_test, - batch_size=self.batch_size, - num_workers=self.num_workers, - pin_memory=self.pin_memory, - shuffle=False, - ) diff --git a/tests/interpretability/test_attribution.py b/tests/interpretability/test_attribution.py deleted file mode 100644 index 923cfe3..0000000 --- a/tests/interpretability/test_attribution.py +++ /dev/null @@ -1,325 +0,0 @@ -"""Contracts for the exact additive attribution. - -The load-bearing property is the identity sum_i A[i,c] * l[i,c] == L_c. If it -does not hold exactly, every niche-level attribution built on it is invalid, -so it is asserted directly rather than inferred. -""" - -from __future__ import annotations - -import numpy as np -import pandas as pd -import pytest - -from grass_mil.interpretability.core.attribution import ( - niche_attribution_summary, - instance_margin_contributions, - instance_ovr_margins, - percentile_bootstrap_ci, -) - - -def _instance_table(n_bags: int = 4, per_bag: int = 6, seed: int = 0): - rng = np.random.default_rng(seed) - rows = [] - for bag in range(n_bags): - attn_logits = rng.normal(size=(per_bag, 2)) - attention = np.exp(attn_logits) / np.exp(attn_logits).sum(axis=0, keepdims=True) - logits = rng.normal(size=(per_bag, 2)) - for i in range(per_bag): - rows.append( - { - "bag_id": f"bag_{bag}", - "attention_c0": attention[i, 0], - "attention_c1": attention[i, 1], - "logit_0": logits[i, 0], - "logit_1": logits[i, 1], - } - ) - table = pd.DataFrame(rows) - labels = np.array([i % 3 for i in range(len(table))]) - return table, labels - - -def test_contributions_sum_exactly_to_the_bag_logit() -> None: - table, labels = _instance_table() - - bag_logits = {} - for bag_id, group in table.groupby("bag_id"): - attention = group[["attention_c0", "attention_c1"]].to_numpy() - logits = group[["logit_0", "logit_1"]].to_numpy() - bag_logits[bag_id] = (attention * logits).sum(axis=0) - - result = niche_attribution_summary( - table, - labels, - attention_columns=["attention_c0", "attention_c1"], - logit_columns=["logit_0", "logit_1"], - bag_logits=bag_logits, - n_bootstrap=20, - ) - assert result.identity_residual < 1e-10 - - -def test_identity_residual_detects_a_broken_bag_logit() -> None: - """A wrong bag logit must surface, not pass silently.""" - table, labels = _instance_table(n_bags=2) - bad = {bag_id: np.array([99.0, -99.0]) for bag_id in table["bag_id"].unique()} - - result = niche_attribution_summary( - table, - labels, - attention_columns=["attention_c0", "attention_c1"], - logit_columns=["logit_0", "logit_1"], - bag_logits=bad, - n_bootstrap=10, - ) - assert result.identity_residual > 1.0 - - -def test_bias_removal_shifts_the_margin_by_exactly_the_bias_term() -> None: - attention = np.array([[0.25, 0.75], [0.75, 0.25]]) - logits = np.array([[1.0, 2.0], [3.0, 4.0]]) - bias = [0.5, -0.25] - - raw = instance_margin_contributions(attention, logits) - debiased = instance_margin_contributions(attention, logits, logit_bias=bias) - - expected_shift = attention * np.asarray(bias).reshape(1, -1) - np.testing.assert_allclose(raw - debiased, expected_shift) - - -def test_margin_contributions_reject_mismatched_shapes() -> None: - with pytest.raises(ValueError, match="same number of rows"): - instance_margin_contributions(np.zeros((3, 2)), np.zeros((4, 2))) - with pytest.raises(ValueError, match="neither 1 nor"): - instance_margin_contributions(np.zeros((3, 3)), np.zeros((3, 2))) - with pytest.raises(ValueError, match="logit_bias"): - instance_margin_contributions(np.zeros((3, 2)), np.zeros((3, 2)), logit_bias=[1.0]) - - -def test_cluster_summaries_are_bounded_and_carry_intervals() -> None: - table, labels = _instance_table(n_bags=6, per_bag=8) - result = niche_attribution_summary( - table, - labels, - attention_columns=["attention_c0", "attention_c1"], - logit_columns=["logit_0", "logit_1"], - n_bootstrap=50, - ) - summary = result.per_niche - - # Indexed by (niche_label, class_index). A binary head summarises the - # positive class only: the class-0 margin is its exact negation. - assert set(summary.index) == {(0, 1), (1, 1), (2, 1)} - # Note: margin_signed_share is NOT bounded to [-1, 1]. It is a fraction of - # the bag margin, so a niche pushing hard in one direction can exceed 1 - # when another niche opposes it. The real invariant is that a full - # partition sums to 1, covered separately below. - assert np.isfinite(summary["margin_signed_share"]).all() - assert summary["margin_sign_consistency"].between(0.0, 1.0).all() - assert summary["prevalence"].between(0.0, 1.0).all() - - for column in ("prevalence",): - assert (summary[f"{column}_lo"] <= summary[column] + 1e-9).all() - assert (summary[column] <= summary[f"{column}_hi"] + 1e-9).all() - - -def test_signed_share_across_clusters_sums_to_one_per_region() -> None: - """Shares are fractions of the bag margin, so a full partition sums to 1.""" - table, labels = _instance_table(n_bags=1, per_bag=9) - result = niche_attribution_summary( - table, - labels, - attention_columns=["attention_c0", "attention_c1"], - logit_columns=["logit_0", "logit_1"], - n_bootstrap=0, - ) - assert result.per_niche["margin_signed_share"].sum() == pytest.approx(1.0, abs=1e-9) - - -def test_bootstrap_ci_brackets_the_point_estimate() -> None: - rng = np.random.default_rng(0) - values = rng.normal(loc=2.0, size=200) - point, lo, hi = percentile_bootstrap_ci(values, n_boot=200, rng=rng) - assert lo <= point <= hi - assert point == pytest.approx(float(np.mean(values))) - - -def test_bootstrap_ci_handles_degenerate_inputs() -> None: - empty = percentile_bootstrap_ci(np.array([]), n_boot=10) - assert all(np.isnan(v) for v in empty) - - single = percentile_bootstrap_ci(np.array([3.0]), n_boot=10) - assert single == (3.0, 3.0, 3.0) - - # NaNs are dropped rather than poisoning the estimate. - point, _, _ = percentile_bootstrap_ci(np.array([1.0, np.nan, 3.0]), n_boot=10) - assert point == pytest.approx(2.0) - - -def _multiclass_table(n_bags: int = 4, per_bag: int = 6, n_classes: int = 3, seed: int = 3): - rng = np.random.default_rng(seed) - rows = [] - for bag in range(n_bags): - attn_logits = rng.normal(size=(per_bag, n_classes)) - attention = np.exp(attn_logits) / np.exp(attn_logits).sum(axis=0, keepdims=True) - logits = rng.normal(size=(per_bag, n_classes)) - for i in range(per_bag): - row = {"bag_id": f"bag_{bag}"} - for c in range(n_classes): - row[f"attention_c{c}"] = attention[i, c] - row[f"logit_{c}"] = logits[i, c] - rows.append(row) - table = pd.DataFrame(rows) - labels = np.array([i % 3 for i in range(len(table))]) - return table, labels - - -def test_ovr_margin_reduces_to_the_binary_margin_when_c_is_two() -> None: - """The binary case must be unchanged by the OvR generalisation.""" - contributions = np.array([[0.2, 0.9], [-0.4, 0.1], [1.0, -1.0]]) - margins = instance_ovr_margins(contributions) - - np.testing.assert_allclose(margins[:, 1], contributions[:, 1] - contributions[:, 0]) - # Class 0 is the exact negation, which is why binary summarises one class. - np.testing.assert_allclose(margins[:, 0], -margins[:, 1]) - - -def test_ovr_margin_contrasts_each_class_against_the_mean_of_the_rest() -> None: - contributions = np.array([[3.0, 0.0, 0.0], [1.0, 1.0, 1.0]]) - margins = instance_ovr_margins(contributions) - - # Instance 0: class 0 stands out against a zero mean of the others. - assert margins[0, 0] == pytest.approx(3.0) - assert margins[0, 1] == pytest.approx(-1.5) - # Instance 1: perfectly balanced, so no class is favoured. - np.testing.assert_allclose(margins[1], np.zeros(3), atol=1e-12) - - -def test_ovr_margin_is_a_single_column_for_scalar_heads() -> None: - """Regression and Cox emit one logit; the contribution is the margin.""" - contributions = np.array([[0.7], [-0.2]]) - np.testing.assert_allclose(instance_ovr_margins(contributions), contributions) - - -def test_multiclass_attribution_summarises_every_class() -> None: - table, labels = _multiclass_table(n_classes=3) - attention_columns = [f"attention_c{c}" for c in range(3)] - logit_columns = [f"logit_{c}" for c in range(3)] - - result = niche_attribution_summary( - table, - labels, - attention_columns=attention_columns, - logit_columns=logit_columns, - n_bootstrap=20, - ) - # One row per (niche, class): no class is silently ignored. - assert set(result.per_niche.index) == {(c, k) for c in (0, 1, 2) for k in (0, 1, 2)} - for c in range(3): - assert f"margin_c{c}" in result.per_instance.columns - assert f"contribution_c{c}" in result.per_instance.columns - - -def test_multiclass_identity_holds_for_every_class() -> None: - table, labels = _multiclass_table(n_classes=4) - attention_columns = [f"attention_c{c}" for c in range(4)] - logit_columns = [f"logit_{c}" for c in range(4)] - - bag_logits = {} - for bag_id, group in table.groupby("bag_id"): - attention = group[attention_columns].to_numpy() - logits = group[logit_columns].to_numpy() - bag_logits[bag_id] = (attention * logits).sum(axis=0) - - result = niche_attribution_summary( - table, - labels, - attention_columns=attention_columns, - logit_columns=logit_columns, - bag_logits=bag_logits, - n_bootstrap=10, - ) - assert result.identity_residual < 1e-10 - - -def test_focus_classes_can_be_selected_explicitly() -> None: - table, labels = _multiclass_table(n_classes=3) - result = niche_attribution_summary( - table, - labels, - attention_columns=[f"attention_c{c}" for c in range(3)], - logit_columns=[f"logit_{c}" for c in range(3)], - focus_classes=[2], - n_bootstrap=10, - ) - assert {k for _, k in result.per_niche.index} == {2} - - with pytest.raises(ValueError, match="out of range"): - niche_attribution_summary( - table, - labels, - attention_columns=[f"attention_c{c}" for c in range(3)], - logit_columns=[f"logit_{c}" for c in range(3)], - focus_classes=[7], - ) - - -def test_head_bias_is_read_from_a_checkpoint(tmp_path) -> None: - """The bias must come from the trained model, not be transcribed by hand.""" - import torch - - from grass_mil.models.components.heads import GraphPredictionHead - from grass_mil.models.training.checkpoint_init import read_graph_head_bias - - head = GraphPredictionHead(input_dim=4, output_dim=3, hidden_dim=4, num_layers=2) - with torch.no_grad(): - head.net[-1].bias.copy_(torch.tensor([0.25, -0.5, 1.5])) - - ckpt = tmp_path / "model.ckpt" - torch.save({"state_dict": {f"graph_head.{k}": v for k, v in head.state_dict().items()}}, ckpt) - - assert read_graph_head_bias(str(ckpt)) == pytest.approx([0.25, -0.5, 1.5]) - # An encoder-only checkpoint has no head to read. - empty = tmp_path / "encoder.ckpt" - torch.save({"state_dict": {"encoder.layers.0.weight": torch.zeros(2, 2)}}, empty) - assert read_graph_head_bias(str(empty)) is None - - -def test_margin_plugin_reads_the_bias_from_a_checkpoint(tmp_path) -> None: - import torch - - from grass_mil.contracts import InterpretabilityDataset - from grass_mil.interpretability.plugins.base import PluginContext - from grass_mil.interpretability.plugins.builtin import create_builtin_registry - from grass_mil.models.components.heads import GraphPredictionHead - - table, labels = _instance_table(n_bags=3, per_bag=4) - dataset = InterpretabilityDataset( - instance_table=table, spatial_table=None, id_column="bag_id", bag_id_column="bag_id" - ) - - head = GraphPredictionHead(input_dim=4, output_dim=2, hidden_dim=4, num_layers=2) - with torch.no_grad(): - head.net[-1].bias.copy_(torch.tensor([5.0, -5.0])) - ckpt = tmp_path / "model.ckpt" - torch.save({"state_dict": {f"graph_head.{k}": v for k, v in head.state_dict().items()}}, ckpt) - - plugin = create_builtin_registry().get("margin_attribution") - context = PluginContext(state={"niche_labels": labels}) - unbiased = plugin.run(dataset, context, n_bootstrap=0) - debiased = plugin.run(dataset, context, n_bootstrap=0, logit_bias_checkpoint=str(ckpt)) - - # A large head bias must visibly change the margins once removed. - assert not np.allclose( - unbiased.payload["per_niche"]["margin_weighted"].to_numpy(), - debiased.payload["per_niche"]["margin_weighted"].to_numpy(), - ) - - # An encoder-only checkpoint must fail loudly rather than silently - # attributing with an un-removed bias. - encoder_only = tmp_path / "encoder.ckpt" - torch.save({"state_dict": {"encoder.layers.0.weight": torch.zeros(2, 2)}}, encoder_only) - with pytest.raises(ValueError, match="No graph head bias found"): - plugin.run(dataset, context, logit_bias_checkpoint=str(encoder_only)) diff --git a/tests/interpretability/test_cell_level.py b/tests/interpretability/test_cell_level.py deleted file mode 100644 index b2cbe38..0000000 --- a/tests/interpretability/test_cell_level.py +++ /dev/null @@ -1,104 +0,0 @@ -"""Cell-level tier-2 analyses. - -The instance-level analyses treat a whole ego-graph as one node, so they measure -how neighbourhood labels relate. These measure the cells inside those graphs, -which is the unit the legacy NSCLC reports use. -""" - -import pandas as pd -import pytest - -from grass_mil.interpretability.tier2.cell_level import ( - add_cell_type_indicators, - attach_niche_labels_to_cells, - differential_cell_type_enrichment, - per_niche_cell_type_enrichment, - per_niche_cell_type_moran, -) - - -def _two_niche_cells(): - """Two niches, each an 8-cell chain. - - Niche 1 is segregated (AAAA|BBBB), niche 2 alternates (ABABABAB), so the - two should give opposite-signed same-type enrichment and Moran's I. - """ - rows, edges = [], [] - for niche, types in ((1, list("AAAABBBB")), (2, list("ABABABAB"))): - ids = [f"n{niche}_{i}" for i in range(8)] - for i, (cid, ct) in enumerate(zip(ids, types)): - rows.append( - { - "cell_id": cid, - "instance_id": f"inst{niche}", - "cell_type": ct, - "x": float(i), - "y": 0.0, - } - ) - for i in range(7): - edges.append({"source_id": ids[i], "target_id": ids[i + 1], "distance": 1.0}) - edges.append({"source_id": ids[i + 1], "target_id": ids[i], "distance": 1.0}) - return pd.DataFrame(rows), pd.DataFrame(edges) - - -def test_cells_inherit_the_niche_of_their_ego_graph(): - cells, _ = _two_niche_cells() - out = attach_niche_labels_to_cells(cells, [1, 2], ["inst1", "inst2"]) - assert out.loc[out["instance_id"] == "inst1", "niche_label"].eq(1).all() - assert out.loc[out["instance_id"] == "inst2", "niche_label"].eq(2).all() - # Every cell is labelled: this join is what makes per-niche analysis possible. - assert out["niche_label"].notna().all() - - -def test_attach_rejects_misaligned_labels(): - cells, _ = _two_niche_cells() - with pytest.raises(ValueError): - attach_niche_labels_to_cells(cells, [1], ["inst1", "inst2"]) - - -def test_one_hot_indicators_cover_every_type(): - cells, _ = _two_niche_cells() - out, cols = add_cell_type_indicators(cells) - assert cols == ["ct_A", "ct_B"] - assert out[cols].sum(axis=1).eq(1.0).all() - - -def test_enrichment_separates_segregated_from_alternating_niches(): - """The statistic must reflect the cells, not a per-niche summary label. - - Both niches have the same composition (4 A, 4 B); they differ only in how - the cells are arranged. An analysis that saw one label per ego-graph could - not tell them apart. - """ - cells, edges = _two_niche_cells() - cells = attach_niche_labels_to_cells(cells, [1, 2], ["inst1", "inst2"]) - result = per_niche_cell_type_enrichment(cells, edges, n_perms=0, min_cells=4) - - assert set(result.enrichment) == {1, 2} - segregated = result.enrichment[1] - alternating = result.enrichment[2] - # Segregated: same-type contacts dominate. Alternating: cross-type do. - assert segregated.loc["A", "A"] > segregated.loc["A", "B"] - assert alternating.loc["A", "B"] > alternating.loc["A", "A"] - - -def test_differential_subtracts_the_reference_niche(): - cells, edges = _two_niche_cells() - cells = attach_niche_labels_to_cells(cells, [1, 2], ["inst1", "inst2"]) - result = per_niche_cell_type_enrichment(cells, edges, n_perms=0, min_cells=4) - diff = differential_cell_type_enrichment(result, reference=1) - assert set(diff) == {2} - expected = result.enrichment[2] - result.enrichment[1] - pd.testing.assert_frame_equal(diff[2], expected) - - -def test_moran_is_higher_where_a_type_is_spatially_clustered(): - cells, edges = _two_niche_cells() - cells = attach_niche_labels_to_cells(cells, [1, 2], ["inst1", "inst2"]) - result = per_niche_cell_type_moran(cells, edges, n_perms=0, min_nodes=3) - stat = result.statistic - assert stat.shape[1] == 2 # one column per cell type - # Segregated niche clusters each type; the alternating one anti-correlates. - assert stat.loc[1, "ct_A"] > stat.loc[2, "ct_A"] - assert stat.loc[2, "ct_A"] < 0 diff --git a/tests/interpretability/test_core.py b/tests/interpretability/test_core.py deleted file mode 100644 index 0eec984..0000000 --- a/tests/interpretability/test_core.py +++ /dev/null @@ -1,120 +0,0 @@ -from __future__ import annotations - -import numpy as np -import pandas as pd -import pytest - -from grass_mil.interpretability.core.biomarkers import ( - niche_composition_summary, -) -from grass_mil.interpretability.core.clustering import run_clustering -from grass_mil.interpretability.core.reduction import run_reduction - - -def test_reduction_and_clustering_deterministic() -> None: - rng = np.random.default_rng(7) - x = rng.normal(size=(50, 6)) - red1 = run_reduction("pca", x, n_components=2, random_state=42) - red2 = run_reduction("pca", x, n_components=2, random_state=42) - assert np.allclose(red1.embedding, red2.embedding) - - cl = run_clustering("agglomerative", red1.embedding, n_clusters=3) - assert cl.labels.shape[0] == 50 - assert sorted(set(cl.labels.tolist())) == [0, 1, 2] - - -def test_cluster_biomarker_and_attention_summary() -> None: - table = pd.DataFrame( - { - "instance_id": [f"i{i}" for i in range(6)], - "bag_id": ["b0", "b0", "b0", "b1", "b1", "b1"], - "cell_type": ["A", "A", "B", "A", "B", "B"], - "comp_A": [1.0, 1.0, 0.0, 1.0, 0.0, 0.0], - "comp_B": [0.0, 0.0, 1.0, 0.0, 1.0, 1.0], - "attention": [0.7, 0.2, 0.1, 0.1, 0.3, 0.6], - "score": [0.9, 0.8, 0.4, 0.2, 0.3, 0.7], - } - ) - labels = np.array([0, 0, 1, 1, 1, 0]) - bio = niche_composition_summary(table, labels) - assert "comp_A" in bio.composition.columns - assert bio.niche_counts.sum() == len(table) - - -def test_cluster_biomarker_requires_composition_columns() -> None: - table = pd.DataFrame( - { - "instance_id": ["i0", "i1"], - "cell_type": ["A", "B"], - } - ) - labels = np.array([0, 1]) - with pytest.raises(ValueError, match="No composition columns found"): - niche_composition_summary(table, labels) - - -def test_cluster_biomarker_variance_estimator_is_tunable() -> None: - table = pd.DataFrame( - { - "instance_id": ["i0", "i1", "i2"], - "comp_A": [0.0, 1.0, 2.0], - "comp_B": [2.0, 1.0, 0.0], - } - ) - labels = np.array([0, 1, 1]) - - unbiased = niche_composition_summary(table, labels, variance_estimator="unbiased") - biased = niche_composition_summary(table, labels, variance_estimator="biased") - - assert unbiased.enrichment.loc[0, "comp_A"] == pytest.approx(-1.0) - assert biased.enrichment.loc[0, "comp_A"] == pytest.approx(-1.224744871391589) - assert abs(unbiased.enrichment.loc[0, "comp_A"]) < abs(biased.enrichment.loc[0, "comp_A"]) - - -def test_cluster_biomarker_variance_estimator_rejects_unknown_values() -> None: - table = pd.DataFrame( - { - "instance_id": ["i0", "i1"], - "comp_A": [1.0, 0.0], - "comp_B": [0.0, 1.0], - } - ) - labels = np.array([0, 1]) - with pytest.raises(ValueError, match="Unsupported variance_estimator"): - niche_composition_summary(table, labels, variance_estimator="invalid_mode") - - -def test_hdbscan_clustering_forwards_selection_and_density_knobs(monkeypatch) -> None: - class _FakeHDBSCAN: - def __init__(self, **kwargs): - self.kwargs = dict(kwargs) - - def fit_predict(self, x: np.ndarray) -> np.ndarray: - return np.zeros((x.shape[0],), dtype=int) - - import sklearn.cluster as sklearn_cluster - - monkeypatch.setattr(sklearn_cluster, "HDBSCAN", _FakeHDBSCAN, raising=False) - x = np.array([[0.0, 0.1], [0.2, 0.3], [0.4, 0.5]], dtype=float) - out = run_clustering( - "hdbscan", - x, - min_cluster_size=42, - min_samples=11, - cluster_selection_method="leaf", - cluster_selection_epsilon=0.15, - alpha=1.2, - metric="manhattan", - allow_single_cluster=True, - n_jobs=3, - ) - assert out.labels.shape[0] == x.shape[0] - assert out.method == "hdbscan" - assert out.fitted_object.kwargs["min_cluster_size"] == 42 - assert out.fitted_object.kwargs["min_samples"] == 11 - assert out.fitted_object.kwargs["cluster_selection_method"] == "leaf" - assert out.fitted_object.kwargs["cluster_selection_epsilon"] == pytest.approx(0.15) - assert out.fitted_object.kwargs["alpha"] == pytest.approx(1.2) - assert out.fitted_object.kwargs["metric"] == "manhattan" - assert out.fitted_object.kwargs["allow_single_cluster"] is True - assert out.fitted_object.kwargs["n_jobs"] == 3 diff --git a/tests/interpretability/test_fidelity_phase6.py b/tests/interpretability/test_fidelity_phase6.py deleted file mode 100644 index c8fcea4..0000000 --- a/tests/interpretability/test_fidelity_phase6.py +++ /dev/null @@ -1,292 +0,0 @@ -"""Phase 6: niche discovery (cosine HDBSCAN, relattn margins, tables, co-localisation, pipeline).""" - -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import pandas as pd -import pytest - -from tests.helpers.synthetic_cohort import cohort_overrides, compose_cfg, make_synthetic_cohort - -pytest.importorskip("sklearn") - - -def _blobs(n_per=200, d=16, seed=0): - rng = np.random.default_rng(seed) - centres = rng.normal(size=(3, d)) * 6 - X = np.concatenate([c + rng.normal(size=(n_per, d)) for c in centres]).astype(np.float32) - y = np.repeat(np.arange(3), n_per) - return X, y - - -def test_cosine_hdbscan_recovers_blobs_and_assigns_unsampled_rows(): - from sklearn.metrics import adjusted_rand_score - - from grass_mil.interpretability.core.clustering import cosine_hdbscan, fit_niche_clustering - - X, y = _blobs() - model, labels = cosine_hdbscan(X, n_components=5, min_cluster_size=20, sample_size=150, seed=0) - assert model.sample_idx.size == 150 and model.n_clusters == 3 - assigned = labels >= 0 - assert assigned.mean() > 0.95 - assert adjusted_rand_score(y[assigned], labels[assigned]) > 0.95 - unsampled = np.setdiff1d(np.arange(len(X)), model.sample_idx) - assert adjusted_rand_score(y[unsampled], labels[unsampled]) > 0.95 - assert np.array_equal(model.predict(X[model.sample_idx]), model.labels_sample) - # any registry backend: k-means with the same sample/assign machinery - km, lab_km = fit_niche_clustering( - X, - reduction={"method": "pca", "n_components": 5}, - clustering={"method": "kmeans", "n_clusters": 3}, - sample_size=None, - ) - assert adjusted_rand_score(y, lab_km) > 0.95 and km.clustering_method == "kmeans" - assert km.params["clustering"]["n_clusters"] == 3 - - -def test_relattn_formula_and_region_sum_identity(): - from grass_mil.interpretability.niches.margins import relattn_margins - - rng = np.random.default_rng(0) - rid = np.array(["a"] * 30 + ["b"] * 50) - LG = rng.normal(size=(80, 2)) - AR = rng.normal(size=(80, 2)) - bias = np.array([0.3, -0.2]) - res = relattn_margins(LG, AR, bias, rid) - for r, n in (("a", 30), ("b", 50)): - m = rid == r - a = np.exp(AR[m] - AR[m].max(0)) - a /= a.sum(0, keepdims=True) - direct = (a[:, 1] * (LG[m, 1] - bias[1]) - a[:, 0] * (LG[m, 0] - bias[0])) * n - assert np.allclose(res.margin[m], direct, atol=1e-5) - assert np.allclose(res.attention[m].sum(0), 1.0, atol=1e-5) - L = (a * LG[m]).sum(0) # bag logits = sum A l - assert res.margin[m].sum() / n == pytest.approx( - (L[1] - bias[1]) - (L[0] - bias[0]), abs=1e-5 - ) - - -def test_niche_table_and_colocalisation(): - from grass_mil.interpretability.niches.colocalisation import cluster_colocalisation - from grass_mil.interpretability.niches.tables import niche_table, spearman_concordance - - rng = np.random.default_rng(0) - n = 600 - rid = np.repeat([f"r{i}" for i in range(6)], 100) - region_labels = {f"r{i}": i % 2 for i in range(6)} - labels = rng.integers(0, 3, n) - labels[:20] = -1 - ct = rng.integers(0, 4, n) - margin = rng.normal(size=n) + (labels == 2) * 1.0 - att = rng.uniform(size=n) - T = niche_table( - labels, - cell_types=ct, - region_ids=rid, - region_labels=region_labels, - margin=margin, - attention_pos=att, - cell_type_names=list("ABCD"), - n_boot=50, - ) - assert set(T.cluster) == {"C0", "C1", "C2", "Noise"} - for col in ( - "share", - "cores_present", - "emp_frac_pos", - "share_auroc", - "relattn", - "relattn_lo", - "relattn_hi", - "signed_share", - "sign_consistency", - "lift", - "comp_A", - ): - assert col in T.columns - assert T.share.sum() == pytest.approx(1.0) - assert ( - T.set_index("cluster").loc["C2", "relattn"] > T.set_index("cluster").loc["C0", "relattn"] - ) - conc = spearman_concordance(T) - assert conc["n_clusters"] == 3 and "spearman_emp_relattn" in conc - - edges = np.stack([rng.integers(0, n, 3000), rng.integers(0, n, 3000)]) - CO, clusters = cluster_colocalisation(labels, edges) - assert clusters == [0, 1, 2] and np.allclose(CO, CO.T) - assert np.abs(CO).max() < 0.5 # random edges: observed ~ expected - # segregated graph: edges only within cluster -> positive diagonal, negative off-diagonal - same = np.array( - [ - (i, j) - for i in range(n) - for j in rng.integers(0, n, 5) - if labels[i] == labels[j] and labels[i] >= 0 - ][:3000] - ).T - CO2, _ = cluster_colocalisation(labels, same) - assert np.all(np.diag(CO2) > 0) and (CO2[~np.eye(3, dtype=bool)] < 0).all() - - -def test_identity_residual_from_bag_table(): - from grass_mil.interpretability.core.attribution import niche_attribution_summary - from grass_mil.interpretability.plugins.builtin import _bag_logits_from_table - - rng = np.random.default_rng(0) - rows = [] - bag_rows = [] - for b in range(3): - n = 10 - lg = rng.normal(size=(n, 2)) - a_logit = rng.normal(size=(n, 2)) - a = np.exp(a_logit - a_logit.max(0)) - a /= a.sum(0) - L = (a * lg).sum(0) - bag_rows.append({"bag_id": f"b{b}", "logit_0": L[0], "logit_1": L[1]}) - for i in range(n): - rows.append( - { - "bag_id": f"b{b}", - "attention_c0": a[i, 0], - "attention_c1": a[i, 1], - "logit_0": lg[i, 0], - "logit_1": lg[i, 1], - } - ) - table = pd.DataFrame(rows) - bag_logits = _bag_logits_from_table( - pd.DataFrame(bag_rows), bag_id_column="bag_id", logit_prefix="logit_" - ) - assert set(bag_logits) == {"b0", "b1", "b2"} - res = niche_attribution_summary( - table, - np.zeros(len(table), dtype=int), - attention_columns=["attention_c0", "attention_c1"], - logit_columns=["logit_0", "logit_1"], - bag_logits=bag_logits, - n_bootstrap=10, - ) - assert res.identity_residual < 1e-6 - assert np.isnan( - niche_attribution_summary( - table, - np.zeros(len(table), dtype=int), - attention_columns=["attention_c0", "attention_c1"], - logit_columns=["logit_0", "logit_1"], - n_bootstrap=10, - ).identity_residual - ) - - -@pytest.fixture(scope="module") -def cache_dir(tmp_path_factory): - pytest.importorskip("torch_geometric") - import hydra - - from grass_mil.extract import run_extract - from grass_mil.train import train - - root = tmp_path_factory.mktemp("cohort6") - c = make_synthetic_cohort(root, n_patients=4, regions_per_patient=2, n_cells=80) - cfg = compose_cfg("train", cohort_overrides(c, root / "pre")) - hydra.utils.instantiate(cfg.data).prepare_data() - train_out = root / "train" - train(compose_cfg("train", cohort_overrides(c, train_out) + ["trainer.max_epochs=1"])) - ckpt = train_out / "checkpoints" / "last.ckpt" - ecfg = compose_cfg( - "extract", - cohort_overrides(c, root / "extract") - + [f"+extract.models.final={ckpt}", "extract.device=cpu"], - ) - result = run_extract(ecfg) - return Path(result["models"][0]), ckpt - - -def test_niche_discovery_pipeline_default_and_kmeans(cache_dir, tmp_path): - from grass_mil.inference.cell_cache import CellCache - from grass_mil.interpretability.niches.pipeline import run_niche_discovery, write_niche_outputs - - model_dir, ckpt = cache_dir - cache = CellCache.open(model_dir) - res = run_niche_discovery( - cache, - regions="all" and cache.regions, - clustering={ - "method": "hdbscan", - "params": { - "min_cluster_size": 15, - "min_samples": 1, - "cluster_selection_method": "eom", - }, - }, - sample_size=None, - n_boot=50, - ) - assert res.labels.size == sum(cache.n_cells(r) for r in cache.regions) - assert res.n_clusters >= 1 and res.table.share.sum() == pytest.approx(1.0) - assert res.colocalisation.shape == (res.n_clusters, res.n_clusters) - paths = write_niche_outputs(res, tmp_path / "default", title="test", umap={"enabled": False}) - for key in ( - "labels", - "cluster_model", - "cells", - "clusters", - "report", - "pca_triptych", - "clusters_overview", - "colocalisation", - ): - assert key in paths and Path(paths[key]).exists(), key - cells = pd.read_csv(paths["cells"]) - assert len(cells) == res.labels.size and set(cells.columns) >= { - "region_id", - "cluster", - "cell_type", - "relattn", - "x", - "y", - } - - km = run_niche_discovery( - cache, - regions=cache.regions, - clustering={"method": "kmeans", "n_clusters": 4}, - reduction={"method": "pca", "n_components": 5}, - sample_size=None, - n_boot=20, - ) - assert km.n_clusters == 4 and km.noise_fraction == 0.0 - - att = run_niche_discovery( - cache, - regions=cache.regions, - space="attention", - attention_ckpt=str(ckpt), - clustering={"method": "kmeans", "n_clusters": 3}, - reduction={"method": "none"}, - sample_size=None, - n_boot=20, - ) - assert att.features.shape[1] == 4 # attention hidden_dim of the synthetic model - - -def test_niches_cli_runs(cache_dir, tmp_path): - from grass_mil.interpretability.niches_cli import run_niches - - model_dir, _ = cache_dir - cfg = compose_cfg( - "niches", - [ - f"cells_dir={model_dir}", - f"output_dir={tmp_path / 'niches'}", - "regions=all", - "clustering.params.min_cluster_size=15", - "sample_size=null", - "n_boot=20", - "umap.enabled=false", - ], - ) - out = run_niches(cfg) - assert out["n_clusters"] >= 1 and Path(out["paths"]["report"]).exists() diff --git a/tests/interpretability/test_fidelity_phase7.py b/tests/interpretability/test_fidelity_phase7.py deleted file mode 100644 index 10d49a9..0000000 --- a/tests/interpretability/test_fidelity_phase7.py +++ /dev/null @@ -1,176 +0,0 @@ -"""Phase 7: cross-fold metaclusters, transfer and shared basis.""" - -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import pandas as pd -import pytest - -from tests.helpers.synthetic_cohort import cohort_overrides, compose_cfg, make_synthetic_cohort - -pytest.importorskip("sklearn") - - -def test_hellinger_range_and_constrained_linkage_never_merges_same_fold(): - from grass_mil.interpretability.metaclusters.match import ( - constrained_average_linkage, - hellinger_matrix, - scan_threshold, - ) - - comp = np.array([[1, 0, 0], [0, 1, 0], [0.5, 0.5, 0], [1, 0, 0], [0, 0, 1]], dtype=float) - D = hellinger_matrix(comp) - assert D.shape == (5, 5) and np.allclose(D, D.T) and D.min() >= 0 and D.max() <= 1 + 1e-9 - assert D[0, 3] == 0 and D[0, 1] == pytest.approx(1.0) - fold = np.array(["f0", "f1", "f0", "f1", "f2"]) - groups = constrained_average_linkage(D, fold, 0.9) - for g in groups: - assert len(set(fold[g])) == len(g) # cannot-link respected - assert sorted(sorted(g) for g in groups)[0] == [0, 3] # identical compositions merge - thr, best, scan = scan_threshold(D, fold, grid=[0.05, 0.5, 0.9], min_folds=2) - assert scan.n_replicated.max() == sum(len(g) >= 2 for g in best) - # ties -> fewest metaclusters - tied = scan[scan.n_replicated == scan.n_replicated.max()] - assert len(best) == tied.n_metaclusters.min() - - -def test_patient_shares_and_tests(): - from grass_mil.interpretability.metaclusters.match import patient_shares, patient_tests - - region_labels = { - "r0": np.array([0, 0, 1, -1]), - "r1": np.array([1, 1, 1, 1]), - "r2": np.array([0, 0, 0, 0]), - } - region_fold = {"r0": "f0", "r1": "f0", "r2": "f1"} - id_to_meta = {"f0_C0": 1, "f0_C1": 2, "f1_C0": 1} - reg, pat = patient_shares( - region_labels, region_fold, id_to_meta, {"r0": "P", "r1": "P", "r2": "Q"}, [1, 2] - ) - assert ( - reg.loc["r0", "M1"] == 0.5 and reg.loc["r0", "M2"] == 0.25 and reg.loc["r2", "M1"] == 1.0 - ) - assert (reg.sum(axis=1) <= 1 + 1e-9).all() and pat.loc["P", "M2"] == pytest.approx( - (0.25 + 1.0) / 2 - ) - shares = pd.DataFrame( - {"M1": np.r_[np.ones(5), np.zeros(5)], "M2": np.random.default_rng(0).uniform(size=10)}, - index=[f"p{i}" for i in range(10)], - ) - y = pd.Series(np.r_[np.ones(5), np.zeros(5)].astype(int), index=shares.index) - T = patient_tests(shares, y) - assert set(T.columns) >= {"meta", "share_pos", "share_neg", "auroc", "p", "q"} - assert T.set_index("meta").loc[1, "auroc"] == 1.0 - - -@pytest.fixture(scope="module") -def two_fold_caches(tmp_path_factory): - pytest.importorskip("torch_geometric") - import hydra - - from grass_mil.extract import run_extract - from grass_mil.train import train - - root = tmp_path_factory.mktemp("cohort7") - c = make_synthetic_cohort( - root, n_patients=6, regions_per_patient=2, n_cells=80, n_outer_folds=2 - ) - hydra.utils.instantiate( - compose_cfg("train", cohort_overrides(c, root / "pre")).data - ).prepare_data() - ckpts = {} - for i in range(2): - test_ids = c.folds[i] - train_ids = [r for j, f in enumerate(c.folds) if j != i for r in f] - out = root / f"train{i}" - train( - compose_cfg( - "train", - cohort_overrides(c, out) - + [ - "trainer.max_epochs=1", - f"data.split.explicit.train=[{','.join(train_ids)}]", - f"data.split.explicit.val=[{','.join(test_ids)}]", - f"data.split.explicit.test=[{','.join(test_ids)}]", - ], - ) - ) - ckpts[f"fold{i}"] = (out / "checkpoints" / "last.ckpt", test_ids) - overrides = cohort_overrides(c, root / "extract") + ["extract.device=cpu"] - for name, (ckpt, held) in ckpts.items(): - overrides += [ - f"+extract.models.{name}={ckpt}", - f"+extract.held_out.{name}=[{','.join(held)}]", - ] - result = run_extract(compose_cfg("extract", overrides)) - return c, Path(result["output_dir"]) - - -def test_metaclusters_cli_all_stages(two_fold_caches, tmp_path): - from grass_mil.interpretability.metaclusters_cli import run_metaclusters - - c, cells_dir = two_fold_caches - cfg = compose_cfg( - "metaclusters", - [ - f"cells_dir={cells_dir}", - f"output_dir={tmp_path / 'meta'}", - f"labels.label_file={c.label_path}", - "labels.id_column=REGION_ID", - "labels.label_column=STAS", - "labels.patient_column=PATIENT", - "clustering.method=kmeans", - "clustering.n_clusters=3", - "sample_size=null", - "n_boot=20", - "match.min_folds=2", - "transfer.min_cells=5", - "basis.k=4", - "basis.sample_size=1000", - "basis.n_init=2", - ], - ) - summary = run_metaclusters(cfg) - assert set(summary["cluster"]) == {"fold0", "fold1"} - out = tmp_path / "meta" - fc = pd.read_csv(out / "cluster" / "fold_clusters.csv") - assert set(fc.fold) == {"fold0", "fold1"} and len(fc) == 6 - mt = pd.read_csv(out / "match" / "metaclusters.csv") - assert mt.n_folds.max() <= 2 and summary["match"]["n_metaclusters"] == len(mt) - matched = pd.read_csv(out / "match" / "fold_clusters_matched.csv") - for _, g in matched.groupby("meta"): - assert g.fold.nunique() == len(g) # cannot-link - shares = pd.read_csv(out / "match" / "metacluster_patient_shares.csv", index_col=0) - assert len(shares) == 6 and (shares.sum(axis=1) <= 1 + 1e-6).all() - assert (out / "match" / "metacluster_patient_tests.csv").exists() - ML = np.load(out / "transfer" / "transfer_metalabels.npy") - assert ML.shape[0] == 2 and ML.shape[1] == sum(80 for _ in c.region_ids) - ari = pd.read_csv(out / "transfer" / "transfer_model_ari.csv", index_col=0) - assert ari.loc["fold0", "fold0"] == pytest.approx(1.0) - assert all(v > 0.99 for v in summary["transfer"]["refit_agreement"].values()) - basis = pd.read_csv(out / "basis" / "basis_attribution.csv") - assert len(basis) == 4 and {"relattn_heldout", "relattn_train", "relattn_all"} <= set( - basis.columns - ) - assert (out / "basis" / "report_basis.md").exists() and ( - out / "metaclusters_summary.json" - ).exists() - - # fixed threshold + stage re-run reuses stage-1 outputs - cfg2 = compose_cfg( - "metaclusters", - [ - f"cells_dir={cells_dir}", - f"output_dir={tmp_path / 'meta'}", - f"labels.label_file={c.label_path}", - "labels.label_column=STAS", - "labels.patient_column=PATIENT", - "stage=match", - "match.threshold=0.4", - "match.min_folds=2", - ], - ) - s2 = run_metaclusters(cfg2) - assert s2["match"]["threshold"] == 0.4 diff --git a/tests/interpretability/test_niches.py b/tests/interpretability/test_niches.py deleted file mode 100644 index 9ef213e..0000000 --- a/tests/interpretability/test_niches.py +++ /dev/null @@ -1,236 +0,0 @@ -"""Niche naming, the Background label, and composition ordering.""" - -from __future__ import annotations - -import numpy as np -import pandas as pd -import pytest - -from grass_mil.interpretability.core.biomarkers import niche_composition_summary -from grass_mil.interpretability.core.niches import ( - BACKGROUND_NICHE_ID, - BACKGROUND_NICHE_NAME, - is_background, - niche_display_name, - niche_display_names, - order_niches_by_composition, - relabel_niches_by_composition, -) - - -def test_background_replaces_the_noise_label() -> None: - """Density-based clustering's -1 is a named category, not "noise".""" - assert niche_display_name(BACKGROUND_NICHE_ID) == BACKGROUND_NICHE_NAME - assert is_background(-1) - assert is_background(BACKGROUND_NICHE_NAME) - assert not is_background(0) - - -def test_niche_names_keep_the_numeric_id_for_joinability() -> None: - assert niche_display_name(3) == "Niche 3" - assert niche_display_names([0, 2, -1]) == ["Niche 0", "Niche 2", "Background"] - - -def test_niches_are_ordered_so_similar_compositions_are_adjacent() -> None: - # Two tumour-rich niches, two stroma-rich ones, interleaved by id. - composition = pd.DataFrame( - { - "comp_tumour": [0.9, 0.1, 0.85, 0.15], - "comp_stroma": [0.1, 0.9, 0.15, 0.85], - }, - index=[0, 1, 2, 3], - ) - order = order_niches_by_composition(composition) - - assert sorted(order) == [0, 1, 2, 3] - positions = {niche: position for position, niche in enumerate(order)} - # The tumour-rich pair (0, 2) must be adjacent, as must the stroma pair. - assert abs(positions[0] - positions[2]) == 1 - assert abs(positions[1] - positions[3]) == 1 - - -def test_background_is_pinned_last_and_excluded_from_the_tree() -> None: - """Background is a catch-all; letting it join the tree distorts it.""" - composition = pd.DataFrame( - { - "comp_a": [0.9, 0.85, 0.1, 0.5], - "comp_b": [0.1, 0.15, 0.9, 0.5], - }, - index=[0, 1, 2, BACKGROUND_NICHE_ID], - ) - order = order_niches_by_composition(composition) - assert order[-1] == BACKGROUND_NICHE_ID - assert set(order[:-1]) == {0, 1, 2} - - -def test_ordering_degrades_gracefully_on_tiny_or_empty_inputs() -> None: - assert order_niches_by_composition(pd.DataFrame()) == [] - - two = pd.DataFrame({"comp_a": [0.1, 0.9]}, index=[0, 1]) - # Fewer than three leaves: linkage adds nothing, order is preserved. - assert order_niches_by_composition(two) == [0, 1] - - constant = pd.DataFrame({"comp_a": [0.5] * 4, "comp_b": [0.5] * 4}, index=[0, 1, 2, 3]) - # Correlation distance is undefined here; must still return a full order. - assert sorted(order_niches_by_composition(constant)) == [0, 1, 2, 3] - - -def test_relabelling_makes_the_niche_id_the_dendrogram_position() -> None: - """Niche 1 is first in dendrogram order, Niche 2 second, and so on.""" - rows = [] - labels = [] - # Ids 0 and 2 are tumour-rich, 1 and 3 stroma-rich, deliberately interleaved. - profiles = {0: (0.9, 0.1), 1: (0.1, 0.9), 2: (0.85, 0.15), 3: (0.15, 0.85)} - for niche, (tumour, stroma) in profiles.items(): - for _ in range(5): - rows.append({"comp_tumour": tumour, "comp_stroma": stroma}) - labels.append(niche) - table = pd.DataFrame(rows) - - relabelled, mapping = relabel_niches_by_composition(table, np.array(labels)) - - # Ids are 1-based and contiguous. - assert sorted(set(relabelled)) == [1, 2, 3, 4] - # The two tumour-rich niches now hold adjacent ids, as do the stroma pair. - assert abs(mapping[0] - mapping[2]) == 1 - assert abs(mapping[1] - mapping[3]) == 1 - - -def test_relabelling_keeps_background_out_of_the_numbering() -> None: - rows = [] - labels = [] - for niche, (a, b) in {0: (0.9, 0.1), 1: (0.1, 0.9), 2: (0.5, 0.5), -1: (0.4, 0.6)}.items(): - for _ in range(4): - rows.append({"comp_a": a, "comp_b": b}) - labels.append(niche) - - relabelled, mapping = relabel_niches_by_composition(pd.DataFrame(rows), np.array(labels)) - - assert mapping[-1] == BACKGROUND_NICHE_ID - assert sorted(set(relabelled)) == [-1, 1, 2, 3] - - -def test_relabelling_is_a_no_op_without_composition_columns() -> None: - table = pd.DataFrame({"score": [0.1, 0.2, 0.3]}) - labels = np.array([2, 0, 1]) - relabelled, mapping = relabel_niches_by_composition(table, labels) - np.testing.assert_array_equal(relabelled, labels) - assert mapping == {0: 0, 1: 1, 2: 2} - - -def test_composition_summary_follows_niche_id_order() -> None: - """Ordering is carried by the ids, not re-derived per summary.""" - rows = [] - labels = [] - for niche, (tumour, stroma) in {1: (0.9, 0.1), 2: (0.1, 0.9), -1: (0.5, 0.5)}.items(): - for _ in range(6): - rows.append({"comp_tumour": tumour, "comp_stroma": stroma}) - labels.append(niche) - - summary = niche_composition_summary(pd.DataFrame(rows), np.array(labels)) - - order = list(summary.composition.index) - assert order == sorted(order) - assert list(summary.enrichment.index) == order - assert list(summary.niche_counts.index) == order - - -def test_report_heatmap_labels_niches_by_name() -> None: - pytest.importorskip("plotly.graph_objects") - from grass_mil.contracts import NicheSummary, ReportBundle - from grass_mil.interpretability.core.data import InterpretabilityDataset - from grass_mil.interpretability.reporting.plotly_builders import ( - build_niche_enrichment_heatmap, - ) - - enrichment = pd.DataFrame( - {"comp_a": [1.0, -1.0], "comp_b": [-1.0, 1.0]}, index=[0, BACKGROUND_NICHE_ID] - ) - summary = NicheSummary( - niche_labels=np.array([0, -1]), - composition=enrichment.copy(), - enrichment=enrichment, - niche_counts=pd.Series({0: 1, -1: 1}), - ) - bundle = ReportBundle( - dataset=InterpretabilityDataset( - instance_table=pd.DataFrame({"instance_id": ["a", "b"]}), - spatial_table=None, - id_column="instance_id", - bag_id_column="bag_id", - cell_type_column="cell_type", - ), - reduction=None, - clustering=None, - niche_summary=summary, - plugin_results={}, - niche_feature_reduction=None, - artifacts_dir=None, - ) - - fig = build_niche_enrichment_heatmap(bundle) - assert list(fig.data[0].y) == ["Niche 0", "Background"] - assert "Niche" in fig.layout.title.text - - -def test_pipeline_exports_niches_already_in_dendrogram_order(tmp_path) -> None: - """The relabelling happens once, in the pipeline, and reaches the exports.""" - import pandas as pd - - from grass_mil.contracts import InterpretabilityDataset - from grass_mil.interpretability.pipeline import run_interpretability_pipeline - - rng = np.random.default_rng(0) - rows = [] - # Four well-separated composition profiles; two tumour-rich, two stroma-rich. - for tumour, stroma in [(0.95, 0.05), (0.05, 0.95), (0.85, 0.15), (0.15, 0.85)]: - for _ in range(12): - jitter = rng.normal(scale=0.005) - rows.append( - { - "comp_tumour": tumour + jitter, - "comp_stroma": stroma - jitter, - "inst_emb_0": tumour + jitter, - "inst_emb_1": stroma - jitter, - } - ) - table = pd.DataFrame(rows) - table.insert(0, "instance_id", [f"i{i}" for i in range(len(table))]) - table["bag_id"] = "b0" - - dataset = InterpretabilityDataset( - instance_table=table, - spatial_table=None, - id_column="instance_id", - bag_id_column="bag_id", - ) - bundle = run_interpretability_pipeline( - dataset, - { - "reduction": {"enabled": True, "method": "pca", "params": {"n_components": 2}}, - "clustering": { - "enabled": True, - "method": "agglomerative", - "params": {"n_clusters": 4, "linkage": "ward"}, - }, - "plugins": {"enabled": []}, - }, - artifacts_dir=tmp_path / "artifacts", - ) - - labels = bundle.clustering.labels - # Ids are 1-based after relabelling, and the summary is already ordered. - assert sorted(set(labels)) == [1, 2, 3, 4] - assert list(bundle.niche_summary.composition.index) == [1, 2, 3, 4] - - # Compositionally similar niches carry adjacent ids. Leaf order groups the - # pairs; it does not sort by magnitude within a pair, so adjacency -- not - # monotonicity -- is the property to assert. - tumour = bundle.niche_summary.composition["comp_tumour"].to_numpy() - tumour_rich = sorted(np.argsort(tumour)[-2:]) - stroma_rich = sorted(np.argsort(tumour)[:2]) - assert tumour_rich[1] - tumour_rich[0] == 1 - assert stroma_rich[1] - stroma_rich[0] == 1 - - exported = pd.read_csv(tmp_path / "artifacts" / "niche_labels.csv") - assert sorted(exported["niche_label"].unique()) == [1, 2, 3, 4] diff --git a/tests/interpretability/test_pipeline_and_report.py b/tests/interpretability/test_pipeline_and_report.py deleted file mode 100644 index ba9b885..0000000 --- a/tests/interpretability/test_pipeline_and_report.py +++ /dev/null @@ -1,390 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import pandas as pd -import pytest -from hydra.core.hydra_config import HydraConfig -from omegaconf import DictConfig, open_dict - -from grass_mil.contracts import ClusteringResult -from grass_mil.interpretability.core.data import load_interpretability_dataset -from grass_mil.interpretability.pipeline import run_interpretability_pipeline -from grass_mil.interpretability.report_cli import run_report - -pytest.importorskip("plotly.graph_objects") - -from grass_mil.interpretability.reporting.plotly_builders import bundle_figures -from grass_mil.interpretability.reporting.render import render_interpretability_report - - -def _write_minimal_tables(tmp_path: Path) -> tuple[Path, Path]: - node = pd.DataFrame( - { - "instance_id": [f"i{i}" for i in range(8)], - "bag_id": ["b0"] * 4 + ["b1"] * 4, - "cell_type": ["A", "A", "B", "B", "A", "B", "B", "A"], - "comp_A": [1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0], - "comp_B": [0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0], - "score": [0.1, 0.2, 0.8, 0.7, 0.2, 0.9, 0.5, 0.3], - "attention": [0.2, 0.3, 0.2, 0.3, 0.1, 0.5, 0.2, 0.2], - "sample_id": ["sx"] * 4 + ["sy"] * 4, - "condition": ["X"] * 4 + ["Y"] * 4, - "center_x": [0.0, 0.1, 0.2, 0.3, 1.0, 1.1, 1.2, 1.3], - "center_y": [0.0, 0.2, 0.1, 0.3, 1.0, 1.2, 1.1, 1.3], - "inst_emb_0": [0.0, 0.1, 0.2, 0.3, 1.0, 1.1, 1.2, 1.3], - "inst_emb_1": [0.0, 0.2, 0.1, 0.3, 1.0, 1.2, 1.1, 1.3], - } - ) - edges = pd.DataFrame( - { - "source_id": ["i0", "i1", "i2", "i3", "i4", "i5", "i6", "i7"], - "target_id": ["i1", "i2", "i3", "i0", "i5", "i6", "i7", "i4"], - "distance": [0.1, 0.2, 0.4, 0.3, 0.7, 0.8, 0.9, 0.6], - "weight": [1, 1, 1, 1, 1, 1, 1, 1], - } - ) - instance_path = tmp_path / "instance_table.csv" - spatial_path = tmp_path / "spatial_table.csv" - node.to_csv(instance_path, index=False) - edges.to_csv(spatial_path, index=False) - return instance_path, spatial_path - - -def test_pipeline_and_report_render(monkeypatch, tmp_path: Path) -> None: - instance_path, spatial_path = _write_minimal_tables(tmp_path) - dataset = load_interpretability_dataset( - instance_table_path=instance_path, - spatial_table_path=spatial_path, - id_column="instance_id", - bag_id_column="bag_id", - cell_type_column="cell_type", - ) - cfg = { - "reduction": {"enabled": True, "method": "pca", "params": {"n_components": 2}}, - "clustering": { - "enabled": True, - "method": "agglomerative", - "params": {"n_clusters": 2, "linkage": "ward"}, - }, - "plugins": { - "enabled": ["niche_profiles"], - }, - } - bundle = run_interpretability_pipeline(dataset, cfg, artifacts_dir=tmp_path / "artifacts") - figs = bundle_figures(bundle) - assert len(figs) >= 2 - - monkeypatch.setattr( - "grass_mil.interpretability.reporting.render.render_pdf_via_playwright", - lambda html_path, pdf_path: pdf_path, # type: ignore[lambda-assign] - ) - monkeypatch.setattr( - "grass_mil.interpretability.reporting.render.export_plotly_snapshots", - lambda figures, output_dir, scale=2.0: {}, # type: ignore[lambda-assign] - ) - out = render_interpretability_report( - bundle, - output_dir=tmp_path / "report", - html_enabled=True, - pdf_enabled=True, - ) - assert out["html_path"] is not None - assert out["pdf_path"] is not None - - -def test_report_render_tolerates_snapshot_export_failure_when_pdf_disabled( - monkeypatch, tmp_path: Path -) -> None: - instance_path, spatial_path = _write_minimal_tables(tmp_path) - dataset = load_interpretability_dataset( - instance_table_path=instance_path, - spatial_table_path=spatial_path, - id_column="instance_id", - bag_id_column="bag_id", - cell_type_column="cell_type", - ) - cfg = { - "reduction": {"enabled": True, "method": "pca", "params": {"n_components": 2}}, - "clustering": { - "enabled": True, - "method": "agglomerative", - "params": {"n_clusters": 2, "linkage": "ward"}, - }, - "plugins": { - "enabled": ["niche_profiles"], - "params": {}, - }, - } - bundle = run_interpretability_pipeline(dataset, cfg, artifacts_dir=tmp_path / "artifacts") - - def _raise_snapshot_error(figures, output_dir, scale=2.0): # type: ignore[unused-argument] - raise RuntimeError("kaleido unavailable") - - monkeypatch.setattr( - "grass_mil.interpretability.reporting.render.export_plotly_snapshots", - _raise_snapshot_error, - ) - out = render_interpretability_report( - bundle, - output_dir=tmp_path / "report", - html_enabled=True, - pdf_enabled=False, - ) - assert out["html_path"] is not None - assert out["pdf_path"] is None - assert out["snapshot_paths"] == {} - - -def test_pipeline_and_report_render_with_tissue_graph_plugin(monkeypatch, tmp_path: Path) -> None: - instance_path, spatial_path = _write_minimal_tables(tmp_path) - dataset = load_interpretability_dataset( - instance_table_path=instance_path, - spatial_table_path=spatial_path, - id_column="instance_id", - bag_id_column="bag_id", - cell_type_column="cell_type", - ) - cfg = { - "reduction": {"enabled": True, "method": "pca", "params": {"n_components": 2}}, - "clustering": { - "enabled": True, - "method": "agglomerative", - "params": {"n_clusters": 2, "linkage": "ward"}, - }, - "plugins": { - "enabled": ["tissue_graph"], - "params": { - "tissue_graph": { - "sample_column": "sample_id", - "sample_value": "sx", - "id_column": "instance_id", - } - }, - }, - } - bundle = run_interpretability_pipeline(dataset, cfg, artifacts_dir=tmp_path / "artifacts") - assert "tissue_graph" in bundle.plugin_results - figs = dict(bundle_figures(bundle)) - assert "tissue_graph_sample_id_sx" in figs - - monkeypatch.setattr( - "grass_mil.interpretability.reporting.render.export_plotly_snapshots", - lambda figures, output_dir, scale=2.0: {}, # type: ignore[lambda-assign] - ) - out = render_interpretability_report( - bundle, - output_dir=tmp_path / "report_tissue_graph", - html_enabled=True, - pdf_enabled=False, - ) - assert out["html_path"] is not None - - -def test_pipeline_fails_fast_when_plugin_required_inputs_missing_cluster_labels( - tmp_path: Path, -) -> None: - instance_path, spatial_path = _write_minimal_tables(tmp_path) - dataset = load_interpretability_dataset( - instance_table_path=instance_path, - spatial_table_path=spatial_path, - id_column="instance_id", - bag_id_column="bag_id", - cell_type_column="cell_type", - ) - cfg = { - "reduction": {"enabled": True, "method": "pca", "params": {"n_components": 2}}, - "clustering": {"enabled": False}, - "plugins": {"enabled": ["niche_profiles"], "params": {}}, - } - with pytest.raises(ValueError, match="missing required inputs: niche_labels"): - run_interpretability_pipeline(dataset, cfg, artifacts_dir=tmp_path / "artifacts") - - -def test_pipeline_fails_fast_when_plugin_required_inputs_missing_spatial_table( - tmp_path: Path, -) -> None: - instance_path, _ = _write_minimal_tables(tmp_path) - dataset = load_interpretability_dataset( - instance_table_path=instance_path, - id_column="instance_id", - bag_id_column="bag_id", - cell_type_column="cell_type", - ) - cfg = { - "reduction": {"enabled": True, "method": "pca", "params": {"n_components": 2}}, - "clustering": { - "enabled": True, - "method": "agglomerative", - "params": {"n_clusters": 2, "linkage": "ward"}, - }, - "plugins": {"enabled": ["niche_label_moran"], "params": {}}, - } - with pytest.raises(ValueError, match="missing required inputs: spatial_table"): - run_interpretability_pipeline(dataset, cfg, artifacts_dir=tmp_path / "artifacts") - - -def test_pipeline_clusters_on_raw_embeddings_when_cluster_on_pca_disabled( - monkeypatch, tmp_path: Path -) -> None: - instance_path, _ = _write_minimal_tables(tmp_path) - dataset = load_interpretability_dataset( - instance_table_path=instance_path, - id_column="instance_id", - bag_id_column="bag_id", - cell_type_column="cell_type", - ) - expected_raw = dataset.instance_table[["inst_emb_0", "inst_emb_1"]].to_numpy(dtype=float) - captured: dict[str, np.ndarray] = {} - - def _capture_run_clustering(method, x, **params): # type: ignore[no-untyped-def] - captured["x"] = np.asarray(x) - return ClusteringResult( - method=str(method), - labels=np.zeros((int(x.shape[0]),), dtype=int), - params=dict(params), - fitted_object=None, - ) - - monkeypatch.setattr( - "grass_mil.interpretability.pipeline.run_clustering", _capture_run_clustering - ) - cfg = { - "reduction": {"enabled": True, "method": "pca", "params": {"n_components": 1}}, - "clustering": { - "enabled": True, - "method": "kmeans", - "cluster_on_pca": False, - "params": {"n_clusters": 2, "random_state": 42}, - }, - "plugins": {"enabled": [], "params": {}}, - } - bundle = run_interpretability_pipeline(dataset, cfg, artifacts_dir=tmp_path / "artifacts") - assert "x" in captured - assert captured["x"].shape == expected_raw.shape - assert np.allclose(captured["x"], expected_raw) - assert bundle.niche_feature_reduction is None - - -def test_pipeline_exposes_cluster_feature_reduction_when_cluster_on_pca_enabled( - tmp_path: Path, -) -> None: - instance_path, _ = _write_minimal_tables(tmp_path) - dataset = load_interpretability_dataset( - instance_table_path=instance_path, - id_column="instance_id", - bag_id_column="bag_id", - cell_type_column="cell_type", - ) - cfg = { - "reduction": {"enabled": True, "method": "pca", "params": {"n_components": 2}}, - "clustering": { - "enabled": True, - "method": "agglomerative", - "cluster_on_pca": True, - "pca_components": 1, - "params": {"n_clusters": 2, "linkage": "ward"}, - }, - "plugins": {"enabled": [], "params": {}}, - } - bundle = run_interpretability_pipeline(dataset, cfg, artifacts_dir=tmp_path / "artifacts") - assert bundle.niche_feature_reduction is not None - assert bundle.niche_feature_reduction.method == "pca" - assert bundle.niche_feature_reduction.embedding.shape[1] == 1 - - -def test_report_cli_smoke( - monkeypatch, - cfg_interpret: DictConfig, - tmp_path: Path, # type: ignore[valid-type] -) -> None: - instance_path, spatial_path = _write_minimal_tables(tmp_path) - monkeypatch.setattr( - "grass_mil.interpretability.reporting.render.render_pdf_via_playwright", - lambda html_path, pdf_path: pdf_path, # type: ignore[lambda-assign] - ) - monkeypatch.setattr( - "grass_mil.interpretability.reporting.render.export_plotly_snapshots", - lambda figures, output_dir, scale=2.0: {}, # type: ignore[lambda-assign] - ) - with open_dict(cfg_interpret): - cfg_interpret.data.instance_table = str(instance_path) - cfg_interpret.data.spatial_table = str(spatial_path) - # Keep smoke test independent of optional umap-learn dependency. - cfg_interpret.reduction.method = "pca" - cfg_interpret.reduction.params = {"n_components": 2, "random_state": 42} - cfg_interpret.plugins.enabled = [ - "niche_profiles", - "niche_label_moran", - ] - cfg_interpret.clustering.method = "agglomerative" - cfg_interpret.clustering.cluster_on_pca = False - cfg_interpret.clustering.pca_components = None - cfg_interpret.clustering.params = {"n_clusters": 2, "linkage": "ward"} - cfg_interpret.report.pdf.enabled = True - HydraConfig().set_config(cfg_interpret) - summary, _ = run_report(cfg_interpret) - assert Path(summary["output_dir"]).exists() - assert Path(summary["output_dir"], "report_summary.json").exists() - - -def test_pdf_renderer_raises_when_unavailable(monkeypatch, tmp_path: Path) -> None: - from grass_mil.interpretability.reporting.pdf import render_pdf_via_playwright - - real_import = __import__ - - def _fake_import(name, *args, **kwargs): - if name.startswith("playwright"): - raise ImportError("missing playwright") - return real_import(name, *args, **kwargs) - - monkeypatch.setattr("builtins.__import__", _fake_import) - with pytest.raises(RuntimeError, match="Playwright is unavailable"): - render_pdf_via_playwright( - html_path=tmp_path / "foo.html", - pdf_path=tmp_path / "foo.pdf", - ) - - -def test_resolve_color_values_handles_numeric_and_categorical_columns() -> None: - from grass_mil.interpretability.reporting.plotly_builders import _resolve_color_values - - numeric_values, numeric_ticks = _resolve_color_values(pd.Series([0.5, 1.5, 2.5])) - np.testing.assert_allclose(numeric_values, [0.5, 1.5, 2.5]) - assert numeric_ticks == {} - - # A string column (``condition`` routinely is one) must be factorized rather - # than raising, with the original labels preserved as colorbar ticks. - cat_values, cat_ticks = _resolve_color_values(pd.Series(["X", "X", "Y"])) - np.testing.assert_allclose(cat_values, [0.0, 0.0, 1.0]) - assert cat_ticks["ticktext"] == ["X", "Y"] - assert cat_ticks["tickvals"] == [0, 1] - - -def test_multi_attribute_scatters_plot_string_condition_column(tmp_path: Path) -> None: - node_path, spatial_path = _write_minimal_tables(tmp_path) - dataset = load_interpretability_dataset( - instance_table_path=node_path, - spatial_table_path=spatial_path, - id_column="instance_id", - bag_id_column="bag_id", - cell_type_column="cell_type", - ) - cfg = { - "reduction": {"enabled": True, "method": "pca", "params": {"n_components": 2}}, - "clustering": { - "enabled": True, - "method": "agglomerative", - "params": {"n_clusters": 2, "linkage": "ward"}, - }, - "plugins": {"enabled": []}, - } - bundle = run_interpretability_pipeline(dataset, cfg, artifacts_dir=tmp_path / "artifacts") - - figs = dict(bundle_figures(bundle)) - condition_figs = [name for name in figs if name.endswith("_condition")] - assert condition_figs, f"expected a condition-colored scatter, got {sorted(figs)}" - - marker = figs[condition_figs[0]].data[0].marker - assert list(marker.colorbar.ticktext) == ["X", "Y"] diff --git a/tests/interpretability/test_plugins_and_tier2.py b/tests/interpretability/test_plugins_and_tier2.py deleted file mode 100644 index f5df608..0000000 --- a/tests/interpretability/test_plugins_and_tier2.py +++ /dev/null @@ -1,861 +0,0 @@ -from __future__ import annotations - -import time - -import numpy as np -import pandas as pd -import pytest - -from grass_mil.contracts import NicheSummary, InterpretabilityDataset -from grass_mil.interpretability.plugins.base import PluginContext -from grass_mil.interpretability.plugins.builtin import create_builtin_registry -from grass_mil.interpretability.plugins.registry import create_plugin_registry -from grass_mil.interpretability.tier2.filtration import compute_filtration_curves -from grass_mil.interpretability.tier2.neighborhood import ( - NeighborhoodEnrichmentResult, - run_diff_neighborhood_enrichment, - run_neighborhood_enrichment, -) - - -def _sample_dataset() -> InterpretabilityDataset: - node_table = pd.DataFrame( - { - "instance_id": [f"n{i}" for i in range(6)], - "bag_id": ["b0", "b0", "b0", "b1", "b1", "b1"], - "cell_type": ["A", "A", "B", "A", "B", "B"], - "comp_A": [1.0, 1.0, 0.0, 1.0, 0.0, 0.0], - "comp_B": [0.0, 0.0, 1.0, 0.0, 1.0, 1.0], - "score": [0.3, 0.4, 0.9, 0.1, 0.8, 0.7], - "attention": [0.2, 0.5, 0.3, 0.1, 0.2, 0.7], - "sample_id": ["sx", "sx", "sx", "sy", "sy", "sy"], - "condition": ["X", "X", "X", "Y", "Y", "Y"], - "center_x": [0.0, 0.1, 0.2, 1.0, 1.1, 1.2], - "center_y": [0.0, 0.2, 0.1, 1.0, 1.2, 1.1], - "inst_emb_0": [0.0, 0.1, 0.2, 1.0, 1.1, 1.2], - "inst_emb_1": [0.0, 0.2, 0.1, 1.0, 1.2, 1.1], - } - ) - edges = pd.DataFrame( - { - "source_id": ["n0", "n1", "n2", "n3", "n4", "n5"], - "target_id": ["n1", "n2", "n0", "n4", "n5", "n3"], - "distance": [0.1, 0.3, 0.2, 0.8, 0.7, 0.9], - "weight": [1, 1, 1, 1, 1, 1], - } - ) - return InterpretabilityDataset( - instance_table=node_table, - spatial_table=edges, - id_column="instance_id", - bag_id_column="bag_id", - cell_type_column="cell_type", - ) - - -def test_plugin_registry_and_builtin_execution() -> None: - registry = create_builtin_registry() - assert "niche_profiles" in registry.list() - assert "tissue_graph" in registry.list() - ds = _sample_dataset() - labels = np.array([0, 0, 1, 1, 1, 0]) - plugin = registry.get("niche_profiles") - out = plugin.run(ds, PluginContext(state={"niche_labels": labels})) - assert out.name == "niche_profiles" - assert "composition" in out.payload - - -def test_builtin_tissue_graph_plugin_execution() -> None: - registry = create_builtin_registry() - ds = _sample_dataset() - labels = np.array([0, 0, 1, 1, 1, 0]) - plugin = registry.get("tissue_graph") - out = plugin.run( - ds, - PluginContext(state={"niche_labels": labels}), - sample_column="sample_id", - sample_value="sx", - id_column="instance_id", - ) - assert out.name == "tissue_graph" - assert "tissue_graph_view" in out.payload - assert "tissue_graph_meta" in out.payload - view = out.payload["tissue_graph_view"] - assert view.metadata["sample_value"] == "sx" - assert len(out.sections) == 1 - - -def test_builtin_tissue_graph_plugin_requires_sample_value() -> None: - registry = create_builtin_registry() - ds = _sample_dataset() - labels = np.array([0, 0, 1, 1, 1, 0]) - plugin = registry.get("tissue_graph") - with pytest.raises(ValueError, match="requires a non-empty sample_value"): - plugin.run( - ds, - PluginContext(state={"niche_labels": labels}), - sample_column="sample_id", - id_column="instance_id", - ) - - -def test_builtin_tissue_graph_plugin_fails_when_coords_missing() -> None: - registry = create_builtin_registry() - ds = _sample_dataset() - missing_coords_ds = InterpretabilityDataset( - instance_table=ds.instance_table.drop(columns=["center_x", "center_y"]), - spatial_table=ds.spatial_table, - id_column=ds.id_column, - bag_id_column=ds.bag_id_column, - cell_type_column=ds.cell_type_column, - ) - labels = np.array([0, 0, 1, 1, 1, 0]) - plugin = registry.get("tissue_graph") - with pytest.raises(ValueError, match="missing required columns"): - plugin.run( - missing_coords_ds, - PluginContext(state={"niche_labels": labels}), - sample_column="sample_id", - sample_value="sx", - id_column="instance_id", - ) - - -def test_cluster_profiles_reuses_context_cluster_summary(monkeypatch) -> None: - registry = create_builtin_registry() - ds = _sample_dataset() - labels = np.array([0, 0, 1, 1, 1, 0]) - summary = NicheSummary( - niche_labels=labels.copy(), - composition=pd.DataFrame({"comp_A": [0.1, 0.2]}, index=[0, 1]), - enrichment=pd.DataFrame({"comp_A": [1.0, -1.0]}, index=[0, 1]), - niche_counts=pd.Series({0: 2, 1: 4}), - ) - - def _should_not_run(*args, **kwargs): # type: ignore[no-untyped-def] - raise AssertionError("niche_composition_summary should not be recomputed") - - monkeypatch.setattr( - "grass_mil.interpretability.plugins.builtin.niche_composition_summary", _should_not_run - ) - plugin = registry.get("niche_profiles") - out = plugin.run( - ds, - PluginContext(state={"niche_labels": labels, "niche_summary": summary}), - ) - assert out.payload["composition"].equals(summary.composition) - assert out.payload["enrichment"].equals(summary.enrichment) - - -def test_plugin_registry_instances_are_isolated() -> None: - class _DummyPlugin: - name = "dummy" - - def required_inputs(self) -> list[str]: - return [] - - def run(self, dataset, context: PluginContext, **params): # type: ignore[no-untyped-def] - raise NotImplementedError - - left = create_plugin_registry() - right = create_plugin_registry() - left.register(_DummyPlugin()) - - assert "dummy" in left.list() - assert "dummy" not in right.list() - - -def test_plugin_registry_rejects_duplicate_names_by_default() -> None: - class _PluginA: - name = "duplicate" - - def required_inputs(self) -> list[str]: - return [] - - def run(self, dataset, context: PluginContext, **params): # type: ignore[no-untyped-def] - raise NotImplementedError - - class _PluginB: - name = "duplicate" - - def required_inputs(self) -> list[str]: - return [] - - def run(self, dataset, context: PluginContext, **params): # type: ignore[no-untyped-def] - raise NotImplementedError - - registry = create_plugin_registry() - registry.register(_PluginA()) - with pytest.raises(ValueError, match="already registered"): - registry.register(_PluginB()) - - -def test_plugin_registry_allows_explicit_replace_for_duplicate_names() -> None: - class _PluginA: - name = "duplicate" - - def required_inputs(self) -> list[str]: - return [] - - def run(self, dataset, context: PluginContext, **params): # type: ignore[no-untyped-def] - raise NotImplementedError - - class _PluginB: - name = "duplicate" - - def required_inputs(self) -> list[str]: - return [] - - def run(self, dataset, context: PluginContext, **params): # type: ignore[no-untyped-def] - raise NotImplementedError - - registry = create_plugin_registry() - first = _PluginA() - second = _PluginB() - registry.register(first) - registry.register(second, replace=True) - assert registry.get("duplicate") is second - - -def test_tier2_neighborhood_and_filtration() -> None: - ds = _sample_dataset() - labels = np.array([0, 0, 1, 1, 1, 0]) - frame = ds.instance_table.copy() - frame["niche_label"] = labels - enr = run_neighborhood_enrichment( - frame, - ds.spatial_table, # type: ignore[arg-type] - label_column="niche_label", - id_column="instance_id", - n_perms=4, - random_state=2, - ) - assert enr.enrichment.shape[0] > 0 - curves = compute_filtration_curves( - frame, - ds.spatial_table, # type: ignore[arg-type] - thresholds=np.array([0.2, 0.5, 1.0]), - niche_column="niche_label", - cell_type_column="cell_type", - id_column="instance_id", - distance_column="distance", - ) - assert len(curves.curves) > 0 - - -def test_filtration_subgraph_centric_counts_both_edge_endpoints() -> None: - node_table = pd.DataFrame( - { - "instance_id": ["n0", "n1", "n2"], - "niche_label": [0, 0, 1], - "cell_type": ["A", "B", "A"], - } - ) - spatial_table = pd.DataFrame( - { - "source_id": ["n0", "n0", "n1"], - "target_id": ["n1", "n2", "n2"], - "distance": [0.10, 0.20, 0.20], - } - ) - curves = compute_filtration_curves( - node_table, - spatial_table, - thresholds=np.array([0.15, 0.25]), - niche_column="niche_label", - cell_type_column="cell_type", - id_column="instance_id", - distance_column="distance", - scale_within_cluster=False, - ) - # At 0.15 only n0->n1 contributes in niche 0: A=1 (n0), B=1 (n1). - # At 0.25: - # - subgraph n0 contributes unique nodes {n0,n1,n2}: A=2, B=1 - # - subgraph n1 contributes unique nodes {n1,n2}: A=1, B=1 - # Total for niche 0: A=3, B=2. - assert np.allclose(curves.curves["0"]["A"], np.array([1.0, 3.0])) - assert np.allclose(curves.curves["0"]["B"], np.array([1.0, 2.0])) - - -def test_neighborhood_undirected_option_adds_reverse_direction() -> None: - node_table = pd.DataFrame( - { - "instance_id": ["n0", "n1"], - "niche_label": ["A", "B"], - } - ) - spatial_table = pd.DataFrame( - { - "source_id": ["n0"], - "target_id": ["n1"], - "weight": [1.0], - } - ) - directed = run_neighborhood_enrichment( - node_table, - spatial_table, - label_column="niche_label", - id_column="instance_id", - n_perms=0, - undirected=False, - warn_analytical=False, - ) - undirected = run_neighborhood_enrichment( - node_table, - spatial_table, - label_column="niche_label", - id_column="instance_id", - n_perms=0, - undirected=True, - warn_analytical=False, - ) - assert directed.observed.loc["A", "B"] == 1.0 - assert directed.observed.loc["B", "A"] == 0.0 - assert undirected.observed.loc["A", "B"] == 1.0 - assert undirected.observed.loc["B", "A"] == 1.0 - - -def test_neighborhood_undirected_deduplicates_already_bidirectional_edges() -> None: - node_table = pd.DataFrame( - { - "instance_id": ["n0", "n1"], - "niche_label": ["A", "B"], - } - ) - spatial_table = pd.DataFrame( - { - "source_id": ["n0", "n1"], - "target_id": ["n1", "n0"], - } - ) - out = run_neighborhood_enrichment( - node_table, - spatial_table, - label_column="niche_label", - id_column="instance_id", - n_perms=0, - undirected=True, - warn_analytical=False, - ) - assert out.observed.loc["A", "B"] == 1.0 - assert out.observed.loc["B", "A"] == 1.0 - - -def test_neighborhood_ignores_weight_semantics() -> None: - node_table = pd.DataFrame( - { - "instance_id": ["n0", "n1", "n2"], - "niche_label": ["A", "B", "B"], - } - ) - spatial_table = pd.DataFrame( - { - "source_id": ["n0", "n0"], - "target_id": ["n1", "n2"], - "weight": [100.0, 0.001], - } - ) - out = run_neighborhood_enrichment( - node_table, - spatial_table, - label_column="niche_label", - id_column="instance_id", - n_perms=0, - undirected=False, - warn_analytical=False, - ) - # Squidpy-style neighborhood enrichment is count-based, not weighted. - assert out.observed.loc["A", "B"] == 2.0 - - -def test_neighborhood_enrichment_mode_is_tunable() -> None: - node_table = pd.DataFrame( - { - "instance_id": ["n0", "n1"], - "niche_label": ["A", "B"], - } - ) - spatial_table = pd.DataFrame( - { - "source_id": ["n0"], - "target_id": ["n1"], - } - ) - zscore = run_neighborhood_enrichment( - node_table, - spatial_table, - label_column="niche_label", - id_column="instance_id", - n_perms=0, - undirected=False, - warn_analytical=False, - enrichment_mode="zscore", - ) - difference = run_neighborhood_enrichment( - node_table, - spatial_table, - label_column="niche_label", - id_column="instance_id", - n_perms=0, - undirected=False, - warn_analytical=False, - enrichment_mode="obs-exp", - ) - log2fc = run_neighborhood_enrichment( - node_table, - spatial_table, - label_column="niche_label", - id_column="instance_id", - n_perms=0, - undirected=False, - warn_analytical=False, - enrichment_mode="log2fc", - ) - assert zscore.enrichment.loc["A", "B"] == pytest.approx(1.7320508075688774) - assert difference.enrichment.loc["A", "B"] == pytest.approx(0.75) - assert log2fc.enrichment.loc["A", "B"] == pytest.approx(2.0) - - -def test_neighborhood_enrichment_mode_rejects_unknown_values() -> None: - node_table = pd.DataFrame( - { - "instance_id": ["n0", "n1"], - "niche_label": ["A", "B"], - } - ) - spatial_table = pd.DataFrame( - { - "source_id": ["n0"], - "target_id": ["n1"], - } - ) - with pytest.raises(ValueError, match="Unsupported enrichment_mode"): - run_neighborhood_enrichment( - node_table, - spatial_table, - label_column="niche_label", - id_column="instance_id", - n_perms=0, - warn_analytical=False, - enrichment_mode="unknown_mode", - ) - - -def test_diff_neighborhood_returns_pairwise_differences_with_pvalues() -> None: - node_table = pd.DataFrame( - { - "instance_id": ["x0", "x1", "y0", "y1"], - "niche_label": ["A", "B", "A", "B"], - "sample_id": ["sx", "sx", "sy", "sy"], - "condition": ["X", "X", "Y", "Y"], - } - ) - spatial_table = pd.DataFrame( - { - "source_id": ["x0", "x1", "y1", "y0"], - "target_id": ["x1", "x0", "y0", "y1"], - "weight": [1.0, 1.0, 1.0, 1.0], - } - ) - out = run_diff_neighborhood_enrichment( - node_table, - spatial_table, - label_column="niche_label", - condition_column="condition", - id_column="instance_id", - n_perms=8, - random_state=3, - undirected=False, - warn_analytical=False, - ) - assert "X_Y" in out - diff = out["X_Y"] - assert diff.enrichment.shape == (2, 2) - assert diff.pvalues is not None - assert diff.pvalues.shape == (2, 2) - - -def test_diff_neighborhood_pvalues_are_centered_on_permutation_mean(monkeypatch) -> None: - calls = {"i": 0} - # First two calls are baseline left/right. Remaining calls are (left,right) per permutation. - per_call_scores = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 2.0, 0.0, 3.0, 0.0] - - def _fake_run_neighborhood(*args, **kwargs): # type: ignore[no-untyped-def] - idx = calls["i"] - calls["i"] += 1 - score = per_call_scores[idx] - frame = pd.DataFrame([[score]], index=["A"], columns=["A"]) - zeros = pd.DataFrame([[0.0]], index=["A"], columns=["A"]) - return NeighborhoodEnrichmentResult( - enrichment=frame, - observed=zeros, - expected=zeros, - pvalues=None, - ) - - monkeypatch.setattr( - "grass_mil.interpretability.tier2.neighborhood.run_neighborhood_enrichment", - _fake_run_neighborhood, - ) - node_table = pd.DataFrame( - { - "instance_id": ["x0", "x1", "y0", "y1"], - "niche_label": ["A", "A", "A", "A"], - "sample_id": ["sx0", "sx1", "sy0", "sy1"], - "condition": ["X", "X", "Y", "Y"], - } - ) - spatial_table = pd.DataFrame({"source_id": ["x0"], "target_id": ["x1"]}) - out = run_diff_neighborhood_enrichment( - node_table, - spatial_table, - label_column="niche_label", - condition_column="condition", - permutation_group_column="sample_id", - n_perms=4, - warn_analytical=False, - ) - p = float(out["X_Y"].pvalues.loc["A", "A"]) # type: ignore[union-attr] - assert p == pytest.approx(1.0) - assert p != pytest.approx(0.8) - - -def test_diff_neighborhood_respects_enrichment_mode() -> None: - node_table = pd.DataFrame( - { - "instance_id": ["x0", "x1", "y0", "y1"], - "niche_label": ["A", "B", "A", "B"], - "condition": ["X", "X", "Y", "Y"], - } - ) - spatial_table = pd.DataFrame( - { - "source_id": ["x0", "y1"], - "target_id": ["x1", "y0"], - } - ) - zscore_out = run_diff_neighborhood_enrichment( - node_table, - spatial_table, - label_column="niche_label", - condition_column="condition", - id_column="instance_id", - n_perms=0, - undirected=False, - warn_analytical=False, - enrichment_mode="zscore", - ) - diff_out = run_diff_neighborhood_enrichment( - node_table, - spatial_table, - label_column="niche_label", - condition_column="condition", - id_column="instance_id", - n_perms=0, - undirected=False, - warn_analytical=False, - enrichment_mode="obs-exp", - ) - assert zscore_out["X_Y"].enrichment.loc["A", "B"] == pytest.approx(2.309401076758503) - assert diff_out["X_Y"].enrichment.loc["A", "B"] == pytest.approx(1.0) - - -def test_diff_neighborhood_requires_permutation_group_column_for_permutation_mode() -> None: - node_table = pd.DataFrame( - { - "instance_id": ["x0", "x1", "y0", "y1"], - "niche_label": ["A", "B", "A", "B"], - "condition": ["X", "X", "Y", "Y"], - } - ) - spatial_table = pd.DataFrame( - { - "source_id": ["x0", "x1", "y1", "y0"], - "target_id": ["x1", "x0", "y0", "y1"], - } - ) - with pytest.raises(ValueError, match="Missing permutation group column"): - run_diff_neighborhood_enrichment( - node_table, - spatial_table, - label_column="niche_label", - condition_column="condition", - n_perms=4, - warn_analytical=False, - ) - - -def test_diff_neighborhood_rejects_mixed_condition_groups() -> None: - node_table = pd.DataFrame( - { - "instance_id": ["x0", "x1", "y0", "y1"], - "niche_label": ["A", "B", "A", "B"], - "sample_id": ["s0", "s0", "s1", "s1"], - "condition": ["X", "Y", "X", "Y"], - } - ) - spatial_table = pd.DataFrame( - { - "source_id": ["x0", "x1", "y1", "y0"], - "target_id": ["x1", "x0", "y0", "y1"], - } - ) - with pytest.raises(ValueError, match="must belong to exactly one condition"): - run_diff_neighborhood_enrichment( - node_table, - spatial_table, - label_column="niche_label", - condition_column="condition", - permutation_group_column="sample_id", - n_perms=4, - warn_analytical=False, - ) - - -def test_neighborhood_warns_in_analytical_mode() -> None: - node_table = pd.DataFrame( - { - "instance_id": ["n0", "n1"], - "niche_label": ["A", "B"], - } - ) - spatial_table = pd.DataFrame( - { - "source_id": ["n0"], - "target_id": ["n1"], - } - ) - with pytest.warns(UserWarning, match="analytical expected/std"): - run_neighborhood_enrichment( - node_table, - spatial_table, - label_column="niche_label", - id_column="instance_id", - n_perms=0, - ) - - -def test_diff_neighborhood_warns_about_analytical_baseline() -> None: - node_table = pd.DataFrame( - { - "instance_id": ["x0", "x1", "y0", "y1"], - "niche_label": ["A", "B", "A", "B"], - "condition": ["X", "X", "Y", "Y"], - } - ) - spatial_table = pd.DataFrame( - { - "source_id": ["x0", "x1", "y1", "y0"], - "target_id": ["x1", "x0", "y0", "y1"], - } - ) - with pytest.warns(UserWarning, match="analytical neighborhood model"): - run_diff_neighborhood_enrichment( - node_table, - spatial_table, - label_column="niche_label", - condition_column="condition", - id_column="instance_id", - n_perms=0, - ) - - -def _legacy_filtration_curves( - node_table: pd.DataFrame, - spatial_table: pd.DataFrame, - *, - thresholds: np.ndarray, - niche_column: str, - cell_type_column: str, - id_column: str, - distance_column: str, -) -> dict[str, dict[str, np.ndarray]]: - node_to_cell_type = node_table.set_index(id_column)[cell_type_column].astype(str).to_dict() - node_to_cluster = node_table.set_index(id_column)[niche_column].astype(str).to_dict() - edges = spatial_table.copy() - edges["source_id"] = edges["source_id"].astype(str) - edges["target_id"] = edges["target_id"].astype(str) - edges = edges[edges["source_id"].isin(node_to_cluster.keys())].copy() - edges["source_cluster"] = edges["source_id"].map(node_to_cluster) - - niches = sorted(node_table[niche_column].astype(str).unique()) - cell_types = sorted(node_table[cell_type_column].astype(str).unique()) - curves: dict[str, dict[str, np.ndarray]] = { - c: {ct: np.zeros_like(thresholds, dtype=float) for ct in cell_types} for c in niches - } - for niche in niches: - c_edges = edges[edges["source_cluster"] == str(niche)] - by_subgraph = { - str(sub_id): frame for sub_id, frame in c_edges.groupby("source_id", sort=False) - } - for i, thr in enumerate(thresholds): - threshold_counts: dict[str, float] = {ct: 0.0 for ct in cell_types} - for subgraph_edges in by_subgraph.values(): - selected = subgraph_edges[subgraph_edges[distance_column] <= float(thr)] - if selected.empty: - continue - unique_node_ids = set(selected["source_id"].tolist()) | set( - selected["target_id"].tolist() - ) - for node_id in unique_node_ids: - ct = node_to_cell_type.get(str(node_id)) - if ct is not None: - threshold_counts[ct] += 1.0 - for ct in cell_types: - curves[niche][ct][i] = float(threshold_counts.get(ct, 0.0)) - return curves - - -def test_filtration_vectorized_matches_legacy_counts() -> None: - rng = np.random.default_rng(19) - n_nodes = 32 - node_ids = [f"n{i}" for i in range(n_nodes)] - node_table = pd.DataFrame( - { - "instance_id": node_ids, - "niche_label": rng.choice(["0", "1", "2"], size=n_nodes), - "cell_type": rng.choice(["A", "B", "C"], size=n_nodes), - } - ) - edge_count = 120 - spatial_table = pd.DataFrame( - { - "source_id": rng.choice(node_ids, size=edge_count), - "target_id": rng.choice(node_ids, size=edge_count), - "distance": rng.uniform(0.0, 1.0, size=edge_count), - } - ) - thresholds = np.array([0.6, 0.2, 0.9, 0.1, 0.4], dtype=float) - - legacy = _legacy_filtration_curves( - node_table, - spatial_table, - thresholds=thresholds, - niche_column="niche_label", - cell_type_column="cell_type", - id_column="instance_id", - distance_column="distance", - ) - out = compute_filtration_curves( - node_table, - spatial_table, - thresholds=thresholds, - niche_column="niche_label", - cell_type_column="cell_type", - id_column="instance_id", - distance_column="distance", - scale_within_cluster=False, - ) - for niche, ct_map in legacy.items(): - for cell_type, values in ct_map.items(): - assert np.allclose(out.curves[niche][cell_type], values) - - -def test_filtration_runtime_sanity() -> None: - rng = np.random.default_rng(11) - n_nodes = 220 - node_ids = [f"n{i}" for i in range(n_nodes)] - node_table = pd.DataFrame( - { - "instance_id": node_ids, - "niche_label": rng.choice(["0", "1", "2", "3"], size=n_nodes), - "cell_type": rng.choice(["A", "B", "C", "D"], size=n_nodes), - } - ) - edge_count = 2400 - spatial_table = pd.DataFrame( - { - "source_id": rng.choice(node_ids, size=edge_count), - "target_id": rng.choice(node_ids, size=edge_count), - "distance": rng.uniform(0.0, 55.0, size=edge_count), - } - ) - thresholds = np.linspace(0.0, 55.0, 300) - t0 = time.perf_counter() - out = compute_filtration_curves( - node_table, - spatial_table, - thresholds=thresholds, - niche_column="niche_label", - cell_type_column="cell_type", - id_column="instance_id", - distance_column="distance", - ) - elapsed = time.perf_counter() - t0 - assert len(out.curves) > 0 - assert elapsed < 8.0 - - -def _attribution_dataset() -> InterpretabilityDataset: - """Sample dataset carrying per-class attention and instance logits.""" - ds = _sample_dataset() - table = ds.instance_table.copy() - # Two attention channels, each normalised within its bag. - table["attention_c0"] = [2 / 3, 1 / 6, 1 / 6, 0.2, 0.3, 0.5] - table["attention_c1"] = [0.2, 0.5, 0.3, 1 / 6, 1 / 6, 2 / 3] - table["logit_0"] = [0.5, -0.2, 0.9, 0.1, -0.4, 0.7] - table["logit_1"] = [-0.3, 0.8, 0.2, 0.6, 0.1, -0.5] - return InterpretabilityDataset( - instance_table=table, - spatial_table=ds.spatial_table, - id_column=ds.id_column, - bag_id_column=ds.bag_id_column, - cell_type_column=ds.cell_type_column, - ) - - -def test_margin_attribution_plugin_reports_per_niche_intervals() -> None: - registry = create_builtin_registry() - ds = _attribution_dataset() - labels = np.array([0, 0, 1, 1, 0, 1]) - - out = registry.get("margin_attribution").run( - ds, PluginContext(state={"niche_labels": labels}), n_bootstrap=25 - ) - assert out.name == "margin_attribution" - summary = out.payload["per_niche"] - # (niche_label, class_index); a binary head summarises the positive class. - assert set(summary.index) == {(0, 1), (1, 1)} - for column in ("margin_signed_share", "prevalence"): - assert column in summary.columns - assert f"{column}_lo" in summary.columns - - -def test_margin_attribution_requires_instance_logits() -> None: - registry = create_builtin_registry() - ds = _sample_dataset() # has no logit_* columns - with pytest.raises(ValueError, match="instance logit columns"): - registry.get("margin_attribution").run( - ds, PluginContext(state={"niche_labels": np.zeros(6, dtype=int)}) - ) - - -def test_cluster_agreement_plugin_compares_stored_labelings() -> None: - registry = create_builtin_registry() - ds = _sample_dataset() - table = ds.instance_table.copy() - table["prior_labels"] = [0, 0, 0, 1, 1, 1] - ds = InterpretabilityDataset( - instance_table=table, - spatial_table=ds.spatial_table, - id_column=ds.id_column, - bag_id_column=ds.bag_id_column, - cell_type_column=ds.cell_type_column, - ) - - out = registry.get("niche_agreement").run( - ds, - PluginContext(state={"niche_labels": np.array([0, 0, 0, 1, 1, 1])}), - label_columns=["prior_labels"], - ) - metrics = out.payload["pairwise_metrics"] - # Identical partitions agree perfectly. - assert metrics["ari"].iloc[0] == pytest.approx(1.0) - - -def test_cluster_agreement_plugin_requires_a_second_labeling() -> None: - registry = create_builtin_registry() - with pytest.raises(ValueError, match="at least one additional labeling"): - registry.get("niche_agreement").run( - _sample_dataset(), PluginContext(state={"niche_labels": np.zeros(6, dtype=int)}) - ) diff --git a/tests/interpretability/test_spatial_statistics.py b/tests/interpretability/test_spatial_statistics.py deleted file mode 100644 index 9b6d5e4..0000000 --- a/tests/interpretability/test_spatial_statistics.py +++ /dev/null @@ -1,242 +0,0 @@ -"""Contracts for Moran's I, Ripley's cross-L, and cross-space agreement. - -Each statistic is checked against a case with a known answer rather than -against a golden value, so the tests state what the statistic means. -""" - -from __future__ import annotations - -import numpy as np -import pandas as pd -import pytest - -from grass_mil.interpretability.core.agreement import ( - niche_jaccard_matrix, - compute_niche_agreement, -) -from grass_mil.interpretability.tier2.autocorrelation import ( - morans_i, - run_morans_i, -) -from grass_mil.interpretability.tier2.ripley import ( - aggregate_ripley, - resolve_radii, - ripley_cross_l, -) - - -def _chain_adjacency(n: int): - from scipy.sparse import coo_matrix - - rows = np.concatenate([np.arange(n - 1), np.arange(1, n)]) - cols = np.concatenate([np.arange(1, n), np.arange(n - 1)]) - return coo_matrix((np.ones(rows.size), (rows, cols)), shape=(n, n)).tocsr() - - -def test_morans_i_is_high_for_a_smooth_field_and_low_when_shuffled() -> None: - n = 60 - adjacency = _chain_adjacency(n) - - smooth = np.arange(n, dtype=float) - clustered = morans_i(smooth, adjacency) - assert clustered > 0.9, "a monotone field along a chain is strongly autocorrelated" - - rng = np.random.default_rng(0) - shuffled = np.mean([morans_i(rng.permutation(smooth), adjacency) for _ in range(50)]) - assert abs(shuffled) < 0.2, "permuting the field destroys autocorrelation" - - -def test_morans_i_is_negative_for_an_alternating_field() -> None: - n = 60 - values = np.array([i % 2 for i in range(n)], dtype=float) - assert morans_i(values, _chain_adjacency(n)) < -0.9 - - -def test_morans_i_is_undefined_for_a_constant_feature() -> None: - assert np.isnan(morans_i(np.ones(10), _chain_adjacency(10))) - - -def _autocorrelation_tables(n_per_group: int = 30): - rows = [] - edges = [] - for group in ("a", "b"): - for i in range(n_per_group): - node = f"{group}_{i}" - # Group a carries a smooth gradient; group b alternates. - value = float(i) if group == "a" else float(i % 2) - rows.append({"instance_id": node, "niche_label": group, "feature": value}) - if i + 1 < n_per_group: - edges.append({"source_id": node, "target_id": f"{group}_{i + 1}"}) - return pd.DataFrame(rows), pd.DataFrame(edges) - - -def test_run_morans_i_separates_smooth_from_alternating_groups() -> None: - nodes, edges = _autocorrelation_tables() - result = run_morans_i( - nodes, - edges, - feature_columns=["feature"], - group_column="niche_label", - n_perms=50, - random_state=0, - ) - assert result.statistic.loc["a", "feature"] > 0.8 - assert result.statistic.loc["b", "feature"] < -0.8 - # Both are far from the permutation null, so both are significant. - assert (result.pvalue["feature"] < 0.1).all() - assert (result.qvalue["feature"] >= result.pvalue["feature"]).all() - assert result.n_nodes.loc["a"] == 30 - - -def test_run_morans_i_skips_groups_below_the_size_threshold() -> None: - nodes = pd.DataFrame( - { - "instance_id": ["n0", "n1", "n2"], - "niche_label": ["small", "small", "other"], - "feature": [1.0, 2.0, 3.0], - } - ) - edges = pd.DataFrame({"source_id": ["n0"], "target_id": ["n1"]}) - result = run_morans_i( - nodes, edges, feature_columns=["feature"], group_column="niche_label", n_perms=5 - ) - assert np.isnan(result.statistic.loc["small", "feature"]) - - -def _poisson_points(n: int, seed: int, size: float = 100.0) -> np.ndarray: - rng = np.random.default_rng(seed) - return rng.uniform(0.0, size, size=(n, 2)) - - -def test_ripley_l_is_near_zero_for_a_poisson_pattern() -> None: - coords = _poisson_points(600, seed=0) - labels = ["x"] * len(coords) - radii = resolve_radii(coords, n_radii=20, max_fraction=0.15) - - result = ripley_cross_l(coords, labels, pairs=[("x", "x")], radii=radii) - curve = result.curves[("x", "x")] - # Complete spatial randomness gives L(r) - r == 0 up to sampling noise. - assert np.nanmax(np.abs(curve)) < 4.0 - - -def test_ripley_l_is_positive_for_a_clustered_pattern() -> None: - rng = np.random.default_rng(1) - centres = rng.uniform(10.0, 90.0, size=(8, 2)) - coords = np.vstack([centre + rng.normal(scale=1.5, size=(60, 2)) for centre in centres]) - labels = ["x"] * len(coords) - radii = resolve_radii(coords, n_radii=20, max_fraction=0.15) - - curve = ripley_cross_l(coords, labels, pairs=[("x", "x")], radii=radii).curves[("x", "x")] - assert np.nanmax(curve) > 5.0, "tight clumps must show strong positive clustering" - - -def test_ripley_self_pair_correction_matches_a_hand_computed_case() -> None: - # Four points on a unit square; every pairwise distance <= sqrt(2). - coords = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]) - radii = np.array([2.0]) - result = ripley_cross_l(coords, ["x"] * 4, pairs=[("x", "x")], radii=radii) - - # With r beyond every distance: counts = 16, minus 4 self pairs = 12, - # denominator 4*3 = 12, area = 1 -> K = 1, L = sqrt(1/pi) - 2. - expected = np.sqrt(1.0 / np.pi) - 2.0 - assert result.curves[("x", "x")][0] == pytest.approx(expected) - - -def test_aggregate_ripley_is_region_equal_and_honours_min_count() -> None: - frames = [] - for group in range(3): - coords = _poisson_points(80, seed=group) - frames.append( - pd.DataFrame( - { - "center_x": coords[:, 0], - "center_y": coords[:, 1], - "cell_type": ["x"] * 78 + ["rare"] * 2, - "region": f"r{group}", - } - ) - ) - table = pd.concat(frames, ignore_index=True) - - result = aggregate_ripley( - table, - label_column="cell_type", - group_column="region", - pairs=[("x", "x"), ("rare", "rare")], - n_radii=10, - min_count=5, - ) - assert result.n_groups[("x", "x")] == 3 - # `rare` never reaches min_count, so it contributes no curve at all. - assert result.n_groups[("rare", "rare")] == 0 - assert np.isnan(result.curves[("rare", "rare")]).all() - - -def test_aggregate_ripley_radius_grid_does_not_depend_on_group_order() -> None: - """The legacy grid came from whichever region was seen first.""" - frames = [] - for group, scale in enumerate((10.0, 100.0, 50.0)): - coords = _poisson_points(40, seed=group, size=scale) - frames.append( - pd.DataFrame( - { - "center_x": coords[:, 0], - "center_y": coords[:, 1], - "cell_type": ["x"] * len(coords), - "region": f"r{group}", - } - ) - ) - table = pd.concat(frames, ignore_index=True) - reversed_table = pd.concat(frames[::-1], ignore_index=True) - - forward = aggregate_ripley(table, label_column="cell_type", group_column="region", n_radii=8) - backward = aggregate_ripley( - reversed_table, label_column="cell_type", group_column="region", n_radii=8 - ) - np.testing.assert_allclose(forward.radii, backward.radii) - - with pytest.raises(ValueError, match="Unsupported radius_source"): - aggregate_ripley( - table, label_column="cell_type", group_column="region", radius_source="nonsense" - ) - - -def test_cluster_agreement_is_one_for_identical_and_near_zero_for_independent() -> None: - rng = np.random.default_rng(0) - labels = rng.integers(0, 4, size=300) - independent = rng.integers(0, 4, size=300) - - result = compute_niche_agreement( - {"encoder": labels, "same": labels.copy(), "random": independent} - ) - metrics = result.pairwise_metrics - assert metrics.loc[("encoder", "same"), "ari"] == pytest.approx(1.0) - assert abs(metrics.loc[("encoder", "random"), "ari"]) < 0.1 - # A relabelled partition IS the same partition: these indices are - # chance-corrected and invariant to how the niches are numbered. - relabelled = compute_niche_agreement({"a": labels, "b": (labels + 1) % 4}) - assert relabelled.pairwise_metrics.loc[("a", "b"), "ari"] == pytest.approx(1.0) - - -def test_cluster_agreement_validates_its_inputs() -> None: - with pytest.raises(ValueError, match="at least two labelings"): - compute_niche_agreement({"only": [0, 1, 2]}) - with pytest.raises(ValueError, match="same instances"): - compute_niche_agreement({"a": [0, 1, 2], "b": [0, 1]}) - with pytest.raises(ValueError, match="Unsupported agreement metric"): - compute_niche_agreement({"a": [0, 1], "b": [1, 0]}, metrics=["nonsense"]) - - -def test_cluster_jaccard_matrix_reports_overlap_fractions() -> None: - a = np.array([0, 0, 1, 1]) - b = np.array([0, 1, 1, 1]) - matrix = niche_jaccard_matrix(a, b) - - # Niche 0 of `a` = {0,1}; niche 1 of `b` = {1,2,3}; overlap {1}, union 4. - assert matrix.loc[0, 1] == pytest.approx(0.25) - # Niche 1 of `a` = {2,3} is fully inside niche 1 of `b`. - assert matrix.loc[1, 1] == pytest.approx(2 / 3) - - with pytest.raises(ValueError, match="must align"): - niche_jaccard_matrix(np.array([0, 1]), np.array([0, 1, 2])) diff --git a/tests/interpretability/test_tissue_graph.py b/tests/interpretability/test_tissue_graph.py deleted file mode 100644 index 11150f7..0000000 --- a/tests/interpretability/test_tissue_graph.py +++ /dev/null @@ -1,109 +0,0 @@ -from __future__ import annotations - -import pandas as pd -import pytest - -from grass_mil.interpretability.tier2.tissue_graph import ( - build_tissue_graph_figure, - prepare_tissue_graph_view, -) - -go = pytest.importorskip("plotly.graph_objects") - - -def _sample_tables() -> tuple[pd.DataFrame, pd.DataFrame]: - node_table = pd.DataFrame( - { - "instance_id": ["n0", "n1", "n2", "n3", "n4"], - "sample_id": ["sx", "sx", "sx", "sy", "sy"], - "center_x": [0.0, 1.0, 2.0, 10.0, 11.0], - "center_y": [0.0, 1.0, 0.5, 10.0, 11.0], - "niche_label": [0, 1, 1, 0, 2], - } - ) - spatial_table = pd.DataFrame( - { - "source_id": ["n0", "n1", "n2", "n3"], - "target_id": ["n1", "n2", "n3", "n4"], - "distance": [1.0, 1.2, 2.0, 0.9], - } - ) - return node_table, spatial_table - - -def test_prepare_tissue_graph_view_happy_path_filters_to_selected_sample() -> None: - node_table, spatial_table = _sample_tables() - view = prepare_tissue_graph_view( - node_table, - spatial_table, - sample_column="sample_id", - sample_value="sx", - id_column="instance_id", - x_column="center_x", - y_column="center_y", - label_column="niche_label", - ) - - assert len(view.nodes) == 3 - assert set(view.nodes["instance_id"].tolist()) == {"n0", "n1", "n2"} - assert len(view.edges) == 2 - assert set(view.edges["source_id"].tolist()) == {"n0", "n1"} - assert view.metadata["sample_column"] == "sample_id" - assert view.metadata["sample_value"] == "sx" - assert view.metadata["node_count"] == 3 - assert view.metadata["edge_count"] == 2 - assert view.metadata["niche_counts"] == {"0": 1, "1": 2} - - -def test_prepare_tissue_graph_view_fails_for_missing_coords() -> None: - node_table, spatial_table = _sample_tables() - with pytest.raises(ValueError, match="missing required columns"): - prepare_tissue_graph_view( - node_table.drop(columns=["center_x"]), - spatial_table, - sample_column="sample_id", - sample_value="sx", - id_column="instance_id", - x_column="center_x", - y_column="center_y", - label_column="niche_label", - ) - - -def test_prepare_tissue_graph_view_fails_for_missing_or_empty_sample_selector() -> None: - node_table, spatial_table = _sample_tables() - with pytest.raises(ValueError, match="sample_value is required"): - prepare_tissue_graph_view( - node_table, - spatial_table, - sample_column="sample_id", - sample_value=None, - id_column="instance_id", - ) - with pytest.raises(ValueError, match="No rows found"): - prepare_tissue_graph_view( - node_table, - spatial_table, - sample_column="sample_id", - sample_value="missing", - id_column="instance_id", - ) - - -def test_build_tissue_graph_figure_returns_expected_traces_and_title() -> None: - node_table, spatial_table = _sample_tables() - view = prepare_tissue_graph_view( - node_table, - spatial_table, - sample_column="sample_id", - sample_value="sx", - id_column="instance_id", - ) - with_edges = build_tissue_graph_figure(view, show_edges=True, reverse_y=True) - assert isinstance(with_edges, go.Figure) - assert len(with_edges.data) == 2 - assert "sample_id=sx" in str(with_edges.layout.title.text) - - without_edges = build_tissue_graph_figure(view, show_edges=False) - assert isinstance(without_edges, go.Figure) - assert len(without_edges.data) == 1 diff --git a/tests/interpretability/test_transfer.py b/tests/interpretability/test_transfer.py deleted file mode 100644 index 7466bff..0000000 --- a/tests/interpretability/test_transfer.py +++ /dev/null @@ -1,210 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import pandas as pd -import pytest -from sklearn.neighbors import KNeighborsClassifier - -from grass_mil.contracts import ( - ClusteringResult, - InterpretabilityDataset, - ReportBundle, -) -from grass_mil.interpretability.core.reduction import run_reduction -from grass_mil.interpretability.core.transfer import ( - apply_niche_transfer, - fit_niche_transfer_from_report_bundle, - load_niche_transfer_bundle, - save_niche_transfer_bundle, -) - - -def _sample_instance_table() -> pd.DataFrame: - return pd.DataFrame( - { - "instance_id": [f"i{i}" for i in range(10)], - "bag_id": ["b0"] * 5 + ["b1"] * 5, - "inst_emb_0": [0.0, 0.1, 0.2, 0.3, 0.4, 1.0, 1.1, 1.2, 1.3, 1.4], - "inst_emb_1": [0.0, 0.15, 0.25, 0.35, 0.45, 1.0, 1.15, 1.25, 1.35, 1.45], - } - ) - - -def _build_report_bundle( - *, - dataset: InterpretabilityDataset, - labels: np.ndarray, - artifacts_dir: Path, - cluster_on_pca: bool, - pca_components: int = 2, -) -> ReportBundle: - emb = dataset.instance_table[["inst_emb_0", "inst_emb_1"]].to_numpy(dtype=float) - pca_result = run_reduction("pca", emb, n_components=pca_components, random_state=11) - clustering = ClusteringResult( - method="hdbscan", - labels=np.asarray(labels), - params={}, - fitted_object=None, - ) - return ReportBundle( - dataset=dataset, - reduction=pca_result, - niche_feature_reduction=pca_result if cluster_on_pca else None, - clustering=clustering, - niche_summary=None, - plugin_results={}, - artifacts_dir=artifacts_dir, - metadata={"cluster_on_pca": cluster_on_pca}, - ) - - -def test_niche_transfer_fit_and_apply_matches_direct_pca_knn_reference(tmp_path: Path) -> None: - table = _sample_instance_table() - dataset = InterpretabilityDataset(instance_table=table, id_column="instance_id") - labels = np.array([0, 0, 0, 1, 1, 2, 2, 2, 2, 1], dtype=int) - report_bundle = _build_report_bundle( - dataset=dataset, - labels=labels, - artifacts_dir=tmp_path, - cluster_on_pca=True, - pca_components=2, - ) - - bundle = fit_niche_transfer_from_report_bundle(dataset, report_bundle, n_neighbors=3) - query = table[["instance_id", "inst_emb_1", "inst_emb_0"]].sample(frac=1.0, random_state=3) - out = apply_niche_transfer(bundle, query, id_column="instance_id") - - query_matrix = query[["inst_emb_0", "inst_emb_1"]].to_numpy(dtype=float) - ref_knn = KNeighborsClassifier(n_neighbors=3, weights="uniform", metric="minkowski") - ref_knn.fit(report_bundle.niche_feature_reduction.embedding, labels) # type: ignore[union-attr] - ref_features = report_bundle.niche_feature_reduction.fitted_object.transform( # type: ignore[union-attr] - query_matrix - ) - ref_labels = ref_knn.predict(ref_features) - ref_conf = np.max(ref_knn.predict_proba(ref_features), axis=1) - - assert out["transferred_cluster_label"].tolist() == ref_labels.tolist() - assert np.allclose(out["transfer_confidence"].to_numpy(dtype=float), ref_conf) - - -def test_niche_transfer_preserves_noise_label_minus_one(tmp_path: Path) -> None: - table = _sample_instance_table() - dataset = InterpretabilityDataset(instance_table=table, id_column="instance_id") - labels = np.array([-1, -1, 0, 0, 0, 1, 1, 1, 2, 2], dtype=int) - report_bundle = _build_report_bundle( - dataset=dataset, - labels=labels, - artifacts_dir=tmp_path, - cluster_on_pca=False, - ) - bundle = fit_niche_transfer_from_report_bundle(dataset, report_bundle, n_neighbors=1) - out = apply_niche_transfer(bundle, table[["instance_id", "inst_emb_0", "inst_emb_1"]]) - - predicted = out["transferred_cluster_label"].to_numpy(dtype=int) - assert -1 in set(predicted.tolist()) - assert np.array_equal(predicted, labels) - - -def test_niche_transfer_bundle_serialization_roundtrip(tmp_path: Path) -> None: - table = _sample_instance_table() - dataset = InterpretabilityDataset(instance_table=table, id_column="instance_id") - labels = np.array([0, 0, 0, 1, 1, 2, 2, 2, 2, 1], dtype=int) - report_bundle = _build_report_bundle( - dataset=dataset, - labels=labels, - artifacts_dir=tmp_path, - cluster_on_pca=True, - pca_components=2, - ) - bundle = fit_niche_transfer_from_report_bundle(dataset, report_bundle, n_neighbors=3) - - bundle_path = tmp_path / "cluster_transfer_bundle.joblib" - save_niche_transfer_bundle(bundle, bundle_path) - restored = load_niche_transfer_bundle(bundle_path) - - query = table[["instance_id", "inst_emb_0", "inst_emb_1"]].copy() - left = apply_niche_transfer(bundle, query) - right = apply_niche_transfer(restored, query) - assert left.equals(right) - - -def test_niche_transfer_query_column_validation_and_reordered_columns(tmp_path: Path) -> None: - table = _sample_instance_table() - dataset = InterpretabilityDataset(instance_table=table, id_column="instance_id") - labels = np.array([0, 0, 0, 1, 1, 2, 2, 2, 2, 1], dtype=int) - report_bundle = _build_report_bundle( - dataset=dataset, - labels=labels, - artifacts_dir=tmp_path, - cluster_on_pca=False, - ) - bundle = fit_niche_transfer_from_report_bundle(dataset, report_bundle, n_neighbors=3) - - ordered = table[["instance_id", "inst_emb_0", "inst_emb_1"]] - reordered = table[["instance_id", "inst_emb_1", "inst_emb_0"]] - out_ordered = apply_niche_transfer(bundle, ordered) - out_reordered = apply_niche_transfer(bundle, reordered) - assert out_ordered.equals(out_reordered) - - with pytest.raises(ValueError, match="missing embedding columns"): - apply_niche_transfer(bundle, table[["instance_id", "inst_emb_0"]]) - - -def test_niche_transfer_fit_fails_fast_for_missing_inputs(tmp_path: Path) -> None: - table = _sample_instance_table() - dataset = InterpretabilityDataset(instance_table=table, id_column="instance_id") - labels = np.array([0, 0, 0, 1, 1, 2, 2, 2, 2, 1], dtype=int) - valid = _build_report_bundle( - dataset=dataset, - labels=labels, - artifacts_dir=tmp_path, - cluster_on_pca=True, - pca_components=2, - ) - - missing_clustering = ReportBundle( - dataset=dataset, - reduction=valid.reduction, - niche_feature_reduction=valid.niche_feature_reduction, - clustering=None, - niche_summary=None, - plugin_results={}, - artifacts_dir=tmp_path, - metadata={"cluster_on_pca": True}, - ) - with pytest.raises(ValueError, match="missing clustering result"): - fit_niche_transfer_from_report_bundle(dataset, missing_clustering, n_neighbors=3) - - missing_cluster_pca = ReportBundle( - dataset=dataset, - reduction=valid.reduction, - niche_feature_reduction=None, - clustering=valid.clustering, - niche_summary=None, - plugin_results={}, - artifacts_dir=tmp_path, - metadata={"cluster_on_pca": True}, - ) - with pytest.raises(ValueError, match="missing niche_feature_reduction"): - fit_niche_transfer_from_report_bundle(dataset, missing_cluster_pca, n_neighbors=3) - - with pytest.raises(ValueError, match="n_neighbors cannot exceed"): - fit_niche_transfer_from_report_bundle(dataset, valid, n_neighbors=len(table) + 1) - - -def test_niche_transfer_apply_fails_for_empty_input(tmp_path: Path) -> None: - table = _sample_instance_table() - dataset = InterpretabilityDataset(instance_table=table, id_column="instance_id") - labels = np.array([0, 0, 0, 1, 1, 2, 2, 2, 2, 1], dtype=int) - report_bundle = _build_report_bundle( - dataset=dataset, - labels=labels, - artifacts_dir=tmp_path, - cluster_on_pca=False, - ) - bundle = fit_niche_transfer_from_report_bundle(dataset, report_bundle, n_neighbors=3) - empty = pd.DataFrame(columns=["instance_id", "inst_emb_0", "inst_emb_1"]) - with pytest.raises(ValueError, match="empty instance table"): - apply_niche_transfer(bundle, empty) diff --git a/tests/models/test_attention.py b/tests/models/test_attention.py deleted file mode 100644 index e513d29..0000000 --- a/tests/models/test_attention.py +++ /dev/null @@ -1,68 +0,0 @@ -import pytest -import torch - - -def test_gated_attention_shape(): - from grass_mil.models.components.attention import AttnNetGated - - module = AttnNetGated(input_dim=32, hidden_dim=16, dropout=False, n_classes=1) - x = torch.randn(20, 32) - logits, out_x = module(x) - assert logits.shape == (20, 1) - assert out_x.shape == x.shape - - -def test_gated_projected_attention_shape(): - from grass_mil.models.components.attention import AttnNetGatedProjected - - module = AttnNetGatedProjected( - input_dim=32, projection_dim=16, hidden_dim=8, dropout=False, n_classes=1 - ) - x = torch.randn(20, 32) - logits, out_x = module(x) - assert logits.shape == (20, 1) - assert out_x.shape == x.shape - - -@pytest.mark.parametrize("enabled", [True, False]) -def test_attention_optional_factory(enabled): - from grass_mil.models.components.factory import build_attention - - att = build_attention(enabled, attention_type="gated", input_dim=16, hidden_dim=8) - if enabled: - assert att is not None - else: - assert att is None - - -def test_attention_head_spaces_expose_the_pre_tanh_activation(): - """The interpretability suite clusters the pre-tanh attention space. - - Reaching into module internals from the analysis code is fragile, so the - head exposes it. Pre-tanh matters because tanh saturation collapses the - extremes that distinguish strongly-attended instances. - """ - import torch - - from grass_mil.models.components.attention import AttnNetGated, AttnNetGatedProjected - - x = torch.randn(7, 12) - - gated = AttnNetGated(input_dim=12, hidden_dim=5, n_classes=2).eval() - spaces = gated.head_spaces(x) - assert set(spaces) == {"a_pre_tanh", "gate_product"} - assert spaces["a_pre_tanh"].shape == (7, 5) - # It really is pre-saturation: tanh of it reproduces the gate branch. - torch.testing.assert_close(torch.tanh(spaces["a_pre_tanh"]), gated.attention_a(x)) - - projected = AttnNetGatedProjected( - input_dim=12, projection_dim=6, hidden_dim=4, n_classes=2 - ).eval() - proj_spaces = projected.head_spaces(x) - assert proj_spaces["a_pre_tanh"].shape == (7, 4) - hidden = projected.projection(x) - torch.testing.assert_close( - torch.tanh(proj_spaces["a_pre_tanh"]), projected.attention_a(hidden) - ) - # The pre-tanh space is unbounded, unlike the saturated gate output. - assert proj_spaces["gate_product"].abs().max() <= 1.0 diff --git a/tests/models/test_augmentations.py b/tests/models/test_augmentations.py deleted file mode 100644 index 3ef132c..0000000 --- a/tests/models/test_augmentations.py +++ /dev/null @@ -1,243 +0,0 @@ -"""Contracts for the BGRL view generators.""" - -from __future__ import annotations - -import pytest -import torch - -from grass_mil.models.training.augmentations import ( - build_augmentation, - degree_importance, - drop_edges, - drop_importance, - feature_noise, - resolve_feature_indices, -) - - -def _star_graph(n_leaves: int = 8): - """A hub with many neighbours and leaves with exactly one.""" - from torch_geometric.data import Data - - leaves = torch.arange(1, n_leaves + 1, dtype=torch.long) - hub = torch.zeros_like(leaves) - edge_index = torch.cat([torch.stack([hub, leaves]), torch.stack([leaves, hub])], dim=1) - n_nodes = n_leaves + 1 - data = Data( - x=torch.ones(n_nodes, 3), - edge_index=edge_index, - edge_attr=torch.ones(edge_index.size(1), 2), - ) - data.num_nodes = n_nodes - data.categorical_codes = torch.zeros((n_nodes, 1), dtype=torch.long) - return data - - -def _generator(seed: int = 0) -> torch.Generator: - generator = torch.Generator() - generator.manual_seed(seed) - return generator - - -def test_degree_importance_ranks_hubs_above_leaves() -> None: - data = _star_graph(n_leaves=8) - importance = degree_importance(data.edge_index, int(data.num_nodes)) - assert importance[0] > importance[1:].max() - # Mean-max normalisation clamps every below-average node to zero. - assert float(importance[1:].max()) == pytest.approx(0.0) - assert float(importance.max()) == pytest.approx(1.0) - - -def test_drop_importance_preserves_high_degree_nodes() -> None: - """Low-degree nodes are dropped far more often than the hub.""" - data = _star_graph(n_leaves=12) - generator = _generator(0) - - hub_dropped = 0 - leaf_dropped = 0 - trials = 40 - for _ in range(trials): - out = drop_importance(data, mu=1.0, p_lambda=1.0, unassigned_index=7, generator=generator) - dropped = out.categorical_codes.view(-1) == 7 - hub_dropped += int(dropped[0]) - leaf_dropped += int(dropped[1:].sum()) - - assert hub_dropped == 0, "the hub has maximal importance and must never drop" - assert leaf_dropped > 0 - - -def test_drop_importance_marks_dropped_nodes_unassigned_and_zeroes_features() -> None: - data = _star_graph(n_leaves=10) - out = drop_importance(data, mu=1.0, p_lambda=1.0, unassigned_index=99, generator=_generator(1)) - dropped = out.categorical_codes.view(-1) == 99 - assert bool(dropped.any()) - assert torch.all(out.x[dropped] == 0) - assert torch.all(out.x[~dropped] == 1) - - -def test_drop_importance_does_not_mutate_the_source_graph() -> None: - """Guards the view-vs-clone class of bug. - - The legacy implementation wrote through a view, so a fixup applied to the - augmented graph silently reverted on the original. - """ - data = _star_graph(n_leaves=10) - x_before = data.x.clone() - codes_before = data.categorical_codes.clone() - edges_before = data.edge_index.clone() - - drop_importance(data, mu=1.0, p_lambda=1.0, unassigned_index=5, generator=_generator(2)) - - torch.testing.assert_close(data.x, x_before) - torch.testing.assert_close(data.categorical_codes, codes_before) - torch.testing.assert_close(data.edge_index, edges_before) - - -def test_drop_importance_keeps_the_graph_undirected() -> None: - """Dropping one direction would make message passing direction-dependent.""" - from torch_geometric.utils import is_undirected - - data = _star_graph(n_leaves=16) - for seed in range(5): - out = drop_importance( - data, mu=1.0, p_lambda=1.0, unassigned_index=3, generator=_generator(seed) - ) - assert is_undirected(out.edge_index), f"asymmetric edges at seed {seed}" - assert out.edge_attr.shape[0] == out.edge_index.shape[1] - - -def test_drop_edges_forces_undirected_by_default() -> None: - from torch_geometric.utils import is_undirected - - data = _star_graph(n_leaves=16) - out = drop_edges(data, p=0.5, generator=_generator(0)) - assert is_undirected(out.edge_index) - assert out.edge_attr.shape[0] == out.edge_index.shape[1] - - -def test_augmentations_are_deterministic_under_a_fixed_generator() -> None: - data = _star_graph(n_leaves=12) - first = drop_importance( - data, mu=0.5, p_lambda=0.5, unassigned_index=4, generator=_generator(7) - ) - second = drop_importance( - data, mu=0.5, p_lambda=0.5, unassigned_index=4, generator=_generator(7) - ) - torch.testing.assert_close(first.x, second.x) - torch.testing.assert_close(first.edge_index, second.edge_index) - - -def test_feature_noise_touches_only_the_selected_columns() -> None: - data = _star_graph(n_leaves=4) - out = feature_noise(data, std=1.0, feature_indices=[1], generator=_generator(0)) - - torch.testing.assert_close(out.x[:, 0], data.x[:, 0]) - torch.testing.assert_close(out.x[:, 2], data.x[:, 2]) - assert not torch.allclose(out.x[:, 1], data.x[:, 1]) - - -def test_feature_noise_rejects_out_of_range_columns() -> None: - data = _star_graph(n_leaves=4) - with pytest.raises(ValueError, match="out of range"): - feature_noise(data, std=1.0, feature_indices=[9], generator=_generator(0)) - - -def test_feature_columns_resolve_by_name() -> None: - names = ["cell_type", "SIZE", "other"] - assert resolve_feature_indices(["SIZE"], feature_names=names) == [1] - # Integer indices remain a usable fallback. - assert resolve_feature_indices([2], feature_names=names) == [2] - with pytest.raises(ValueError, match="not found"): - resolve_feature_indices(["MISSING"], feature_names=names) - - -def test_build_augmentation_dispatches_on_mode() -> None: - data = _star_graph(n_leaves=8) - - importance = build_augmentation( - {"mode": "importance", "mu": 0.5, "p_lambda": 0.5}, - unassigned_index=6, - ) - assert importance(data, _generator(0)) is not data - - uniform = build_augmentation({"mode": "uniform", "drop_edge_p": 0.5}) - assert uniform(data, _generator(0)) is not data - - with pytest.raises(ValueError, match="Unsupported augmentation mode"): - build_augmentation({"mode": "nonsense"}) - - -@pytest.mark.parametrize("mu,p_lambda", [(0.0, 0.5), (1.5, 0.5), (0.5, 0.0), (0.5, 1.5)]) -def test_drop_importance_validates_its_hyperparameters(mu: float, p_lambda: float) -> None: - with pytest.raises(ValueError): - drop_importance(_star_graph(), mu=mu, p_lambda=p_lambda) - - -def test_drop_edges_honours_its_generator() -> None: - """The generator must actually drive the draw. - - torch_geometric's dropout_edge takes no generator and consumes the global - RNG, so routing through it made seeded runs silently irreproducible. - """ - data = _star_graph(n_leaves=24) - - first = drop_edges(data, p=0.5, generator=_generator(11)) - second = drop_edges(data, p=0.5, generator=_generator(11)) - torch.testing.assert_close(first.edge_index, second.edge_index) - - different = drop_edges(data, p=0.5, generator=_generator(12)) - assert not ( - first.edge_index.shape == different.edge_index.shape - and torch.equal(first.edge_index, different.edge_index) - ) - - # The global RNG must not influence the result at all. - torch.manual_seed(0) - a = drop_edges(data, p=0.5, generator=_generator(11)) - torch.manual_seed(999) - b = drop_edges(data, p=0.5, generator=_generator(11)) - torch.testing.assert_close(a.edge_index, b.edge_index) - - -def test_bgrl_views_are_reproducible_and_distinct() -> None: - """Two views must differ from each other but repeat across runs.""" - import lightning as L - - from grass_mil.models.bgrl_module import BGRLModule - - def _views(seed: int): - module = BGRLModule( - encoder={"input_dim": 3, "hidden_dim": 8, "out_dim": 8, "num_layers": 2}, - ssl={ - "method": "bgrl", - "predictor": { - "_target_": "grass_mil.models.components.ssl.MLPPredictor", - "input_size": 8, - "output_size": 8, - "hidden_size": 16, - }, - }, - task={ - "drop_edge_p1": 0.4, - "drop_feat_p1": 0.4, - "drop_edge_p2": 0.4, - "drop_feat_p2": 0.4, - "augmentation_seed": seed, - # Set explicitly so setup() never reaches for a Trainer. - "total_steps": 10, - }, - ) - L.seed_everything(123, workers=True) - module.setup("fit") - return module._make_views(_star_graph(n_leaves=20)) - - first_a, first_b = _views(5) - second_a, second_b = _views(5) - - # Same augmentation seed reproduces both views exactly. - torch.testing.assert_close(first_a.edge_index, second_a.edge_index) - torch.testing.assert_close(first_b.edge_index, second_b.edge_index) - - # The two views of one batch are independent draws, not copies. - same_shape = first_a.edge_index.shape == first_b.edge_index.shape - assert not (same_shape and torch.equal(first_a.edge_index, first_b.edge_index)) diff --git a/tests/models/test_backbones.py b/tests/models/test_backbones.py deleted file mode 100644 index 886d446..0000000 --- a/tests/models/test_backbones.py +++ /dev/null @@ -1,80 +0,0 @@ -import pytest -import torch - - -def _toy_graph(num_nodes: int = 12): - from torch_geometric.data import Data - - row = torch.randint(0, num_nodes, (36,), dtype=torch.long) - col = torch.randint(0, num_nodes, (36,), dtype=torch.long) - edge_index = torch.stack([row, col], dim=0) - x = torch.randn(num_nodes, 16) - edge_attr = torch.rand(edge_index.size(1), 2) - batch = torch.zeros(num_nodes, dtype=torch.long) - return Data(x=x, edge_index=edge_index, edge_attr=edge_attr, batch=batch) - - -@pytest.mark.parametrize("conv_type", ["gin", "gcn", "gat", "graphsage", "gine"]) -def test_encoder_forward_node_and_graph(conv_type): - pytest.importorskip("torch_geometric") - from grass_mil.models.components.backbones import EncoderConfig, GNNEncoder - - data = _toy_graph() - cfg = EncoderConfig( - input_dim=16, - hidden_dim=32, - out_dim=32, - num_layers=3, - conv_type=conv_type, - pooling="mean", - use_edge_attr=True if conv_type in {"gcn", "gine"} else False, - edge_attr_dim=2 if conv_type == "gine" else None, - ) - encoder = GNNEncoder(cfg) - node_emb, graph_emb = encoder( - x=data.x, - edge_index=data.edge_index, - edge_attr=data.edge_attr, - batch=data.batch, - return_graph_embedding=True, - ) - assert node_emb.shape == (data.num_nodes, 32) - assert graph_emb.shape == (1, encoder.graph_pool.output_dim) - - -def test_encoder_warns_when_edge_attr_enabled_on_unsupported_backbone(): - pytest.importorskip("torch_geometric") - from grass_mil.models.components.backbones import EncoderConfig, GNNEncoder - - with pytest.warns(UserWarning, match="use_edge_attr=True is ignored"): - GNNEncoder( - EncoderConfig( - input_dim=16, - hidden_dim=32, - out_dim=32, - num_layers=2, - conv_type="gin", - use_edge_attr=True, - ) - ) - - -@pytest.mark.parametrize("jk", ["last", "concat", "max", "sum"]) -def test_encoder_jk_modes(jk): - pytest.importorskip("torch_geometric") - from grass_mil.models.components.backbones import EncoderConfig, GNNEncoder - - data = _toy_graph() - cfg = EncoderConfig( - input_dim=16, - hidden_dim=24, - out_dim=None, - num_layers=2, - conv_type="gin", - jk=jk, - pooling=None, - ) - encoder = GNNEncoder(cfg) - node_emb = encoder(data.x, data.edge_index, data.edge_attr) - expected_dim = 24 * 3 if jk == "concat" else 24 - assert node_emb.shape == (data.num_nodes, expected_dim) diff --git a/tests/models/test_bagging.py b/tests/models/test_bagging.py deleted file mode 100644 index 0f7934d..0000000 --- a/tests/models/test_bagging.py +++ /dev/null @@ -1,106 +0,0 @@ -"""Bag-pooling contracts, including the attribution identity. - -The published MIL head pools instance logits with per-class attention: - - A[i,c] = softmax_i(alpha[i,c]), L_c = sum_i A[i,c] * l[i,c] - -The exactness of that sum is what the whole attribution suite rests on, so it -is asserted directly here rather than inferred from downstream behaviour. -""" - -from __future__ import annotations - -import pytest -import torch -from torch import nn - -from grass_mil.contracts import validate_attention_width -from grass_mil.models.training.bagging import aggregate_bag_logits_attention - - -class _FixedAttention(nn.Module): - """Attention returning preset logits, so pooling is checked in isolation.""" - - def __init__(self, logits: torch.Tensor) -> None: - super().__init__() - self.register_buffer("logits", logits) - self.attention_c = nn.Linear(1, logits.shape[1]) - - def forward(self, x): - return self.logits[: x.shape[0]], x - - -def _pool(*, logits: torch.Tensor, attn_logits: torch.Tensor): - return aggregate_bag_logits_attention( - logits=logits, - embeddings=torch.zeros((logits.shape[0], 1)), - attention=_FixedAttention(attn_logits), - bag_groups={"bag": list(range(logits.shape[0]))}, - max_instances_per_bag=0, - instance_sampling="all", - ) - - -def test_bag_logit_is_the_attention_weighted_sum_of_instance_logits() -> None: - """The identity sum_i A[i,c] * l[i,c] == L_c, asserted exactly.""" - torch.manual_seed(0) - logits = torch.randn(6, 2) - attn_logits = torch.randn(6, 2) - - bag_logits, bag_ids, bag_attention, _ = _pool(logits=logits, attn_logits=attn_logits) - - assert bag_ids == ["bag"] - attention = bag_attention["bag"] - assert attention.shape == (6, 2) - - expected = (logits * attention).sum(dim=0, keepdim=True) - torch.testing.assert_close(bag_logits, expected, rtol=0, atol=1e-12) - - -def test_attention_is_normalised_per_class_over_instances() -> None: - torch.manual_seed(1) - _, _, bag_attention, _ = _pool(logits=torch.randn(5, 3), attn_logits=torch.randn(5, 3)) - - attention = bag_attention["bag"] - # Each class column is a distribution over instances; columns are independent. - torch.testing.assert_close(attention.sum(dim=0), torch.ones(3), rtol=0, atol=1e-6) - - -def test_per_class_attention_can_select_different_instances_per_class() -> None: - """The point of per-class attention: class 0 and class 1 may disagree.""" - logits = torch.tensor([[1.0, 100.0], [2.0, 200.0]]) - # Class 0 concentrates on instance 0, class 1 on instance 1. - attn_logits = torch.tensor([[20.0, 0.0], [0.0, 20.0]]) - - bag_logits, _, _, _ = _pool(logits=logits, attn_logits=attn_logits) - - assert bag_logits[0, 0] == pytest.approx(1.0, abs=1e-4) - assert bag_logits[0, 1] == pytest.approx(200.0, abs=1e-2) - - -def test_pooling_rejects_a_shared_attention_channel() -> None: - """One channel broadcast across classes cannot express opposing evidence. - - A niche that pushes toward one class and away from the other needs two - independent attention distributions, so a binary head requires two - channels rather than one. - """ - with pytest.raises(ValueError, match="Attention emits 1 channel"): - _pool(logits=torch.randn(4, 2), attn_logits=torch.randn(4, 1)) - - -def test_pooling_rejects_attention_width_that_does_not_match_the_head() -> None: - with pytest.raises(ValueError, match="Attention emits 3 channel"): - _pool(logits=torch.randn(4, 2), attn_logits=torch.randn(4, 3)) - - -def test_validate_attention_width_accepts_one_channel_per_class() -> None: - validate_attention_width(attention_width=4, num_classes=4) - # A scalar head (regression, Cox log-hazard) is one class, one channel. - validate_attention_width(attention_width=1, num_classes=1) - - -@pytest.mark.parametrize("width", [0, 1, 2, 5]) -def test_validate_attention_width_rejects_other_widths(width: int) -> None: - with pytest.raises(ValueError): - validate_attention_width(attention_width=width, num_classes=4) diff --git a/tests/models/test_fidelity_phase1.py b/tests/models/test_fidelity_phase1.py deleted file mode 100644 index 2d79975..0000000 --- a/tests/models/test_fidelity_phase1.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Phase 1 fidelity: graph edge features, encoder epsilon/init, attention init and space.""" - -import numpy as np -import torch -from torch import nn - -from grass_mil.data.components.graph_builders import DelaunayGraphBuilder, log_distance_feature -from grass_mil.models.components.attention import AttnNetGatedProjected -from grass_mil.models.components.backbones import EncoderConfig, GNNEncoder -from grass_mil.models.components.embeddings import ( - CategoricalEmbeddingConfig, - NodeInputEmbedding, - resolve_categorical_embedding_config, -) -from grass_mil.models.components.factory import build_loss - - -def test_log_distance_matches_reference_formula(): - d = np.array([0.0, 7.389, 20.0, 148.41, 1000.0]) - got = log_distance_feature(d, log_lower=2.0, log_upper=5.0) - expected = np.clip((np.log(d + 1e-5) - 2.0) / 3.0, 0.0, 1.0) - assert np.allclose(got, expected, atol=1e-6) - assert got[0] == 0.0 and got[-1] == 1.0 - - -def test_delaunay_builder_emits_log_distance_and_stores_edges_twice(): - rng = np.random.default_rng(0) - coords = rng.uniform(0, 500, size=(60, 2)) - builder = DelaunayGraphBuilder( - edge_features=["log_distance", "neighbor", "distance"], neighbor_cutoff_um=55.0 - ) - res = builder.build(coords) - assert res.edge_attr_names == ["log_distance", "neighbor", "distance"] - ei = res.edge_index.numpy() - pairs = set(map(tuple, ei.T.tolist())) - assert all((b, a) in pairs for a, b in pairs) - dist = res.edge_attr[:, 2].numpy() - assert np.allclose(res.edge_attr[:, 0].numpy(), log_distance_feature(dist), atol=1e-6) - assert np.array_equal(res.edge_attr[:, 1].numpy(), (dist <= 55.0).astype(np.float32)) - - -def _gine_cfg(**kw): - base = dict( - input_dim=3, - hidden_dim=16, - out_dim=16, - num_layers=2, - conv_type="gine", - norm="layernorm", - pooling="max", - edge_attr_dim=1, - edge_feature_index=0, - use_edge_attr=True, - ) - base.update(kw) - return EncoderConfig(**base) - - -def test_gine_eps_is_trainable_by_default_and_fixed_when_disabled(): - enc = GNNEncoder(_gine_cfg()) - assert all(isinstance(layer.eps, nn.Parameter) for layer in enc.layers) - enc_fixed = GNNEncoder(_gine_cfg(train_eps=False)) - assert all(not isinstance(layer.eps, nn.Parameter) for layer in enc_fixed.layers) - - -def test_input_embedding_xavier_init_changes_scale_and_default_keeps_torch(): - torch.manual_seed(0) - cat = CategoricalEmbeddingConfig(label="cell_type", num_embeddings=10, column_index=0) - xav = NodeInputEmbedding(input_dim=4, hidden_dim=64, categorical=cat, init="xavier") - torch.manual_seed(0) - dflt = NodeInputEmbedding(input_dim=4, hidden_dim=64, categorical=cat, init="default") - # nn.Embedding default is N(0, 1); xavier bound for (11, 64) is sqrt(6 / 75) ~ 0.28. - assert xav.embedding.weight.abs().max() < 0.5 - assert dflt.embedding.weight.abs().max() > 1.0 - - -def test_unassigned_label_resolves_to_existing_category(): - cfg = {"label": "cell_type", "reserve_unassigned": False, "unassigned_label": "other"} - resolved = resolve_categorical_embedding_config( - cfg, - cardinalities={"cell_type": 3}, - slices={"cell_type": 0}, - label_maps={"cell_type": {"tumor": 0, "other": 1, "immune": 2}}, - ) - assert resolved.unassigned_index == 1 - emb = NodeInputEmbedding(input_dim=0, hidden_dim=8, categorical=resolved) - assert emb.unassigned_index == 1 - assert emb.embedding.num_embeddings == 3 - - -def test_attention_kaiming_init_and_attention_space(): - torch.manual_seed(0) - att = AttnNetGatedProjected( - input_dim=32, projection_dim=8, hidden_dim=4, dropout=True, n_classes=2 - ) - att.eval() - x = torch.randn(20, 32) - space = att.attention_space(x) - assert space.shape == (20, 4) - manual = att.attention_a[0](att.projection[1](att.projection[0](x))) - assert torch.allclose(space, manual) - logits, passthrough = att(x) - assert logits.shape == (20, 2) and passthrough is x - # kaiming_uniform (a=0): bound = sqrt(2) * sqrt(3 / fan_in) = sqrt(6 / 32) ~ 0.433 - assert att.projection[0].weight.abs().max() <= np.sqrt(6.0 / 32) + 1e-6 - assert att.projection[0].weight.abs().max() > 1.0 / np.sqrt(32) # wider than torch default - - -def test_loss_weights_accept_lists(): - loss = build_loss("categorical_ce", class_weight=[1.0, 3.0]) - assert isinstance(loss.class_weight, torch.Tensor) - logits = torch.randn(5, 2) - target = torch.tensor([0, 1, 1, 0, 1]) - assert torch.isfinite(loss(logits, target)) diff --git a/tests/models/test_fidelity_phase2.py b/tests/models/test_fidelity_phase2.py deleted file mode 100644 index 5e4d165..0000000 --- a/tests/models/test_fidelity_phase2.py +++ /dev/null @@ -1,284 +0,0 @@ -"""Phase 2 fidelity: region bags, root sampling, parameter groups, schedules, class weights, -fine-tune augmentation, regime presets and region-level metrics.""" - -from pathlib import Path - -import pytest -import torch -from hydra import compose, initialize_config_module -from hydra.core.global_hydra import GlobalHydra -from omegaconf import OmegaConf, open_dict - -from grass_mil.data.components.samplers import ( - UNLIMITED_NEIGHBORS, - RuntimeShadowConfig, - ShadowCustomStrategy, -) -from grass_mil.models.training import ( - build_finetune_augmentation, - build_param_groups, - drop_features, - inverse_frequency_class_weights, - resolve_t_max_auto, -) - -CONFIG_MODULE = "grass_mil.configs" - - -def _compose(overrides: list[str]): - GlobalHydra.instance().clear() - with initialize_config_module(version_base="1.3", config_module=CONFIG_MODULE): - cfg = compose(config_name="train", overrides=overrides) - GlobalHydra.instance().clear() - return cfg - - -def _units(n_units: int = 3, n_nodes: int = 40): - from torch_geometric.data import Data - - out = [] - g = torch.Generator().manual_seed(0) - for i in range(n_units): - edge_index = torch.randint(0, n_nodes, (2, 160), generator=g) - data = Data( - x=torch.randn(n_nodes, 4, generator=g), - edge_index=edge_index, - edge_attr=torch.rand(160, 2, generator=g), - graph_y=torch.tensor([[float(i % 2)]]), - graph_w=torch.ones(1, 1), - num_nodes=n_nodes, - ) - data.region_id = f"region_{i}" - data.sample_id = f"sample_{i}" - data.patch_id = f"patch_{i}" - out.append(data) - return out - - -def _runtime(**kw): - base = dict( - enabled=True, - depth=2, - num_neighbors=-1, - subgraph_batch_size=8, - proportional_root_sampling=False, - ) - base.update(kw) - return base - - -def test_num_neighbors_minus_one_means_unlimited(): - cfg = RuntimeShadowConfig.from_dict(_runtime()) - assert cfg.num_neighbors == UNLIMITED_NEIGHBORS - - -def test_unit_as_bag_yields_one_batch_per_unit_with_subsample_count_roots(): - strategy = ShadowCustomStrategy(runtime=_runtime(subsample_count=12, unit_as_bag=True)) - loader = strategy.build_dataset_loader(dataset=_units(), batch_size=1, shuffle=False) - batches = list(loader) - assert len(batches) == 3 == len(loader) - for batch in batches: - assert int(batch.ptr.numel()) - 1 == 12 - assert len(set(batch.region_id)) == 1 - - -def test_fixed_per_unit_seed_gives_identical_roots_every_pass(): - strategy = ShadowCustomStrategy( - runtime=_runtime(subsample_count=10, unit_as_bag=True, fixed_per_unit_seed=7) - ) - loader = strategy.build_dataset_loader(dataset=_units(), batch_size=1, shuffle=False) - first = [b.root_n_id.sort().values.tolist() for b in loader] - second = [b.root_n_id.sort().values.tolist() for b in loader] - assert first == second - unseeded = ShadowCustomStrategy(runtime=_runtime(subsample_count=10, unit_as_bag=True)) - loader2 = unseeded.build_dataset_loader(dataset=_units(), batch_size=1, shuffle=False) - a = [b.root_n_id.sort().values.tolist() for b in loader2] - b = [b.root_n_id.sort().values.tolist() for b in loader2] - assert a != b - - -def test_all_roots_mode_covers_every_node_once(): - strategy = ShadowCustomStrategy(runtime=_runtime(subsample_count=None, subgraph_batch_size=16)) - loader = strategy.build_dataset_loader( - dataset=_units(n_units=1, n_nodes=40), batch_size=1, shuffle=False - ) - roots = torch.cat([b.root_n_id for b in loader]).sort().values.tolist() - assert roots == list(range(40)) - - -def test_build_param_groups_scales_and_freezes(): - enc = torch.nn.Linear(4, 4) - head = torch.nn.Linear(4, 2) - att = torch.nn.Linear(4, 2) - groups = build_param_groups( - base_lr=1e-4, - base_weight_decay=1e-3, - encoder=enc, - head=head, - attention=att, - encoder_lr_scale=0.0, - head_lr_scale=0.01, - attention_lr_scale=2.0, - head_weight_decay=5e-4, - ) - names = [g["name"] for g in groups] - assert names == ["head", "attention"] - assert groups[0]["lr"] == pytest.approx(1e-6) and groups[0]["weight_decay"] == pytest.approx( - 5e-4 - ) - assert groups[1]["lr"] == pytest.approx(2e-4) and groups[1]["weight_decay"] == pytest.approx( - 1e-3 - ) - with pytest.raises(ValueError): - build_param_groups( - base_lr=1e-4, - base_weight_decay=0.0, - encoder=enc, - head=head, - attention=att, - encoder_lr_scale=0.0, - head_lr_scale=0.0, - attention_lr_scale=0.0, - ) - - -def test_t_max_auto_counts_optimizer_steps(): - assert resolve_t_max_auto(num_train_units=340, max_epochs=25, hyperbatch_size=12) == 25 * 29 - assert resolve_t_max_auto(num_train_units=5, max_epochs=3, hyperbatch_size=12) == 3 - - -def test_inverse_frequency_class_weights_have_mean_one(): - w = inverse_frequency_class_weights([0] * 70 + [1] * 30, 2) - assert w.mean().item() == pytest.approx(1.0) - assert w[1] > w[0] - assert (w[1] / w[0]).item() == pytest.approx(70 / 30) - - -def test_elementwise_feature_drop_writes_unassigned_token(): - from torch_geometric.data import Data - - data = Data( - x=torch.ones(200, 3), edge_index=torch.empty(2, 0, dtype=torch.long), num_nodes=200 - ) - data.categorical_codes = torch.zeros(200, 1, dtype=torch.long) - gen = torch.Generator().manual_seed(0) - out = drop_features( - data, p=0.3, generator=gen, mode="elementwise", categorical_column=0, unassigned_index=9 - ) - zero_frac = float((out.x == 0).float().mean()) - assert 0.2 < zero_frac < 0.4 - assert 0.2 < float((out.categorical_codes[:, 0] == 9).float().mean()) < 0.4 - assert torch.equal(data.x, torch.ones(200, 3)) # input untouched - assert build_finetune_augmentation({"drop_edge_p": 0.0, "drop_feat_p": 0.0}) is None - aug = build_finetune_augmentation({"drop_edge_p": 0.1, "drop_feat_p": 0.1}, unassigned_index=9) - assert aug(data).x.shape == data.x.shape - - -@pytest.mark.parametrize( - ("regime", "extra"), - [ - ("mil", []), - ("mil_frozen", []), - ("linear_probehead", ["task.head_init_path=/tmp/none.pt"]), - ("linear_ftall", ["task.head_init_path=/tmp/none.pt"]), - ("linear_ftall_reg", ["task.head_init_path=/tmp/none.pt"]), - ], -) -def test_regime_presets_compose(regime: str, extra: list[str]): - cfg = _compose([f"regime={regime}", *extra]) - assert cfg.regime_name == regime - assert cfg.model.encoder.conv_type == "gine" and int(cfg.model.encoder.num_layers) == 4 - assert cfg.data.sampler.runtime.unit_as_bag is True - assert int(cfg.data.val_sampler.runtime.subsample_count) == 1024 - assert str(cfg.scheduler.T_max) == "auto" and cfg.task.scheduler_interval == "step" - if regime.startswith("linear"): - assert int(cfg.model.graph_head.num_layers) == 1 - else: - assert int(cfg.model.graph_head.num_layers) == 2 - if regime == "linear_ftall_reg": - assert float(cfg.task.optimization.encoder_lr_scale) == pytest.approx(0.3) - assert float(cfg.task.finetune_augmentation.drop_edge_p) == pytest.approx(0.05) - assert float(cfg.model.encoder.dropout) == pytest.approx(0.2) - if regime == "linear_probehead": - assert ( - cfg.task.freeze_head is True and float(cfg.task.optimization.encoder_lr_scale) == 0.0 - ) - - -def test_regime_defaults_can_be_overridden(): - cfg = _compose( - [ - "regime=mil", - "model.encoder.num_layers=3", - "data.sampler.runtime.depth=3", - "model.graph_head.num_layers=1", - "data.sampler.runtime.subsample_count=64", - ] - ) - assert int(cfg.model.encoder.num_layers) == 3 and int(cfg.data.sampler.runtime.depth) == 3 - assert ( - int(cfg.model.graph_head.num_layers) == 1 - and int(cfg.data.sampler.runtime.subsample_count) == 64 - ) - - -def test_region_level_metrics_binary_and_categorical(): - from grass_mil.models.supervised_module import SupervisedModule - - logits = torch.tensor([[2.0, -1.0], [-1.0, 2.0], [1.5, 0.0], [-2.0, 1.0]]) - targets = torch.tensor([0, 1, 0, 1]) - m = SupervisedModule._region_level_metrics("categorical", logits, targets) - assert m["auroc"] == pytest.approx(1.0) and m["bal_acc"] == pytest.approx(1.0) - assert m["loss_region"] > 0 - mb = SupervisedModule._region_level_metrics( - "binary", torch.tensor([[2.0], [-2.0]]), torch.tensor([1.0, 0.0]) - ) - assert mb["auroc"] == pytest.approx(1.0) - - -def test_mil_training_with_scaled_groups_and_auto_weights(cfg_train, tmp_path): - """End-to-end: region accumulation, scaled parameter groups, auto class weights, - per-step cosine with T_max auto, region-level val metrics and test predictions.""" - from grass_mil.train import train - - cfg = _compose(["task=finetune_mil", "model=supervised_module", "scheduler=cosine_mil_step"]) - with open_dict(cfg): - cfg.data = OmegaConf.create( - { - "_target_": "tests.helpers.synthetic_datamodule.SyntheticBagDataModule", - "batch_size": 4, - "num_workers": 0, - "pin_memory": False, - "input_dim": 8, - } - ) - cfg.model.encoder.input_dim = 8 - cfg.task.optimization.encoder_lr_scale = 0.3 - cfg.task.optimization.head_lr_scale = 0.01 - cfg.task.optimization.head_weight_decay = 1.0e-3 - cfg.task.region_accumulation.hyperbatch_size = 2 - cfg.task.grad_clip_norm = 1.0 - cfg.callbacks = None - cfg.logger = None - cfg.paths.output_dir = str(tmp_path) - cfg.paths.root_dir = str(tmp_path) - cfg.trainer.max_epochs = 1 - cfg.trainer.min_epochs = 1 - cfg.trainer.accelerator = "cpu" - cfg.trainer.devices = 1 - cfg.trainer.default_root_dir = str(tmp_path) - cfg.extras.print_config = False - cfg.extras.enforce_tags = False - cfg.train = True - cfg.test = True - metrics, objects = train(cfg) - module = objects["model"] - assert module.loss_fn.class_weight is not None - assert module.loss_fn.class_weight.mean().item() == pytest.approx(1.0) - opt = objects["trainer"].optimizers[0] - group_lrs = sorted(g["initial_lr"] if "initial_lr" in g else g["lr"] for g in opt.param_groups) - assert len(opt.param_groups) == 3 - assert "val/auroc" in metrics and "test/auroc" in metrics - assert (Path(tmp_path) / "metrics_per_epoch.jsonl").exists() - assert (Path(tmp_path) / "test_predictions.csv").exists() - assert group_lrs[0] < group_lrs[-1] diff --git a/tests/models/test_fidelity_phase3.py b/tests/models/test_fidelity_phase3.py deleted file mode 100644 index 0f9b543..0000000 --- a/tests/models/test_fidelity_phase3.py +++ /dev/null @@ -1,295 +0,0 @@ -"""Phase 3: explicit splits, probe pre-fit, nested cross-validation.""" - -from __future__ import annotations - -import json -from pathlib import Path - -import numpy as np -import pandas as pd -import pytest -import torch - -from tests.helpers.synthetic_cohort import cohort_overrides, compose_cfg, make_synthetic_cohort - -pytest.importorskip("torch_geometric") -pytest.importorskip("sklearn") - - -# ---------------------------------------------------------------------------- -# splits -# ---------------------------------------------------------------------------- - - -def test_explicit_split_from_lists_and_files(tmp_path): - from grass_mil.data.components.splits import ExplicitSplit - - (tmp_path / "tr.json").write_text(json.dumps(["a", "b"])) - (tmp_path / "va.txt").write_text("c\n\nd\n") - split = ExplicitSplit.from_dict( - {"train": str(tmp_path / "tr.json"), "val": str(tmp_path / "va.txt"), "test": ["e"]} - ) - assert split.train == ["a", "b"] and split.val == ["c", "d"] and split.test == ["e"] - assert split.lookup() == {"a": "train", "b": "train", "c": "val", "d": "val", "e": "test"} - assert ExplicitSplit.from_dict({"train": None, "val": None, "test": None}) is None - with pytest.raises(ValueError): - ExplicitSplit.from_dict({"train": ["a"], "val": ["a"], "test": ["b"]}) - # val == test is the refit convention and is tolerated - ExplicitSplit.from_dict({"train": ["a"], "val": ["b"], "test": ["b"]}) - - -def test_fold_json_roundtrip_and_duplicates(tmp_path): - from grass_mil.data.components.splits import load_fold_json, write_fold_json - - write_fold_json([["r1", "r2"], ["r3"]], tmp_path / "f.json") - assert load_fold_json(tmp_path / "f.json") == [["r1", "r2"], ["r3"]] - (tmp_path / "bad.json").write_text(json.dumps([["r1"], ["r1"]])) - with pytest.raises(ValueError): - load_fold_json(tmp_path / "bad.json") - - -def test_inner_folds_are_patient_safe_and_deterministic(tmp_path): - from grass_mil.data.components.splits import load_labels, make_inner_folds - - rows = [] - for p in range(9): - for r in range(2): - rows.append({"REGION_ID": f"P{p}_{r}", "STAS": p % 2, "PATIENT": f"P{p}"}) - pd.DataFrame(rows).to_csv(tmp_path / "labels.csv", index=False) - table = load_labels( - tmp_path / "labels.csv", - id_column="REGION_ID", - label_column="STAS", - patient_column="PATIENT", - ) - ids = [r["REGION_ID"] for r in rows] - folds_a = make_inner_folds(ids, table, k=3, seed=0) - folds_b = make_inner_folds(list(reversed(ids)), table, k=3, seed=0) - assert folds_a == folds_b # order of the id set does not matter - for tr, va in folds_a: - assert not (set(tr) & set(va)) - assert not ({table.group_of(r) for r in tr} & {table.group_of(r) for r in va}) - assert len({table.label_of(r) for r in tr}) == 2 - - -# ---------------------------------------------------------------------------- -# probe folding -# ---------------------------------------------------------------------------- - - -def test_probe_folding_reproduces_decision_function(): - from sklearn.linear_model import LogisticRegression - from sklearn.preprocessing import StandardScaler - - from grass_mil.models.training.probe import fold_probe_into_linear_head - - rng = np.random.default_rng(0) - X = rng.normal(size=(200, 12)) * rng.uniform(0.5, 3.0, size=12) + rng.normal(size=12) - y = (X[:, 0] - X[:, 3] > 0).astype(int) - sc = StandardScaler().fit(X) - clf = LogisticRegression(C=1.0, class_weight="balanced", max_iter=3000).fit(sc.transform(X), y) - W, B = fold_probe_into_linear_head( - clf.coef_, clf.intercept_, sc.mean_, sc.scale_, num_classes=2 - ) - assert torch.all(W[0] == 0) and B[0] == 0 - logits = torch.from_numpy(X).float() @ W.T + B - ref = clf.decision_function(sc.transform(X)) - assert np.allclose((logits[:, 1] - logits[:, 0]).numpy(), ref, atol=1e-4) - - -# ---------------------------------------------------------------------------- -# cohort fixture: precompute once per module -# ---------------------------------------------------------------------------- - - -@pytest.fixture(scope="module") -def cohort(tmp_path_factory): - root = tmp_path_factory.mktemp("cohort") - c = make_synthetic_cohort(root, n_patients=8, regions_per_patient=2, n_cells=120) - cfg = compose_cfg("train", cohort_overrides(c, root / "out_precompute")) - import hydra - - dm = hydra.utils.instantiate(cfg.data) - dm.prepare_data() - return c - - -def test_precompute_records_patient_and_label(cohort): - index = json.loads((cohort.processed_dir / "processed_index.json").read_text()) - entries = index["entries"] - assert len(entries) == len(cohort.region_ids) - for e in entries: - assert e["patient_id"] == cohort.patient_of[e["region_id"]] - assert int(round(e["graph_y"][0])) == cohort.label_of[e["region_id"]] - data = torch.load(entries[0]["path"], weights_only=False) - assert data.patient_id == cohort.patient_of[entries[0]["region_id"]] - assert list(data.graph_label_names) == ["STAS"] - - -def test_explicit_split_repartitions_without_reprocessing(cohort, tmp_path): - import hydra - - ids = cohort.region_ids - cfg = compose_cfg( - "train", - cohort_overrides(cohort, tmp_path) - + [ - f"data.split.explicit.train=[{','.join(ids[:10])}]", - f"data.split.explicit.val=[{','.join(ids[10:13])}]", - f"data.split.explicit.test=[{','.join(ids[13:15])}]", - ], - ) - dm = hydra.utils.instantiate(cfg.data) - dm.setup("fit") - assert sorted(e.region_id for e in dm.dataset_train.entries) == sorted(ids[:10]) - assert sorted(e.region_id for e in dm.dataset_val.entries) == sorted(ids[10:13]) - assert sorted(e.region_id for e in dm.dataset_test.entries) == sorted(ids[13:15]) - cfg2 = compose_cfg( - "train", - cohort_overrides(cohort, tmp_path) + ["data.split.explicit.train=[nope]"], - ) - dm2 = hydra.utils.instantiate(cfg2.data) - with pytest.raises(ValueError, match="no processed graph"): - dm2.setup("fit") - - -def test_probe_cli_writes_folded_head(cohort, tmp_path): - from grass_mil.probe import run_probe - - ids = cohort.region_ids - assert len(ids) == 16 - cfg = compose_cfg( - "probe", - cohort_overrides(cohort, tmp_path) - + [ - "regime=linear_probehead", - "task.head_init_path=null", - f"data.split.explicit.train=[{','.join(ids[:10])}]", - f"data.split.explicit.val=[{','.join(ids[10:])}]", - f"data.split.explicit.test=[{','.join(ids[10:])}]", - "probe.roots_per_unit=16", - f"probe.output_path={tmp_path}/probe_head.pt", - ], - ) - summary = run_probe(cfg) - head = torch.load(tmp_path / "probe_head.pt", weights_only=False) - W = head["graph_head.net.0.weight"] - assert W.shape == (2, 16) and torch.all(W[0] == 0) - assert "val_region_auroc" in summary["metrics"] - assert summary["settings"]["roots_per_unit"] == 16 - assert summary["settings"]["n_train_regions"] == 10 - assert (tmp_path / "probe_head.json").exists() - - -# ---------------------------------------------------------------------------- -# nested CV -# ---------------------------------------------------------------------------- - - -def test_best_epoch_from_jsonl(tmp_path): - from grass_mil.nested_cv import best_epoch_from_jsonl - - p = tmp_path / "m.jsonl" - p.write_text( - "\n".join( - json.dumps(r) - for r in [ - {"epoch": 0, "val/auroc": 0.5}, - {"epoch": 1, "val/auroc": 0.7}, - {"epoch": 2, "val/auroc": 0.7}, - {"epoch": 3, "val/auroc": 0.6}, - ] - ) - ) - assert best_epoch_from_jsonl(p, metric="val/auroc") == ( - 2, - "val/auroc", - ) # earliest tie, 1-based - assert best_epoch_from_jsonl(p, metric="val/loss", mode="min") == (4, "last") - assert best_epoch_from_jsonl(p, metric="val/x", fallback_metric="val/auroc") == ( - 1, - "val/auroc", - ) - - -@pytest.mark.parametrize("regime", ["mil", "linear_probehead"]) -def test_nested_cv_smoke(cohort, tmp_path, regime): - from grass_mil.nested_cv import run_nested_cv - - out = tmp_path / regime - cfg = compose_cfg( - "nested_cv", - cohort_overrides(cohort, out) - + [ - f"regime={regime}", - f"nested_cv.folds_file={cohort.folds_path}", - "nested_cv.inner_k=2", - "nested_cv.outer_folds=[0,1]", - "nested_cv.probe.roots_per_unit=16", - "trainer.max_epochs=2", - ], - ) - summary = run_nested_cv(cfg) - assert summary["n_outer"] == 2 and not summary["complete"] - assert len(summary["auc_folds"]) == 2 - assert all(np.isfinite(a) for a in summary["auc_folds"]) - assert summary["probe"] == (regime == "linear_probehead") - oof = pd.read_csv(out / "oof_predictions.csv") - expected = sorted(cohort.folds[0] + cohort.folds[1]) - assert sorted(oof["region_id"]) == expected # one OOF prediction per held-out region - assert set(oof["patient_id"]) == {cohort.patient_of[r] for r in expected} - pat = pd.read_csv(out / "oof_patient_predictions.csv") - assert len(pat) == len({cohort.patient_of[r] for r in expected}) - for i in (0, 1): - res = json.loads((out / f"outer_{i}" / "result.json").read_text()) - assert 1 <= res["E"] <= 2 and len(res["inner_epochs"]) == 2 - assert (out / f"outer_{i}" / "refit" / "test_predictions.csv").exists() - if regime == "linear_probehead": - assert res["head_init_path"] and Path(res["head_init_path"]).exists() - # resume: a second call reuses every finished run - summary2 = run_nested_cv(cfg) - assert summary2["auc_folds"] == summary["auc_folds"] - - -def test_explicit_split_patch_id_kind_and_unit_key(): - from grass_mil.data.components.datasets import ProcessedIndexEntry - from grass_mil.data.components.splits import ExplicitSplit - from grass_mil.data.spatial_omics_datamodule import unit_key - - e = ProcessedIndexEntry(path="x", sample_id="s1", region_id="r1", patch_id="s1_t3") - assert ( - unit_key(e, "region") == "r1" - and unit_key(e, "patch") == "s1_t3" - and unit_key(e, "sample") == "s1" - ) - split = ExplicitSplit.from_dict( - {"train": ["s1_t3"], "val": None, "test": None, "id_kind": "patch"} - ) - assert split.id_kind == "patch" - with pytest.raises(ValueError): - ExplicitSplit.from_dict({"train": ["a"], "id_kind": "tile"}) - - -def test_patch_bags_log_region_level_metrics(cohort, tmp_path): - """Tiles as bags (task.bag_key=patch_id): validation AUROC is computed per tile.""" - from grass_mil.train import train - - cfg = compose_cfg( - "train", - cohort_overrides(cohort, tmp_path) - + [ - f"data.processed_dir={tmp_path / 'processed_tiles'}", - "data.sample_unit=tile", - "data.tiling.tile_size_um=300", - "data.tiling.stride_um=300", - "data.tiling.min_cells=5", - "data.split.split_by=region", - "task.bag_key=patch_id", - "trainer.max_epochs=1", - ], - ) - metrics, objects = train(cfg) - assert "val/auroc" in metrics and "test/auroc" in metrics - preds = pd.read_csv(tmp_path / "test_predictions.csv") - assert preds["bag_id"].str.contains("patch").all() - assert preds["bag_id"].nunique() == len(preds) > len(set(preds["region_id"])) diff --git a/tests/models/test_fidelity_phase4.py b/tests/models/test_fidelity_phase4.py deleted file mode 100644 index 3bfc6b2..0000000 --- a/tests/models/test_fidelity_phase4.py +++ /dev/null @@ -1,236 +0,0 @@ -"""Phase 4: BGRL "A3" pretraining fidelity (merged-unit sampling, momentum schedule, -legacy checkpoint import, preset end to end).""" - -from __future__ import annotations - - -import numpy as np -import pytest -import torch - -from tests.helpers.synthetic_cohort import cohort_overrides, compose_cfg, make_synthetic_cohort -from tests.models.test_fidelity_phase2 import _runtime, _units - -pytest.importorskip("torch_geometric") -pytest.importorskip("torch_sparse") - - -def test_merge_graph_units_offsets_edges_and_tracks_units(): - from grass_mil.data.components.merge import merge_graph_units - - units = _units(3, 40) - merged = merge_graph_units(units) - assert merged.num_nodes == 120 - assert merged.x.shape == (120, 4) - assert merged.edge_index.shape[1] == 3 * 160 - assert merged.edge_index[:, 160:320].min() >= 40 and merged.edge_index[:, 160:320].max() < 80 - assert merged.unit_index.tolist() == [0] * 40 + [1] * 40 + [2] * 40 - assert merged.unit_region_ids == ["region_0", "region_1", "region_2"] - assert merged.unit_graph_y.shape == (3, 1) - - -def test_merge_units_sampler_mixes_regions_and_covers_all_roots(): - from grass_mil.data.components.samplers import get_sampler_strategy - - units = _units(3, 40) - strategy = get_sampler_strategy( - { - "name": "shadow_custom", - "kwargs": {}, - "runtime": _runtime( - merge_units=True, subgraph_batch_size=32, depth=1, num_neighbors=-1 - ), - } - ) - loader = strategy.build_dataset_loader(dataset=units, batch_size=1, shuffle=True) - assert len(loader) == int(np.ceil(120 / 32)) - batches = list(loader) - assert len(batches) == len(loader) - n_graphs = sum(int(b.ptr.numel() - 1) for b in batches) - assert n_graphs == 120 # every root of every unit exactly once - # A full-size batch spans more than one unit: units are pooled, not iterated. - first = batches[0] - unit_of_graph = first.unit_index[first.ptr[:-1]] - assert len(set(unit_of_graph.tolist())) > 1 - - -def test_momentum_schedule_cosine_to_final_without_warmup(): - from grass_mil.models.training.ssl_runtime import CosineWarmup - - sched = CosineWarmup(base_value=0.99, warmup_steps=0, total_steps=1000, min_value=1.0) - assert sched.value(0) == pytest.approx(0.99) - assert sched.value(500) == pytest.approx(0.995) - assert sched.value(1000) == pytest.approx(1.0) - assert all(sched.value(s) <= sched.value(s + 1) + 1e-12 for s in range(999)) - - -def test_bgrl_module_reads_momentum_final_and_zero_momentum_warmup(): - from grass_mil.models.bgrl_module import BGRLModule - - module = BGRLModule( - encoder={"input_dim": 4, "hidden_dim": 8, "out_dim": 8, "num_layers": 1}, - ssl={"method": "bgrl", "predictor": {"hidden_size": 8}}, - task={ - "warmup_steps": 100, - "total_steps": 1000, - "momentum": 0.99, - "momentum_final": 1.0, - "momentum_warmup_steps": 0, - }, - ) - assert module._momentum_min == 1.0 - assert module._momentum_warmup_steps == 0 - legacy = BGRLModule( - encoder={"input_dim": 4, "hidden_dim": 8, "out_dim": 8, "num_layers": 1}, - ssl={"method": "bgrl", "predictor": {"hidden_size": 8}}, - task={"warmup_steps": 100, "total_steps": 1000, "momentum": 0.99, "momentum_min": 0.99}, - ) - assert legacy._momentum_min == 0.99 and legacy._momentum_warmup_steps == 100 - - -def test_spacegm_bgrl_remap_loads_into_public_encoder(): - from grass_mil.models.components.backbones import EncoderConfig - from grass_mil.models.components.factory import build_encoder - from grass_mil.models.training.checkpoint_init import remap_encoder_keys - - cfg = EncoderConfig( - input_dim=1, - hidden_dim=16, - out_dim=16, - num_layers=2, - conv_type="gine", - norm="layernorm", - pooling="max", - use_edge_attr=True, - edge_attr_dim=1, - edge_feature_index=0, - train_eps=True, - categorical_embedding={ - "label": "cell_type", - "num_embeddings": 5, - "column_index": 0, - "reserve_unassigned": True, - }, - ) - enc = build_encoder(cfg) - ref = {k: v.clone() for k, v in enc.state_dict().items()} - legacy = {} - for k, v in ref.items(): - if k == "input_proj.embedding.weight": - legacy["online_encoder.x_embedding.weight"] = v + 1 - elif k.startswith("input_proj.linear."): - legacy[k.replace("input_proj.linear.", "online_encoder.feat_embedding.")] = v + 1 - elif k.startswith("layers."): - i, rest = k.split(".", 2)[1], k.split(".", 2)[2] - legacy[f"online_encoder.gnns.{i}.gine.{rest}"] = v + 1 - legacy[f"online_encoder.gnns.{i}.mlp.0.weight"] = torch.zeros(1) - legacy[f"online_encoder.gnns.{i}.edge_embedding.weight"] = torch.zeros(4, 16) - elif k.startswith("norms."): - legacy[k.replace("norms.", "online_encoder.batch_norms.")] = v + 1 - legacy["target_encoder.x_embedding.weight"] = torch.zeros(1) - legacy["predictor.net.0.weight"] = torch.zeros(1) - remapped = remap_encoder_keys(legacy, encoder_init_map="spacegm_bgrl") - assert set(remapped) == {f"encoder.{k}" for k in ref} - stripped = {k[len("encoder.") :]: v for k, v in remapped.items()} - enc.load_state_dict(stripped, strict=True) - assert torch.allclose(enc.state_dict()["layers.0.eps"], ref["layers.0.eps"] + 1) - with pytest.raises(ValueError): - remap_encoder_keys({"foo": torch.zeros(1)}, encoder_init_map="spacegm_bgrl") - - -@pytest.fixture(scope="module") -def cohort(tmp_path_factory): - root = tmp_path_factory.mktemp("cohort4") - c = make_synthetic_cohort(root, n_patients=4, regions_per_patient=2, n_cells=100) - import hydra - - cfg = compose_cfg("train", cohort_overrides(c, root / "out_precompute")) - dm = hydra.utils.instantiate(cfg.data) - dm.prepare_data() - return c - - -def test_pretrain_a3_preset_then_finetune_from_checkpoint(cohort, tmp_path): - from grass_mil.train import train - - pre_out = tmp_path / "pretrain" - base = [ - o - for o in cohort_overrides(cohort, pre_out) - if not o.startswith(("regime=", "data.val_sampler", "model.attention", "task.")) - ] - cfg = compose_cfg( - "train", - base - + [ - "experiment=grass_mil_pretrain_a3", - "task.total_steps=4", - "task.warmup_steps=1", - "task.lr_warmup.warmup_steps=1", - "task.augmentation.feature_noise_columns=[area]", - "callbacks.model_checkpoint.every_n_train_steps=2", - "trainer.precision=32", - "data.sampler.runtime.depth=2", - "data.sampler.runtime.subgraph_batch_size=16", - ], - ) - assert cfg.task.momentum_final == 1.0 and cfg.data.sampler.runtime.merge_units is True - assert cfg.trainer.max_steps == 4 and cfg.scheduler.T_max == "auto" - metrics, objects = train(cfg) - module = objects["model"] - assert module.encoder.dropout.p == pytest.approx(0.1) - assert "train/loss" in metrics - ckpts = sorted((pre_out / "checkpoints").glob("*.ckpt")) - names = {p.name for p in ckpts} - assert "last.ckpt" in names and any(n.startswith("step_") for n in names) - # the scheduler covered total_steps - warmup and momentum reached its final value - # last update used the momentum of zero-based step 3 of 4 (cosine 0.99 -> 1.0) - import math - - expected = 1.0 + (0.99 - 1.0) * 0.5 * (1.0 + math.cos(math.pi * 3 / 4)) - assert float(metrics["train/momentum"]) == pytest.approx(expected, abs=1e-5) - pre_state = torch.load(pre_out / "checkpoints" / "last.ckpt", weights_only=False)["state_dict"] - enc_keys = [ - k for k in pre_state if k.startswith("ssl_model.online_encoder.layers.0.nn.0.weight") - ] - assert enc_keys - - ft_out = tmp_path / "finetune" - ft_cfg = compose_cfg( - "train", - cohort_overrides(cohort, ft_out) - + [ - f"model.init_from_ckpt={pre_out / 'checkpoints' / 'last.ckpt'}", - "trainer.max_epochs=1", - ], - ) - _, ft_objects = train(ft_cfg) - ft_module = ft_objects["model"] - # the fine-tuned model started from the pretrained encoder weights - loaded = ft_module.encoder.state_dict()["norms.0.weight"].detach().cpu() - pretrained = pre_state["ssl_model.online_encoder.norms.0.weight"].detach().cpu() - assert torch.allclose(loaded, pretrained, atol=5e-3) - assert ft_module.encoder.dropout.p == 0.0 - - -def test_zero_validation_callback_filter_keeps_interpolations_resolvable(): - from omegaconf import OmegaConf - - from grass_mil.train import _filter_val_monitor_callbacks - - cfg = OmegaConf.create( - { - "paths": {"output_dir": "/tmp/run"}, - "callbacks": { - "model_checkpoint": { - "_target_": "x", - "dirpath": "${paths.output_dir}/checkpoints", - "monitor": None, - }, - "early_stopping": {"_target_": "y", "monitor": "val/loss"}, - }, - } - ) - filtered, removed = _filter_val_monitor_callbacks(cfg.callbacks) - assert removed == ["early_stopping"] - assert filtered.model_checkpoint.dirpath == "/tmp/run/checkpoints" diff --git a/tests/models/test_fidelity_phase5.py b/tests/models/test_fidelity_phase5.py deleted file mode 100644 index be19dd3..0000000 --- a/tests/models/test_fidelity_phase5.py +++ /dev/null @@ -1,273 +0,0 @@ -"""Phase 5: cell cache extraction, patient aggregation, evaluation statistics.""" - -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import pandas as pd -import pytest -import torch - -from tests.helpers.synthetic_cohort import cohort_overrides, compose_cfg, make_synthetic_cohort - -pytest.importorskip("torch_geometric") -pytest.importorskip("sklearn") - - -# ---------------------------------------------------------------------------- -# statistics -# ---------------------------------------------------------------------------- - - -def test_auroc_matches_sklearn_and_delong_identity_and_symmetry(): - from sklearn.metrics import roc_auc_score - - from grass_mil.evaluation import auroc, delong_test - - rng = np.random.default_rng(0) - y = rng.integers(0, 2, 80) - s1 = y + rng.normal(0, 0.5, 80) - s2 = rng.normal(size=80) - assert auroc(y, s1) == pytest.approx(roc_auc_score(y, s1)) - same = delong_test(y, s1, s1) - assert same.delta == 0.0 and same.p_value == pytest.approx(1.0) - d12 = delong_test(y, s1, s2) - d21 = delong_test(y, s2, s1) - assert d12.delta == pytest.approx(-d21.delta) and d12.p_value == pytest.approx(d21.p_value) - assert d12.auroc_1 == pytest.approx(roc_auc_score(y, s1)) - assert d12.p_value < 0.05 # informative vs random score on 80 units - - -def test_bootstrap_ci_brackets_point_and_group_bootstrap_runs(): - from grass_mil.evaluation import bootstrap_auroc_ci - - rng = np.random.default_rng(1) - y = rng.integers(0, 2, 60) - s = y + rng.normal(0, 1, 60) - ci = bootstrap_auroc_ci(y, s, n_boot=300, seed=0) - assert ci.lower <= ci.auroc <= ci.upper and ci.n_valid > 0 - groups = np.repeat(np.arange(30), 2) - ci_g = bootstrap_auroc_ci(y, s, n_boot=300, seed=0, groups=groups) - assert ci_g.lower <= ci_g.auroc <= ci_g.upper - - -def test_within_fold_standardisation_and_per_fold_auroc(): - from grass_mil.evaluation import per_fold_auroc, rank_within_fold, z_within_fold - - folds = np.array([0] * 5 + [1] * 5) - s = np.array([1, 2, 3, 4, 5, 10, 20, 30, 40, 50], dtype=float) - z = z_within_fold(s, folds) - assert z[:5].mean() == pytest.approx(0) and z[5:].std() == pytest.approx(1) - assert np.allclose(z[:5], z[5:]) # scale-free - r = rank_within_fold(s, folds) - assert np.allclose(r[:5], (np.arange(5) + 0.5) / 5) - y = np.array([0, 0, 1, 1, 1, 0, 1, 0, 1, 1]) - pf = per_fold_auroc(y, s, folds) - assert pf["n_folds"] == 2 and 0 <= pf["mean"] <= 1 and "fold_0" in pf - - -def test_operating_points_from_confusion_and_derived_quantities(): - from grass_mil.evaluation import ( - confusion_counts, - flag_top_within_fold, - likelihood_ratios, - metrics_from_confusion, - ppv_npv_at_prevalence, - threshold_points, - ) - - y = np.array([1, 1, 1, 1, 0, 0, 0, 0, 0, 0]) - pred = np.array([1, 1, 1, 0, 1, 1, 0, 0, 0, 0]) - c = confusion_counts(pred, y) - assert (c.tp, c.fp, c.fn, c.tn) == (3, 2, 1, 4) - m = metrics_from_confusion(c) - assert m["sensitivity"] == pytest.approx(0.75) and m["specificity"] == pytest.approx(4 / 6) - assert m["ppv"] == pytest.approx(0.6) and m["npv"] == pytest.approx(0.8) - lr = likelihood_ratios(0.75, 4 / 6) - assert lr["lr_plus"] == pytest.approx(0.75 / (1 - 4 / 6)) - r = ppv_npv_at_prevalence(0.75, 4 / 6, 0.4) - assert r["ppv"] == pytest.approx(0.75 * 0.4 / (0.75 * 0.4 + (1 / 3) * 0.6)) - folds = np.array([0] * 5 + [1] * 5) - score = np.arange(10, dtype=float) - flagged = flag_top_within_fold(score, folds, 0.4) - assert flagged.tolist() == [0, 0, 0, 1, 1, 0, 0, 0, 1, 1] # top 40 % of EACH fold - t = threshold_points(score, y) - assert set(t) == {"youden", "rule_out", "rule_in"} - - -# ---------------------------------------------------------------------------- -# patient aggregation -# ---------------------------------------------------------------------------- - - -def _payload(): - from grass_mil.inference.schemas import BatchPredictionPayload - - torch.manual_seed(0) - regions = ["P1_a", "P1_b", "P2_a"] - counts = [3, 2, 4] - inst_logits = torch.randn(sum(counts), 2) - inst_attn = torch.randn(sum(counts), 2) - inst_region = [r for r, n in zip(regions, counts) for _ in range(n)] - bag_logits = torch.stack( - [ - ( - inst_logits[[i for i, rr in enumerate(inst_region) if rr == r]] - * torch.softmax( - inst_attn[[i for i, rr in enumerate(inst_region) if rr == r]], dim=0 - ) - ).sum(0) - for r in regions - ] - ) - return ( - BatchPredictionPayload( - bag_ids=regions, - bag_logits=bag_logits, - bag_targets=torch.tensor([[1.0], [1.0], [0.0]]), - bag_attention=None, - row_region_ids=regions, - row_sample_ids=regions, - instance_logits=inst_logits, - instance_attention_logits=inst_attn, - instance_patch_ids=[f"p{i}" for i in range(sum(counts))], - instance_bag_ids=inst_region, - instance_region_ids=inst_region, - instance_sample_ids=inst_region, - ), - inst_logits, - inst_attn, - inst_region, - ) - - -def test_patient_mean_probability_and_merged_bag_aggregation(): - from grass_mil.inference.aggregation import aggregate_group_logits - - payload, inst_logits, inst_attn, inst_region = _payload() - gmap = {"P1_a": "P1", "P1_b": "P1", "P2_a": "P2"} - mean_p = aggregate_group_logits( - payload, mode="mean_probability", bag_scope="patient", group_map=gmap - ) - assert mean_p.bag_ids == ["P1", "P2"] - probs = torch.softmax(mean_p.bag_logits, dim=1) - expected = torch.softmax(payload.bag_logits[:2], dim=1).mean(0) - assert torch.allclose(probs[0], expected, atol=1e-6) - assert mean_p.bag_targets[0].item() == 1.0 and mean_p.bag_targets[1].item() == 0.0 - - merged = aggregate_group_logits( - payload, mode="attention_weighted", bag_scope="patient", group_map=gmap - ) - idx = [i for i, r in enumerate(inst_region) if gmap[r] == "P1"] - direct = (inst_logits[idx] * torch.softmax(inst_attn[idx], dim=0)).sum(0) - assert torch.allclose(merged.bag_logits[0], direct, atol=1e-6) - with pytest.raises(ValueError, match="patient"): - aggregate_group_logits(payload, mode="mean_probability", bag_scope="patient") - - -# ---------------------------------------------------------------------------- -# extraction + scores on the synthetic cohort -# ---------------------------------------------------------------------------- - - -@pytest.fixture(scope="module") -def cohort(tmp_path_factory): - root = tmp_path_factory.mktemp("cohort5") - c = make_synthetic_cohort(root, n_patients=4, regions_per_patient=2, n_cells=80) - import hydra - - cfg = compose_cfg("train", cohort_overrides(c, root / "out_precompute")) - dm = hydra.utils.instantiate(cfg.data) - dm.prepare_data() - return c - - -def test_extract_cell_cache_covers_every_cell(cohort, tmp_path): - from grass_mil.extract import run_extract - from grass_mil.inference.cell_cache import CellCache - from grass_mil.train import train - - train_out = tmp_path / "train" - cfg = compose_cfg("train", cohort_overrides(cohort, train_out) + ["trainer.max_epochs=1"]) - train(cfg) - ckpt = train_out / "checkpoints" / "last.ckpt" - assert ckpt.exists() - - ext_out = tmp_path / "extract" - ecfg = compose_cfg( - "extract", - cohort_overrides(cohort, ext_out) - + [f"+extract.models.final={ckpt}", "extract.batch_size=32", "extract.device=cpu"], - ) - result = run_extract(ecfg) - cache = CellCache.open(Path(result["models"][0])) - assert sorted(cache.regions) == sorted(cohort.region_ids) - assert len(cache.bias) == 2 and set(cache.cell_type_names) == {"A", "B", "C", "D"} - rid = cache.regions[0] - arrays = cache.load(rid) - n = cache.n_cells(rid) - assert arrays["E"].shape == (n, 16) and arrays["E"].dtype == np.float16 - assert arrays["LG"].shape == (n, 2) and arrays["AR"].shape == (n, 2) - assert arrays["ct"].shape == (n,) and arrays["pos"].shape == (n, 2) - assert arrays["edge_index"].shape[0] == 2 and arrays["edge_index"].max() < n - assert np.isfinite(arrays["E"].astype(np.float32)).all() - assert cache.label_of(rid) == float(cohort.label_of[rid]) - # row order = cell order of the region table (first cell's type matches the CSV) - csv = pd.read_csv(cohort.manifest_path).set_index("region_id").loc[rid, "input_path"] - types = pd.read_csv(csv)["cell_type"].tolist() - assert [cache.cell_type_names[c] for c in arrays["ct"]] == types - - -def test_scores_cli_on_synthetic_oof_table(tmp_path): - from grass_mil.evaluation.cli import run_scores - - rng = np.random.default_rng(0) - rows = [] - for outer in range(3): - for p in range(20): - pid = f"o{outer}_P{p}" - label = int(p % 2 == 0) - for core in range(2): - rows.append( - { - "outer": outer, - "region_id": f"{pid}_c{core}", - "patient_id": pid, - "label": label, - "prob_class_1": 0.3 + 0.3 * label + rng.normal(0, 0.15) + 0.1 * outer, - "other": rng.uniform(), - "subtype": "A" if p < 10 else "B", - } - ) - table = tmp_path / "oof.csv" - pd.DataFrame(rows).to_csv(table, index=False) - cfg = compose_cfg( - "scores", - [ - f"input={table}", - f"output_dir={tmp_path / 'scores'}", - "scores=[prob_class_1,other]", - "strata=[subtype]", - "n_boot=200", - ], - ) - summary = run_scores(cfg) - assert summary["n"] == 60 # patient level - first = summary["scores"][0] - assert first["auroc"] > 0.8 and first["lower"] <= first["auroc"] <= first["upper"] - assert first["perfold_n_folds"] == 3 - assert summary["delong"][0]["p_value"] < 0.05 - out = tmp_path / "scores" - for name in ( - "scores_summary.csv", - "delong.csv", - "operating_points.csv", - "strata.csv", - "ppv_npv_by_prevalence.csv", - "report_scores.md", - ): - assert (out / name).exists(), name - ops = pd.read_csv(out / "operating_points.csv") - top30 = ops[(ops.score == "prob_class_1") & (ops.rule.str.startswith("top 30"))].iloc[0] - assert abs(top30["flagged"] - 0.3) < 0.06 diff --git a/tests/models/test_losses.py b/tests/models/test_losses.py deleted file mode 100644 index 4551a3f..0000000 --- a/tests/models/test_losses.py +++ /dev/null @@ -1,63 +0,0 @@ -import pytest -import torch - - -def test_weighted_cross_entropy(): - from grass_mil.models.components.losses import WeightedCrossEntropyLoss - - loss_fn = WeightedCrossEntropyLoss() - logits = torch.randn(12, 3) - target = torch.randint(0, 3, (12,)) - weight = torch.rand(12) - loss = loss_fn(logits, target, sample_weight=weight) - assert loss.ndim == 0 - assert torch.isfinite(loss) - - -def test_weighted_bce(): - from grass_mil.models.components.losses import WeightedBCEWithLogitsLoss - - loss_fn = WeightedBCEWithLogitsLoss() - logits = torch.randn(12, 1) - target = torch.randint(0, 2, (12, 1)).float() - weight = torch.rand(12, 1) - loss = loss_fn(logits, target, sample_weight=weight) - assert loss.ndim == 0 - assert torch.isfinite(loss) - - -@pytest.mark.parametrize("loss_name", ["mse", "huber"]) -def test_regression_losses(loss_name): - from grass_mil.models.components.losses import WeightedHuberLoss, WeightedMSELoss - - pred = torch.randn(12, 2) - target = torch.randn(12, 2) - sample_weight = torch.rand(12) - loss_fn = WeightedMSELoss() if loss_name == "mse" else WeightedHuberLoss(delta=1.0) - loss = loss_fn(pred, target, sample_weight=sample_weight) - assert loss.ndim == 0 - assert torch.isfinite(loss) - - -def test_coxsgd_loss(): - from grass_mil.models.components.losses import CoxSGDLoss - - loss_fn = CoxSGDLoss(top_n=2, regularizer_weight=0.01) - y_pred = torch.randn(16, 1) - length = torch.arange(16).float().flip(0) - event = torch.randint(0, 2, (16,)).float() - loss = loss_fn(y_pred, length, event) - assert loss.ndim == 0 - assert torch.isfinite(loss) - - -def test_coxsgd_top_n_handles_tiny_batches(): - from grass_mil.models.components.losses import CoxSGDLoss - - loss_fn = CoxSGDLoss(top_n=10, regularizer_weight=0.01) - y_pred = torch.randn(1, 1) - length = torch.tensor([1.0]) - event = torch.tensor([1.0]) - loss = loss_fn(y_pred, length, event) - assert loss.ndim == 0 - assert torch.isfinite(loss) diff --git a/tests/models/test_node_embeddings.py b/tests/models/test_node_embeddings.py deleted file mode 100644 index d6a0663..0000000 --- a/tests/models/test_node_embeddings.py +++ /dev/null @@ -1,181 +0,0 @@ -"""Contracts for the encoder's input embeddings. - -The published encoder sums a learned cell-type embedding with a projection of -the continuous node features, and conditions message passing on the scalar edge -length alone. -""" - -from __future__ import annotations - -import pytest -import torch - -from grass_mil.models.components.embeddings import ( - CategoricalEmbeddingConfig, - NodeInputEmbedding, - resolve_categorical_embedding_config, -) - - -def _categorical(**overrides): - params = {"label": "cell_type", "num_embeddings": 5, "column_index": 0} - params.update(overrides) - return CategoricalEmbeddingConfig(**params) - - -def test_node_embedding_sums_rather_than_concatenates() -> None: - layer = NodeInputEmbedding(input_dim=3, hidden_dim=8, categorical=_categorical()) - x = torch.randn(6, 3) - codes = torch.randint(0, 5, (6, 1)) - - out = layer(x, codes) - assert out.shape == (6, 8), "embedding must be summed into the hidden dim, not concatenated" - - expected = layer.linear(x) + layer.embedding(codes[:, 0]) - torch.testing.assert_close(out, expected) - - -def test_node_embedding_supports_cell_type_only_cohorts() -> None: - """`x` legitimately has zero columns when no continuous features exist.""" - layer = NodeInputEmbedding(input_dim=0, hidden_dim=8, categorical=_categorical()) - codes = torch.randint(0, 5, (4, 1)) - - out = layer(torch.zeros(4, 0), codes) - assert out.shape == (4, 8) - torch.testing.assert_close(out, layer.embedding(codes[:, 0])) - - -def test_node_embedding_without_categorical_matches_plain_projection() -> None: - """Existing configs must be unaffected by the new embedding path.""" - identity = NodeInputEmbedding(input_dim=8, hidden_dim=8, categorical=None) - x = torch.randn(5, 8) - torch.testing.assert_close(identity(x), x) - - projected = NodeInputEmbedding(input_dim=4, hidden_dim=8, categorical=None) - torch.testing.assert_close(projected(x[:, :4]), projected.linear(x[:, :4])) - - -def test_node_embedding_reserves_an_unassigned_row() -> None: - reserved = NodeInputEmbedding( - input_dim=2, hidden_dim=4, categorical=_categorical(reserve_unassigned=True) - ) - assert reserved.embedding.num_embeddings == 6 - assert reserved.unassigned_index == 5 - - plain = NodeInputEmbedding( - input_dim=2, hidden_dim=4, categorical=_categorical(reserve_unassigned=False) - ) - assert plain.embedding.num_embeddings == 5 - assert plain.unassigned_index is None - - -def test_node_embedding_rejects_out_of_range_codes() -> None: - """Clamping would silently train on the wrong cell type.""" - layer = NodeInputEmbedding( - input_dim=2, hidden_dim=4, categorical=_categorical(reserve_unassigned=False) - ) - codes = torch.tensor([[0], [99]]) - with pytest.raises(ValueError, match="out of range"): - layer(torch.randn(2, 2), codes) - - -def test_node_embedding_requires_codes_when_configured() -> None: - layer = NodeInputEmbedding(input_dim=2, hidden_dim=4, categorical=_categorical()) - with pytest.raises(ValueError, match="no `categorical_codes`"): - layer(torch.randn(2, 2), None) - - -def test_zero_input_dim_without_categorical_is_rejected() -> None: - with pytest.raises(ValueError, match="no input features"): - NodeInputEmbedding(input_dim=0, hidden_dim=4, categorical=None) - - -def test_resolve_categorical_config_prefers_explicit_values() -> None: - resolved = resolve_categorical_embedding_config( - {"label": "cell_type", "num_embeddings": 3}, - cardinalities={"cell_type": 17}, - slices={"cell_type": 2}, - ) - assert resolved.num_embeddings == 3, "explicit config must win over inference" - assert resolved.column_index == 2, "unset fields are inferred" - - -def test_resolve_categorical_config_reports_unresolvable_label() -> None: - with pytest.raises(ValueError, match="vocabulary size"): - resolve_categorical_embedding_config({"label": "missing"}, cardinalities={}, slices={}) - - -def _encoder(**overrides): - from grass_mil.models.components.backbones import EncoderConfig, GNNEncoder - - params = { - "input_dim": 2, - "hidden_dim": 8, - "out_dim": 8, - "num_layers": 2, - "conv_type": "gine", - "norm": "layernorm", - "pooling": "max", - "use_edge_attr": True, - "edge_attr_dim": 1, - "edge_feature_index": 0, - } - params.update(overrides) - return GNNEncoder(EncoderConfig(**params)) - - -def _toy_graph(num_edge_features: int = 2): - edge_index = torch.tensor([[0, 1, 2, 3], [1, 2, 3, 0]], dtype=torch.long) - x = torch.randn(4, 2) - edge_attr = torch.rand(edge_index.size(1), num_edge_features) - codes = torch.randint(0, 5, (4, 1)) - return x, edge_index, edge_attr, codes - - -def test_gine_reads_only_the_selected_edge_column() -> None: - """The conv is conditioned on edge length; other columns must not leak in. - - This is the fidelity assertion for the edge path: perturbing any - non-selected column must leave the output bit-identical. - """ - torch.manual_seed(0) - encoder = _encoder().eval() - x, edge_index, edge_attr, _ = _toy_graph(num_edge_features=3) - - with torch.no_grad(): - baseline = encoder(x, edge_index, edge_attr=edge_attr) - - perturbed = edge_attr.clone() - perturbed[:, 1:] += 17.0 - after = encoder(x, edge_index, edge_attr=perturbed) - assert torch.equal(baseline, after) - - with torch.no_grad(): - length_changed = edge_attr.clone() - length_changed[:, 0] += 17.0 - moved = encoder(x, edge_index, edge_attr=length_changed) - assert not torch.allclose(baseline, moved), "the selected length column must matter" - - -def test_encoder_embeds_cell_type_from_categorical_codes() -> None: - torch.manual_seed(0) - encoder = _encoder( - categorical_embedding={"label": "cell_type", "num_embeddings": 5, "column_index": 0} - ).eval() - x, edge_index, edge_attr, codes = _toy_graph() - - with torch.no_grad(): - baseline = encoder(x, edge_index, edge_attr=edge_attr, categorical_codes=codes) - shuffled = encoder(x, edge_index, edge_attr=edge_attr, categorical_codes=(codes + 1) % 5) - assert baseline.shape == (4, 8) - assert not torch.allclose(baseline, shuffled), "cell type must influence the encoding" - - -def test_encoder_rejects_edge_feature_index_for_non_gine_convs() -> None: - with pytest.raises(ValueError, match="conv_type='gin'"): - _encoder(conv_type="gin", edge_attr_dim=None, use_edge_attr=False) - - -def test_encoder_requires_unit_edge_dim_when_selecting_one_column() -> None: - with pytest.raises(ValueError, match="edge_attr_dim"): - _encoder(edge_attr_dim=2) diff --git a/tests/models/test_optimization.py b/tests/models/test_optimization.py deleted file mode 100644 index aa90e3d..0000000 --- a/tests/models/test_optimization.py +++ /dev/null @@ -1,153 +0,0 @@ -import torch -import pytest - - -def _encoder_cfg(): - return { - "input_dim": 8, - "hidden_dim": 16, - "out_dim": 16, - "num_layers": 2, - "dropout": 0.1, - "conv_type": "gin", - "norm": "batchnorm", - "jk": "last", - "act": "relu", - "pooling": "mean", - "gat_heads": 2, - "use_edge_attr": False, - "edge_weight_index": 0, - "edge_attr_dim": None, - } - - -def _optim_cfg(): - return {"_target_": "torch.optim.AdamW", "lr": 1e-3, "weight_decay": 1e-2} - - -def _cosine_epoch_cfg(): - return { - "_target_": "torch.optim.lr_scheduler.CosineAnnealingLR", - "T_max": 10, - "eta_min": 1e-6, - } - - -def _cosine_step_cfg(): - return { - "_target_": "torch.optim.lr_scheduler.CosineAnnealingLR", - "T_max": 100, - "eta_min": 1e-7, - } - - -def test_supervised_mil_uses_explicit_backbone_and_attention_lrs(): - pytest.importorskip("torch_geometric") - from grass_mil.models.supervised_module import SupervisedModule - - module = SupervisedModule( - encoder=_encoder_cfg(), - graph_head={ - "input_dim": 16, - "output_dim": 1, - "hidden_dim": 16, - "num_layers": 2, - "dropout": 0.0, - }, - attention={ - "attention_type": "gated_projected", - "input_dim": 16, - "projection_dim": 8, - "hidden_dim": 4, - "dropout": False, - "n_classes": 1, - }, - loss={"loss_type": "categorical_bce"}, - optim=_optim_cfg(), - scheduler=_cosine_epoch_cfg(), - task={ - "aggregation": "mil_attention", - "target_type": "binary", - "loss": "categorical_bce", - "optimization": {"backbone_lr": 1e-3, "attention_lr": 4e-3}, - "lr_warmup": {"enabled": False}, - "region_accumulation": {"enabled": False}, - }, - ) - module.setup("fit") - out = module.configure_optimizers() - optimizer = out["optimizer"] if isinstance(out, dict) else out - - lrs = sorted([group["lr"] for group in optimizer.param_groups]) - assert len(optimizer.param_groups) == 2 - assert lrs == [1e-3, 4e-3] - - -def test_supervised_mean_warmup_is_opt_in(): - pytest.importorskip("torch_geometric") - from grass_mil.models.supervised_module import SupervisedModule - - base_kwargs = dict( - encoder=_encoder_cfg(), - graph_head={ - "input_dim": 16, - "output_dim": 1, - "hidden_dim": 16, - "num_layers": 2, - "dropout": 0.0, - }, - loss={"loss_type": "categorical_bce"}, - optim=_optim_cfg(), - scheduler=_cosine_epoch_cfg(), - ) - base_task = {"target_type": "binary", "loss": "categorical_bce"} - - no_warmup = SupervisedModule( - **base_kwargs, - task={ - **base_task, - "lr_warmup": {"enabled": False, "warmup_steps": 0, "start_factor": 0.1}, - }, - ) - no_warmup.setup("fit") - no_warmup_out = no_warmup.configure_optimizers() - no_warmup_scheduler = no_warmup_out["lr_scheduler"]["scheduler"] - assert isinstance(no_warmup_scheduler, torch.optim.lr_scheduler.CosineAnnealingLR) - - with_warmup = SupervisedModule( - **base_kwargs, - task={ - **base_task, - "lr_warmup": {"enabled": True, "warmup_steps": 2, "start_factor": 0.1}, - }, - ) - with_warmup.setup("fit") - with_warmup_out = with_warmup.configure_optimizers() - with_warmup_scheduler = with_warmup_out["lr_scheduler"]["scheduler"] - assert isinstance(with_warmup_scheduler, torch.optim.lr_scheduler.SequentialLR) - - -def test_bgrl_uses_warmup_then_base_scheduler(): - pytest.importorskip("torch_geometric") - from grass_mil.models.bgrl_module import BGRLModule - - module = BGRLModule( - encoder=_encoder_cfg(), - ssl={"method": "bgrl", "predictor": {"hidden_size": 32}}, - optim=_optim_cfg(), - scheduler=_cosine_step_cfg(), - task={ - "total_steps": 100, - "warmup_steps": 10, - "momentum": 0.99, - "momentum_min": 0.99, - "lr_warmup": {"enabled": True, "warmup_steps": 5, "start_factor": 0.1}, - }, - ) - module.setup("fit") - out = module.configure_optimizers() - scheduler = out["lr_scheduler"]["scheduler"] - - assert isinstance(scheduler, torch.optim.lr_scheduler.SequentialLR) - assert len(scheduler._schedulers) == 2 - assert isinstance(scheduler._schedulers[0], torch.optim.lr_scheduler.LinearLR) diff --git a/tests/models/test_runtime.py b/tests/models/test_runtime.py deleted file mode 100644 index 570a2f0..0000000 --- a/tests/models/test_runtime.py +++ /dev/null @@ -1,615 +0,0 @@ -import pytest -import torch - - -def test_remap_encoder_keys_auto_bgrl_or_identity(): - from grass_mil.models.training.checkpoint_init import remap_encoder_keys - - sd = { - "online_encoder.layers.0.weight": torch.randn(4, 4), - "ssl_model.online_encoder.layers.0.bias": torch.randn(4), - "graph_head.net.0.weight": torch.randn(4, 4), - } - out = remap_encoder_keys(sd, encoder_init_map="auto_bgrl_or_identity") - assert "encoder.layers.0.weight" in out - assert "encoder.layers.0.bias" in out - assert "graph_head.net.0.weight" in out - - -def test_compute_supervised_loss_routes(): - from grass_mil.models.components.losses import ( - CoxSGDLoss, - WeightedBCEWithLogitsLoss, - WeightedCrossEntropyLoss, - WeightedMSELoss, - ) - from grass_mil.models.training.loss_utils import compute_supervised_loss - - bce = compute_supervised_loss( - loss_fn=WeightedBCEWithLogitsLoss(), - task_cfg={"target_type": "binary"}, - bag_logits=torch.randn(4, 1), - bag_targets=torch.randint(0, 2, (4, 1)).float(), - bag_weights=torch.ones(4, 1), - ) - assert bce.ndim == 0 - - mse = compute_supervised_loss( - loss_fn=WeightedMSELoss(), - task_cfg={"target_type": "regression"}, - bag_logits=torch.randn(4, 1), - bag_targets=torch.randn(4, 1), - bag_weights=torch.ones(4, 1), - ) - assert mse.ndim == 0 - - ce = compute_supervised_loss( - loss_fn=WeightedCrossEntropyLoss(), - task_cfg={"target_type": "categorical"}, - bag_logits=torch.randn(5, 3), - bag_targets=torch.tensor([[0], [1], [2], [1], [0]], dtype=torch.long), - bag_weights=torch.ones(5), - ) - assert ce.ndim == 0 - - cox = compute_supervised_loss( - loss_fn=CoxSGDLoss(), - task_cfg={"target_type": "survival"}, - bag_logits=torch.randn(6, 1), - bag_targets=torch.stack( - [ - torch.linspace(1, 6, 6), - torch.tensor([1, 0, 1, 1, 0, 1], dtype=torch.float32), - ], - dim=1, - ), - bag_weights=None, - ) - assert cox.ndim == 0 - - -def test_validate_task_config_accepts_categorical(): - from grass_mil.models.training.builders import validate_task_config - - validate_task_config({"target_type": "categorical", "instance_sampling": "all"}) - - -def test_supervised_module_validation_logs_step_loss_not_acc(): - pytest.importorskip("lightning") - from grass_mil.models.supervised_module import SupervisedModule - - module = SupervisedModule( - encoder={}, - graph_head={}, - loss={}, - optim={}, - task={"aggregation": "mean", "target_type": "categorical"}, - ) - - logged = {} - - def _capture_log(name, value, **kwargs): - logged[name] = (value, kwargs) - - module.log = _capture_log # type: ignore[method-assign] - module._log_stage_metrics( - stage="val", - loss_value=torch.tensor(0.2), - bag_logits=torch.tensor([[5.0, 1.0, -1.0], [-2.0, 0.2, 3.0], [0.1, 0.2, 0.3]]), - bag_targets=torch.tensor([0, 2, 1]), - ) - assert "val/acc" not in logged - assert "val/loss" in logged - _, kwargs = logged["val/loss"] - assert kwargs["on_step"] is True - assert kwargs["on_epoch"] is True - - -def test_supervised_module_logs_epoch_validation_accuracy_from_aggregated_regions(): - pytest.importorskip("lightning") - from grass_mil.models.supervised_module import SupervisedModule - - module = SupervisedModule( - encoder={}, - graph_head={}, - loss={}, - optim={}, - task={"aggregation": "mean", "target_type": "binary"}, - ) - logged = {} - - def _capture_log(name, value, **kwargs): - logged[name] = (value, kwargs) - - module.log = _capture_log # type: ignore[method-assign] - module.on_validation_epoch_start() - module._val_bag_ids = ["p0", "p1", "p2"] - module._val_bag_logits_chunks = [torch.tensor([[0.0], [0.0], [0.0]])] - module._val_bag_targets_chunks = [torch.tensor([[1.0], [1.0], [0.0]])] - module._val_row_region_ids = ["rA", "rA", "rB"] - module._val_row_sample_ids = ["s0", "s0", "s1"] - module._val_instance_logits_chunks = [torch.tensor([[2.0], [-2.0], [-2.0]])] - module._val_instance_patch_ids = ["p0", "p1", "p2"] - module._val_instance_region_ids = ["rA", "rA", "rB"] - module._val_instance_sample_ids = ["s0", "s0", "s1"] - - module.on_validation_epoch_end() - assert "val/acc" in logged - acc_value, kwargs = logged["val/acc"] - assert torch.isclose(acc_value, torch.tensor(1.0)) - assert kwargs["on_step"] is False - assert kwargs["on_epoch"] is True - - -def test_supervised_module_logs_epoch_validation_regression_metrics(): - pytest.importorskip("lightning") - from grass_mil.models.supervised_module import SupervisedModule - - module = SupervisedModule( - encoder={}, - graph_head={}, - loss={}, - optim={}, - task={"aggregation": "mean", "target_type": "regression"}, - ) - logged = {} - - def _capture_log(name, value, **kwargs): - logged[name] = (value, kwargs) - - module.log = _capture_log # type: ignore[method-assign] - module.on_validation_epoch_start() - module._val_bag_ids = ["p0", "p1"] - module._val_bag_logits_chunks = [torch.tensor([[1.0], [3.0]])] - module._val_bag_targets_chunks = [torch.tensor([[1.5], [2.5]])] - module._val_row_region_ids = ["rA", "rB"] - module._val_row_sample_ids = ["s0", "s1"] - module._val_instance_logits_chunks = [torch.tensor([[1.0], [3.0]])] - module._val_instance_patch_ids = ["p0", "p1"] - module._val_instance_region_ids = ["rA", "rB"] - module._val_instance_sample_ids = ["s0", "s1"] - - module.on_validation_epoch_end() - assert "val/mae" in logged - assert "val/rmse" in logged - assert "val/r2" in logged - assert logged["val/mae"][1]["on_epoch"] is True - assert logged["val/rmse"][1]["on_epoch"] is True - assert logged["val/r2"][1]["on_epoch"] is True - - -def test_supervised_module_logs_epoch_validation_survival_c_index(): - pytest.importorskip("lightning") - from grass_mil.models.supervised_module import SupervisedModule - - module = SupervisedModule( - encoder={}, - graph_head={}, - loss={}, - optim={}, - task={"aggregation": "mean", "target_type": "survival"}, - ) - logged = {} - - def _capture_log(name, value, **kwargs): - logged[name] = (value, kwargs) - - module.log = _capture_log # type: ignore[method-assign] - module.on_validation_epoch_start() - module._val_bag_ids = ["p0", "p1", "p2"] - module._val_bag_logits_chunks = [torch.tensor([[3.0], [2.0], [1.0]])] - module._val_bag_targets_chunks = [torch.tensor([[1.0, 1.0], [2.0, 1.0], [3.0, 1.0]])] - module._val_row_region_ids = ["rA", "rB", "rC"] - module._val_row_sample_ids = ["s0", "s1", "s2"] - module._val_instance_logits_chunks = [torch.tensor([[3.0], [2.0], [1.0]])] - module._val_instance_patch_ids = ["p0", "p1", "p2"] - module._val_instance_region_ids = ["rA", "rB", "rC"] - module._val_instance_sample_ids = ["s0", "s1", "s2"] - - module.on_validation_epoch_end() - assert "val/c_index" in logged - c_index_value, kwargs = logged["val/c_index"] - assert torch.isclose(c_index_value, torch.tensor(1.0)) - assert kwargs["on_step"] is False - assert kwargs["on_epoch"] is True - - -def test_attention_bag_aggregation_shapes(): - import torch - from grass_mil.models.components.attention import AttnNetGatedProjected - from grass_mil.models.training.bagging import aggregate_bag_logits_attention - - logits = torch.randn(5, 1) - emb = torch.randn(5, 8) - att = AttnNetGatedProjected(input_dim=8, projection_dim=4, hidden_dim=2, n_classes=1) - grouped = {"r0": [0, 1, 2], "r1": [3, 4]} - bag_logits, bag_ids, bag_attn, bag_indices = aggregate_bag_logits_attention( - logits=logits, - embeddings=emb, - attention=att, - bag_groups=grouped, - max_instances_per_bag=0, - instance_sampling="all", - ) - assert bag_logits.shape == (2, 1) - assert bag_ids == ["r0", "r1"] - assert len(bag_attn["r0"]) == 3 - assert len(bag_indices) == 2 - - -def test_select_target_columns_handles_single_label_list(): - from types import SimpleNamespace - - from grass_mil.models.training.bagging import select_target_columns - - graph_y = torch.tensor([[1.0]]) - batch = SimpleNamespace(graph_label_names=["label"]) - selected = select_target_columns(graph_y, batch, target_columns=["label"]) - assert selected.shape == (1, 1) - assert torch.equal(selected, graph_y) - - -def test_supervised_predict_step_allows_missing_targets(): - pytest.importorskip("lightning") - pytest.importorskip("torch_geometric") - from grass_mil.models.supervised_module import SupervisedModule - - module = SupervisedModule( - encoder={}, - graph_head={}, - loss={}, - optim={}, - task={"aggregation": "mean", "target_type": "binary"}, - ) - - def _forward_bags(_batch, *, allow_missing_targets=False): - assert allow_missing_targets is True - return { - "patch_logits": torch.randn(4, 1), - "bag_logits": torch.randn(2, 1), - "ordered_bag_ids": ["b0", "b1"], - "bag_attention": None, - "bag_indices": [[0, 1], [2, 3]], - "bag_targets": None, - "bag_weights": None, - } - - module._forward_bags = _forward_bags # type: ignore[method-assign] - out = module.predict_step(batch=None, batch_idx=0) - assert out["bag_ids"] == ["b0", "b1"] - assert out["bag_logits"].shape == (2, 1) - assert "bag_targets" not in out - - -def test_supervised_predict_step_can_disable_instance_payload(): - pytest.importorskip("lightning") - pytest.importorskip("torch_geometric") - from grass_mil.models.supervised_module import SupervisedModule - - module = SupervisedModule( - encoder={}, - graph_head={}, - loss={}, - optim={}, - task={"aggregation": "mean", "target_type": "binary"}, - ) - - def _forward_bags(_batch, *, allow_missing_targets=False): - return { - "patch_logits": torch.randn(4, 1), - "bag_logits": torch.randn(2, 1), - "ordered_bag_ids": ["b0", "b1"], - "bag_attention": None, - "bag_indices": [[0, 1], [2, 3]], - "bag_targets": None, - "bag_weights": None, - } - - captured = {"include_instance_payload": None} - - def _build_predict_group_metadata( - _batch, - *, - patch_logits, - graph_emb, - ordered_bag_ids, - bag_indices, - include_instance_payload=True, - ): - captured["include_instance_payload"] = include_instance_payload - payload = { - "row_region_ids": ["r0", "r1"], - "row_sample_ids": ["s0", "s1"], - } - if include_instance_payload: - payload.update( - { - "instance_logits": torch.randn(4, 1), - "instance_patch_ids": ["p0", "p1", "p2", "p3"], - "instance_region_ids": ["r0", "r0", "r1", "r1"], - "instance_sample_ids": ["s0", "s0", "s1", "s1"], - } - ) - return payload - - module._forward_bags = _forward_bags # type: ignore[method-assign] - module._build_predict_group_metadata = _build_predict_group_metadata # type: ignore[method-assign] - module._predict_emit_instance_payload = False - - out = module.predict_step(batch=None, batch_idx=0) - assert captured["include_instance_payload"] is False - assert "instance_logits" not in out - assert "instance_patch_ids" not in out - assert "instance_region_ids" not in out - assert "instance_sample_ids" not in out - - -def test_supervised_predict_step_can_emit_embedding_payload(): - pytest.importorskip("lightning") - pytest.importorskip("torch_geometric") - from grass_mil.models.supervised_module import SupervisedModule - - module = SupervisedModule( - encoder={}, - graph_head={}, - loss={}, - optim={}, - task={"aggregation": "mean", "target_type": "binary"}, - ) - - def _forward_bags(_batch, *, allow_missing_targets=False): - return { - "patch_logits": torch.randn(4, 2), - "graph_emb": torch.tensor([[1.0, 1.0], [3.0, 3.0], [10.0, 10.0], [14.0, 14.0]]), - "node_emb": torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6], [0.7, 0.8]]), - "bag_logits": torch.randn(2, 1), - "ordered_bag_ids": ["b0", "b1"], - "bag_attention": None, - "bag_indices": [[0, 1], [2, 3]], - "bag_targets": None, - "bag_weights": None, - } - - def _build_predict_group_metadata( - _batch, - *, - patch_logits, - graph_emb, - ordered_bag_ids, - bag_indices, - include_instance_payload=True, - ): - return { - "row_region_ids": ["r0", "r1"], - "row_sample_ids": ["s0", "s1"], - } - - module._forward_bags = _forward_bags # type: ignore[method-assign] - module._build_predict_group_metadata = _build_predict_group_metadata # type: ignore[method-assign] - module._predict_emit_instance_payload = False - module._predict_emit_embeddings_payload = True - module._predict_emit_node_embeddings = True - - class _Batch: - region_id = ["r0", "r0", "r1", "r1"] - sample_id = ["s0", "s0", "s1", "s1"] - batch = torch.tensor([0, 1, 2, 3], dtype=torch.long) - - out = module.predict_step(batch=_Batch(), batch_idx=0) - assert out["embedding_bag_ids"] == ["b0", "b1"] - assert torch.allclose(out["graph_embeddings"][0], torch.tensor([2.0, 2.0])) - assert torch.allclose(out["graph_embeddings"][1], torch.tensor([12.0, 12.0])) - assert out["embedding_bag_counts"] == [2, 2] - assert "node_embeddings" in out - assert out["node_bag_ids"] == ["s0::r0", "s0::r0", "s1::r1", "s1::r1"] - - -def test_supervised_predict_step_omits_embedding_payload_by_default(): - pytest.importorskip("lightning") - pytest.importorskip("torch_geometric") - from grass_mil.models.supervised_module import SupervisedModule - - module = SupervisedModule( - encoder={}, - graph_head={}, - loss={}, - optim={}, - task={"aggregation": "mean", "target_type": "binary"}, - ) - - def _forward_bags(_batch, *, allow_missing_targets=False): - return { - "patch_logits": torch.randn(2, 1), - "graph_emb": torch.randn(2, 2), - "node_emb": torch.randn(2, 2), - "bag_logits": torch.randn(1, 1), - "ordered_bag_ids": ["b0"], - "bag_attention": None, - "bag_indices": [[0, 1]], - "bag_targets": None, - "bag_weights": None, - } - - def _build_predict_group_metadata( - _batch, - *, - patch_logits, - graph_emb, - ordered_bag_ids, - bag_indices, - include_instance_payload=True, - ): - return { - "row_region_ids": ["r0"], - "row_sample_ids": ["s0"], - } - - module._forward_bags = _forward_bags # type: ignore[method-assign] - module._build_predict_group_metadata = _build_predict_group_metadata # type: ignore[method-assign] - module._predict_emit_instance_payload = False - - class _Batch: - region_id = ["r0", "r0"] - sample_id = ["s0", "s0"] - - out = module.predict_step(batch=_Batch(), batch_idx=0) - assert "embedding_bag_ids" not in out - assert "graph_embeddings" not in out - assert "embedding_bag_counts" not in out - - -def test_extract_bag_ids_falls_back_per_item_when_primary_missing(): - from types import SimpleNamespace - - from grass_mil.models.training.bagging import extract_bag_ids - - batch = SimpleNamespace(region_id=[None, "", "r2"], sample_id=["s0", "s1", "s2"]) - bag_ids = extract_bag_ids(batch, bag_key="region_id", bag_fallback_key="sample_id") - assert bag_ids == ["s0", "s1", "s2::r2"] - - -def test_extract_bag_ids_namespaces_primary_by_fallback_for_uniqueness(): - from types import SimpleNamespace - - from grass_mil.models.training.bagging import extract_bag_ids - - batch = SimpleNamespace(region_id=["r0", "r0"], sample_id=["s0", "s1"]) - bag_ids = extract_bag_ids(batch, bag_key="region_id", bag_fallback_key="sample_id") - assert bag_ids == ["s0::r0", "s1::r0"] - - -def test_extract_bag_ids_patch_key_uses_sample_region_patch_hierarchy(): - from types import SimpleNamespace - - from grass_mil.models.training.bagging import extract_bag_ids - - batch = SimpleNamespace( - patch_id=["p0", "p1"], - region_id=["r0", "r1"], - sample_id=["s0", "s0"], - ) - bag_ids = extract_bag_ids(batch, bag_key="patch_id", bag_fallback_key="sample_id") - assert bag_ids == ["s0::r0::p0", "s0::r1::p1"] - - -def test_extract_bag_ids_patch_key_fallback_uses_sample_region_when_patch_missing(): - from types import SimpleNamespace - - from grass_mil.models.training.bagging import extract_bag_ids - - batch = SimpleNamespace( - patch_id=[None, ""], - region_id=["r0", "r1"], - sample_id=["s0", "s1"], - ) - bag_ids = extract_bag_ids(batch, bag_key="patch_id", bag_fallback_key="sample_id") - assert bag_ids == ["s0::r0", "s1::r1"] - - -def test_predict_group_metadata_uses_batch_patch_ids_for_instance_patch_ids(): - pytest.importorskip("lightning") - pytest.importorskip("torch_geometric") - from types import SimpleNamespace - - from grass_mil.models.supervised_module import SupervisedModule - - module = SupervisedModule( - encoder={}, - graph_head={}, - loss={}, - optim={}, - task={"aggregation": "mean", "target_type": "binary"}, - ) - batch = SimpleNamespace( - region_id=["r0", "r1", "r0", "r1"], - sample_id=["s0", "s0", "s0", "s0"], - patch_id=["p0", "p1", "p2", None], - ) - metadata = module._build_predict_group_metadata( - batch, - patch_logits=torch.randn(4, 1), - graph_emb=torch.randn(4, 8), - ordered_bag_ids=["bag_a", "bag_b"], - bag_indices=[[0, 2], [1, 3]], - ) - assert metadata["instance_patch_ids"] == ["p0", "p2", "p1", "bag_b"] - - -def test_gather_bag_targets_requires_consistent_labels_within_bag(): - from types import SimpleNamespace - - from grass_mil.models.training.bagging import gather_bag_targets - - batch = SimpleNamespace( - graph_y=torch.tensor([[0.0], [1.0], [1.0]], dtype=torch.float32), - graph_w=torch.ones(3, 1), - ) - with pytest.raises(ValueError, match="Inconsistent graph_y"): - gather_bag_targets( - batch=batch, - bag_indices=[[0, 1], [2]], - target_columns=None, - ) - - -def test_region_buffer_flushes_in_hyperbatch_chunks(): - pytest.importorskip("lightning") - pytest.importorskip("torch_geometric") - from grass_mil.models.supervised_module import SupervisedModule - - module = SupervisedModule( - encoder={}, - graph_head={}, - attention={"input_dim": 1}, - loss={}, - optim={}, - task={ - "aggregation": "mil_attention", - "target_type": "binary", - "region_accumulation": {"enabled": True, "hyperbatch_size": 2}, - }, - ) - - class _Opt: - def __init__(self): - self.steps = 0 - self.zeroes = 0 - - def step(self): - self.steps += 1 - - def zero_grad(self, set_to_none=True): - self.zeroes += 1 - - class _Sched: - def __init__(self): - self.steps = 0 - - def step(self): - self.steps += 1 - - opt = _Opt() - sch = _Sched() - module.optimizers = lambda: opt # type: ignore[method-assign] - module.lr_schedulers = lambda: sch # type: ignore[method-assign] - module.manual_backward = lambda loss: None # type: ignore[method-assign] - module.log = lambda *args, **kwargs: None # type: ignore[method-assign] - - module._region_loss_buffer = [ - torch.tensor(1.0), - torch.tensor(2.0), - torch.tensor(3.0), - torch.tensor(4.0), - torch.tensor(5.0), - ] - flushes = module._flush_region_buffer_if_needed(force=False) - assert len(flushes) == 2 - assert len(module._region_loss_buffer) == 1 - assert opt.steps == 2 - assert sch.steps == 2 - - flushes = module._flush_region_buffer_if_needed(force=True) - assert len(flushes) == 1 - assert len(module._region_loss_buffer) == 0 - assert opt.steps == 3 - assert sch.steps == 3 diff --git a/tests/models/test_ssl.py b/tests/models/test_ssl.py deleted file mode 100644 index 25c851a..0000000 --- a/tests/models/test_ssl.py +++ /dev/null @@ -1,79 +0,0 @@ -import pytest -import torch -from torch import nn - - -class _TinyEncoder(nn.Module): - def __init__(self): - super().__init__() - self.lin = nn.Linear(8, 8) - - def forward(self, x): - return self.lin(x) - - -def test_bgrl_forward_and_update(): - from grass_mil.models.components.ssl import BGRL, MLPPredictor - - encoder = _TinyEncoder() - predictor = MLPPredictor(input_size=8, output_size=8, hidden_size=16) - model = BGRL(encoder=encoder, predictor=predictor) - - x1 = torch.randn(10, 8) - x2 = torch.randn(10, 8) - q, y = model(x1, x2) - assert q.shape == (10, 8) - assert y.shape == (10, 8) - - before = [p.clone() for p in model.target_encoder.parameters()] - model.update_target_network(momentum=0.9) - after = list(model.target_encoder.parameters()) - assert any(not torch.equal(b, a) for b, a in zip(before, after)) - - -@pytest.mark.parametrize("enabled", [True, False]) -def test_ssl_optional_factory(enabled): - from grass_mil.models.components.backbones import EncoderConfig, GNNEncoder - from grass_mil.models.components.factory import build_ssl - - encoder = GNNEncoder(EncoderConfig(input_dim=4, hidden_dim=8, out_dim=8, num_layers=1)) - ssl_model = build_ssl( - enabled, - encoder, - { - "method": "bgrl", - "predictor": {"hidden_size": 16}, - }, - ) - if enabled: - assert ssl_model is not None - else: - assert ssl_model is None - - -def test_ssl_factory_requires_hidden_size(): - from grass_mil.models.components.backbones import EncoderConfig, GNNEncoder - from grass_mil.models.components.factory import build_ssl - - encoder = GNNEncoder(EncoderConfig(input_dim=4, hidden_dim=8, out_dim=8, num_layers=1)) - with pytest.raises(ValueError, match="ssl.predictor.hidden_size is required"): - build_ssl(True, encoder, {"method": "bgrl", "predictor": {}}) - - -def test_ssl_factory_accepts_hydra_target_in_predictor_cfg(): - from grass_mil.models.components.backbones import EncoderConfig, GNNEncoder - from grass_mil.models.components.factory import build_ssl - - encoder = GNNEncoder(EncoderConfig(input_dim=4, hidden_dim=8, out_dim=8, num_layers=1)) - ssl_model = build_ssl( - True, - encoder, - { - "method": "bgrl", - "predictor": { - "_target_": "grass_mil.models.components.ssl.MLPPredictor", - "hidden_size": 16, - }, - }, - ) - assert ssl_model is not None diff --git a/tests/spatial_stats/__init__.py b/tests/spatial_stats/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests/spatial_stats/test_fidelity_phase8.py b/tests/spatial_stats/test_fidelity_phase8.py deleted file mode 100644 index 27513eb..0000000 --- a/tests/spatial_stats/test_fidelity_phase8.py +++ /dev/null @@ -1,171 +0,0 @@ -"""Phase 8: classical spatial statistics.""" - -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import pandas as pd -import pytest - -from tests.helpers.synthetic_cohort import cohort_overrides, compose_cfg, make_synthetic_cohort - -pytest.importorskip("scipy") - - -def _delaunay_edges(coords): - from scipy.spatial import Delaunay - - tri = Delaunay(coords) - e = set() - for s in tri.simplices: - for i in range(3): - a, b = int(s[i]), int(s[(i + 1) % 3]) - e.add((min(a, b), max(a, b))) - return np.array(sorted(e)).T - - -def test_nhood_zscores_random_vs_segregated(): - from grass_mil.spatial_stats import nhood_enrichment_zscores - - rng = np.random.default_rng(0) - coords = rng.uniform(0, 100, (400, 2)) - ei = _delaunay_edges(coords) - random_labels = rng.integers(0, 3, 400) - z = nhood_enrichment_zscores(random_labels, ei, n_groups=3, n_perms=100, seed=0) - assert z.shape == (3, 3) and np.allclose(z, z.T) and np.abs(z).max() < 3.5 - segregated = (coords[:, 0] > 50).astype(int) - z2 = nhood_enrichment_zscores(segregated, ei, n_groups=3, n_perms=100, seed=0) - assert z2[0, 0] > 3 and z2[1, 1] > 3 and z2[0, 1] < -3 - assert np.isnan(z2[2]).all() # absent group - - -def test_window_composition_and_cellular_neighbourhoods(): - from grass_mil.spatial_stats import cellular_neighbourhoods, window_composition - - rng = np.random.default_rng(1) - coords = rng.uniform(0, 100, (300, 2)) - labels = (coords[:, 0] > 50).astype(int) - W = window_composition(coords, labels, n_groups=2, k=10) - assert W.shape == (300, 2) and np.allclose(W.sum(1), 1.0) - assert W[labels == 0, 0].mean() > 0.8 - cn, centres = cellular_neighbourhoods([W, W[:50]], n_cn=2, sample_size=100, seed=0, n_init=2) - assert len(cn) == 2 and cn[0].shape == (300,) and centres.shape == (2, 2) - - -def test_crossk_csr_vs_clustered_and_contacts(): - from grass_mil.spatial_stats import contact_fractions, cross_k - - rng = np.random.default_rng(2) - coords = rng.uniform(0, 1000, (2000, 2)) - labels = rng.integers(0, 2, 2000) - k = cross_k( - coords, labels, group_names=["T", "X"], reference=["T"], partners=["X"], radii=[50, 100] - ) - assert abs(k["K_X_r50"] - 1.0) < 0.3 and abs(k["K_X_r100"] - 1.0) < 0.2 - # every X cell placed next to a T cell: strong attraction at short range - coords2 = coords.copy() - x_idx = np.where(labels == 1)[0] - t_idx = rng.choice(np.where(labels == 0)[0], x_idx.size) - coords2[x_idx] = coords[t_idx] + rng.normal(0, 3, (x_idx.size, 2)) - k2 = cross_k( - coords2, labels, group_names=["T", "X"], reference=["T"], partners=["X"], radii=[20] - ) - assert k2["K_X_r20"] > 1.5 - ei = _delaunay_edges(coords2) - c = contact_fractions( - labels, ei, group_names=["T", "X"], reference=["T"], partners=["X"], infiltrating=["X"] - ) - assert 0 <= c["contact_X"] <= 1 and 0 <= c["homotypic"] <= 1 and "infiltration" in c - - -def test_mann_whitney_table_and_bh(): - from grass_mil.spatial_stats import bh_adjust, mann_whitney_table, patient_mean - - rng = np.random.default_rng(3) - df = pd.DataFrame( - {"a": np.r_[rng.normal(1, 1, 20), rng.normal(0, 1, 20)], "b": rng.normal(size=40)}, - index=[f"r{i}" for i in range(40)], - ) - y = pd.Series(np.r_[np.ones(20), np.zeros(20)].astype(int), index=df.index) - T = mann_whitney_table(df, y) - assert list(T.columns) == [ - "metric", - "n_pos", - "n_neg", - "mean_pos", - "mean_neg", - "diff", - "p", - "q", - ] - assert T.set_index("metric").loc["a", "p"] < 0.05 - q = bh_adjust([0.01, 0.02, 0.5, np.nan]) - assert np.isnan(q[3]) and q[0] <= q[2] - pm = patient_mean(df, {f"r{i}": f"p{i // 2}" for i in range(40)}) - assert len(pm) == 20 - - -@pytest.fixture(scope="module") -def cache_dir(tmp_path_factory): - pytest.importorskip("torch_geometric") - import hydra - - from grass_mil.extract import run_extract - from grass_mil.train import train - - root = tmp_path_factory.mktemp("cohort8") - c = make_synthetic_cohort(root, n_patients=12, regions_per_patient=2, n_cells=80) - hydra.utils.instantiate( - compose_cfg("train", cohort_overrides(c, root / "pre")).data - ).prepare_data() - train_out = root / "train" - train(compose_cfg("train", cohort_overrides(c, train_out) + ["trainer.max_epochs=1"])) - ecfg = compose_cfg( - "extract", - cohort_overrides(c, root / "extract") - + [ - f"+extract.models.final={train_out / 'checkpoints' / 'last.ckpt'}", - "extract.device=cpu", - ], - ) - return c, Path(run_extract(ecfg)["models"][0]) - - -def test_spatial_stats_cli_finds_planted_niche(cache_dir, tmp_path): - from grass_mil.spatial_stats.cli import run - - c, model_dir = cache_dir - cfg = compose_cfg( - "spatial_stats", - [ - f"cells_dir={model_dir}", - f"output_dir={tmp_path / 'spatial'}", - f"labels.label_file={c.label_path}", - "labels.label_column=STAS", - "labels.patient_column=PATIENT", - "reference=[C]", - "partners=[D,A]", - "infiltrating=[D]", - "nhood.n_perms=50", - "cn.n_cn=3", - "cn.sample_size=500", - "min_per_group=3", - "crossk.radii=[30,60]", - ], - ) - run(cfg) - out = tmp_path / "spatial" - for name in ( - "nhood_tests.csv", - "cn_tests.csv", - "contact_tests.csv", - "crossk_tests.csv", - "report_spatial.md", - ): - assert (out / name).exists(), name - nh = pd.read_csv(out / "nhood_tests.csv").set_index("metric") - # the planted niche: D cells cluster together (and near C) only in positive regions - assert nh.loc["D|D", "diff"] > 0 and nh.loc["D|D", "p"] < 0.05 - ck = pd.read_csv(out / "crossk_tests.csv").set_index("metric") - assert ck.loc["K_D_r30", "mean_pos"] > ck.loc["K_D_r30", "mean_neg"] diff --git a/tests/test_configs.py b/tests/test_configs.py index d7041dc..4692d7c 100644 --- a/tests/test_configs.py +++ b/tests/test_configs.py @@ -1,37 +1,94 @@ -import hydra -from hydra.core.hydra_config import HydraConfig -from omegaconf import DictConfig +"""Configuration tree and packaging.""" +from __future__ import annotations -def test_train_config(cfg_train: DictConfig) -> None: - """Tests the training configuration provided by the `cfg_train` pytest fixture. +import importlib +import os +import re +import shutil +import subprocess +from importlib.metadata import entry_points +from pathlib import Path - :param cfg_train: A DictConfig containing a valid training configuration. - """ - assert cfg_train - assert cfg_train.data - assert cfg_train.model - assert cfg_train.trainer +import pytest - HydraConfig().set_config(cfg_train) +from tests.helpers.cohort import compose_cfg - hydra.utils.instantiate(cfg_train.data) - hydra.utils.instantiate(cfg_train.model) - hydra.utils.instantiate(cfg_train.trainer) +CONFIG_DIR = Path(__file__).parents[1] / "src" / "grass_mil" / "configs" +REGIMES = ["finetune", "finetune_reg", "attention_only", "mlp_head", "mlp_head_frozen", "pretrain"] +COMMANDS = { + "grass-mil-train": "grass_mil.train:main", + "grass-mil-nested-cv": "grass_mil.nested_cv:main", + "grass-mil-predict": "grass_mil.predict:main", + "grass-mil-extract": "grass_mil.extract:main", + "grass-mil-niches": "grass_mil.niches.cli:main", + "grass-mil-metaclusters": "grass_mil.niches.replication.cli:main", + "grass-mil-spatial-stats": "grass_mil.spatial.cli:main", +} -def test_eval_config(cfg_eval: DictConfig) -> None: - """Tests the evaluation configuration provided by the `cfg_eval` pytest fixture. +def test_finetune_is_the_default_regime(): + cfg = compose_cfg("train", []) + assert cfg.regime_name == "finetune" + assert cfg.model._target_ == "grass_mil.models.SupervisedModule" + assert cfg.finetune.head_init == "probe" and cfg.model.head.num_layers == 1 + assert cfg.finetune.lr_scales == {"encoder": 1.0, "head": 0.01, "attention": 2.0} - :param cfg_train: A DictConfig containing a valid evaluation configuration. - """ - assert cfg_eval - assert cfg_eval.data - assert cfg_eval.model - assert cfg_eval.trainer - HydraConfig().set_config(cfg_eval) +@pytest.mark.parametrize("regime", REGIMES) +def test_every_regime_composes(regime): + cfg = compose_cfg("train", [f"regime={regime}"]) + assert cfg.regime_name == regime + if regime == "pretrain": + assert cfg.model._target_ == "grass_mil.models.BGRLModule" + assert cfg.data.sampler.merge_units and not cfg.data.labelled_only + else: + assert cfg.callbacks.model_checkpoint.monitor == "val/auroc" - hydra.utils.instantiate(cfg_eval.data) - hydra.utils.instantiate(cfg_eval.model) - hydra.utils.instantiate(cfg_eval.trainer) + +@pytest.mark.parametrize( + "task,metric", + [("classification", "val/auroc"), ("regression", "val/rmse"), ("survival", "val/c_index")], +) +def test_selection_follows_the_task(task, metric): + cfg = compose_cfg("train", [f"task={task}", "task.target_columns=[a,b]"]) + assert cfg.callbacks.early_stopping.monitor == metric + + +@pytest.mark.parametrize( + "name", ["train", "nested_cv", "predict", "extract", "niches", "metaclusters", "spatial_stats"] +) +def test_every_root_config_composes(name): + compose_cfg(name, []) + + +def test_entry_points_match_the_packaged_commands(): + found = { + ep.name: ep.value + for ep in entry_points(group="console_scripts") + if ep.name.startswith("grass-mil") + } + assert found == COMMANDS + for target in COMMANDS.values(): + module, func = target.split(":") + assert callable(getattr(importlib.import_module(module), func)) + + +def test_cli_runs_from_another_directory(tmp_path): + command = shutil.which("grass-mil-train") + if command is None: + pytest.skip("the package is not installed") + # mkl-service sets MKL_THREADING_LAYER in this process; a fresh one picks its own. + env = {k: v for k, v in os.environ.items() if k != "MKL_THREADING_LAYER"} + result = subprocess.run( + [command, "--help"], cwd=tmp_path, capture_output=True, text=True, env=env + ) + assert result.returncode == 0 and "regime" in result.stdout + + +def test_configs_describe_the_method_not_its_history(): + forbidden = re.compile( + r"preprint|space-?gm|\ba3\b|probehead|ftall|fidelity|legacy", re.IGNORECASE + ) + for path in CONFIG_DIR.rglob("*.yaml"): + assert not forbidden.search(path.read_text()), path diff --git a/tests/test_data.py b/tests/test_data.py new file mode 100644 index 0000000..98de7e0 --- /dev/null +++ b/tests/test_data.py @@ -0,0 +1,245 @@ +"""Inputs, graphs, units, labels, splits and the datamodule.""" + +from __future__ import annotations + +import json + +import numpy as np +import pandas as pd +import pytest +import torch + +from grass_mil.data import LabelTable, SpatialOmicsDataModule +from grass_mil.data.graphs import GraphConfig, delaunay_edges, edge_features, log_distance_feature +from grass_mil.data.labels import class_index +from grass_mil.data.loaders import InputConfig, input_type_of, load_cell_table +from grass_mil.data.precompute import ProcessedIndex, read_manifest +from grass_mil.data.splits import inner_folds, load_folds, random_split +from grass_mil.data.units import TilingConfig, grid_tiles + + +def _datamodule(cohort, tmp_path, **kwargs): + defaults = dict( + manifest_path=str(cohort.manifest), + processed_dir=str(tmp_path / "processed"), + input={"cell_id_column": "cell_id", "features": ["area"]}, + graph={"neighbor_cutoff": 55}, + labels={"file": str(cohort.labels), "columns": ["label"], "patient_column": "patient_id"}, + sampler={"depth": 2, "roots_per_unit": 16, "batch_size": 8}, + ) + defaults.update(kwargs) + return SpatialOmicsDataModule(**defaults) + + +# ---------------------------------------------------------------- inputs +def test_input_type_from_suffix_or_manifest(): + assert input_type_of("a/b.csv") == "csv" + assert input_type_of("a/b.h5ad") == "h5ad" + assert input_type_of("a/b.rds") == "sce" + assert input_type_of("a/b.txt", "csv") == "csv" + with pytest.raises(ValueError): + input_type_of("a/b.parquet") + + +def test_csv_features_all_excludes_ids_coordinates_and_cell_type(cohort): + path = cohort.manifest.parent / f"{cohort.regions[0]}.csv" + table = load_cell_table(path, "csv", InputConfig.from_dict({"cell_id_column": "cell_id"})) + assert table.feature_names == ["area"] + assert table.cell_types is not None and set(table.cell_types) <= {"A", "B", "C", "D"} + assert table.cell_ids[0] == "c0" + + +def test_h5ad_input(tmp_path): + anndata = pytest.importorskip("anndata") + rng = np.random.default_rng(0) + adata = anndata.AnnData( + X=rng.random((30, 3)).astype(np.float32), + obs=pd.DataFrame( + {"cell_type": rng.choice(["a", "b"], 30)}, index=[f"c{i}" for i in range(30)] + ), + ) + adata.var_names = ["g1", "g2", "g3"] + adata.obsm["spatial"] = rng.random((30, 2)) * 100 + adata.write_h5ad(tmp_path / "s.h5ad") + table = load_cell_table( + tmp_path / "s.h5ad", "h5ad", InputConfig.from_dict({"features": ["g2"]}) + ) + assert table.features.shape == (30, 1) and table.feature_names == ["g2"] + assert table.coords.shape == (30, 2) + + +def test_manifest_ids_stay_strings_and_region_defaults_to_sample(tmp_path): + (tmp_path / "m.csv").write_text("sample_id,input_path,region_id\n001,a.csv,\n002,b.csv,7\n") + df = read_manifest(tmp_path / "m.csv") + assert list(df["sample_id"]) == ["001", "002"] + assert list(df["region_id"]) == ["001", "7"] + assert df["input_path"].iloc[0] == str((tmp_path / "a.csv").resolve()) + + +def test_manifest_rejects_repeated_regions(tmp_path): + (tmp_path / "m.csv").write_text("sample_id,input_path,region_id\na,a.csv,r\nb,b.csv,r\n") + with pytest.raises(ValueError, match="unique"): + read_manifest(tmp_path / "m.csv") + + +# ---------------------------------------------------------------- graphs +def test_delaunay_edges_are_symmetric_and_unique(): + coords = np.random.default_rng(0).random((50, 2)) + e = delaunay_edges(coords).numpy() + pairs = set(map(tuple, e.T)) + assert len(pairs) == e.shape[1] + assert all((b, a) in pairs for a, b in pairs) + from scipy.spatial import Delaunay + + expected = set() + for simplex in Delaunay(coords).simplices: + for i in range(3): + for j in range(3): + if i != j: + expected.add((int(simplex[i]), int(simplex[j]))) + assert pairs == expected + + +def test_edge_feature_columns(): + coords = np.array([[0.0, 0.0], [10.0, 0.0], [0.0, 200.0]]) + e = delaunay_edges(coords) + cfg = GraphConfig.from_dict( + {"edge_features": ["log_distance", "neighbor", "distance"], "neighbor_cutoff": 15} + ) + attr = edge_features(coords, e, cfg).numpy() + lengths = attr[:, 2] + assert np.allclose(attr[:, 1], lengths <= 15) + assert np.allclose(attr[:, 0], log_distance_feature(lengths, 2.0, 5.0)) + assert log_distance_feature(np.array([np.exp(2.0), np.exp(5.0), 1e6]), 2.0, 5.0).tolist() == [ + 0.0, + 1.0, + 1.0, + ] + + +def test_grid_tiles_cover_every_cell_once(): + coords = np.random.default_rng(0).random((200, 2)) * 500 + tiles = grid_tiles(coords, TilingConfig(size=100, stride=100)) + covered = np.concatenate(tiles) + assert sorted(covered.tolist()) == list(range(200)) + + +# ---------------------------------------------------------------- labels +def test_label_table_encodes_text_labels_and_patients(tmp_path): + (tmp_path / "l.csv").write_text("region_id,grade,patient_id\nr1,high,p1\nr2,low,p1\nr3,,p2\n") + table = LabelTable.read(tmp_path / "l.csv", id_column="region_id", patient_column="patient_id") + assert table.codes == {"grade": ["high", "low"]} + assert table.target("r2").tolist() == [1.0] + assert "r3" not in table # unlabelled + assert table.patient("r1") == "p1" + assert table.resolve("x", "r2") == "r2" + + +def test_class_index_rounds_halves_up(): + assert [class_index(v) for v in (0.0, 0.49, 0.5, 1.0, 2.0)] == [0, 0, 1, 1, 2] + + +# ---------------------------------------------------------------- splits +def test_random_split_keeps_groups_together_and_stratifies(): + groups = {f"u{i}": f"g{i // 2}" for i in range(20)} + strata = {f"g{j}": str(j % 2) for j in range(10)} + split = random_split(groups, (0.6, 0.2, 0.2), seed=1, strata=strata) + for g in set(groups.values()): + assert len({split[u] for u in groups if groups[u] == g}) == 1 + for name in ("train", "val", "test"): + classes = {strata[groups[u]] for u in split if split[u] == name} + assert classes == {"0", "1"} + + +def test_inner_folds_are_group_aware_and_stratified(): + ids = [f"u{i}" for i in range(24)] + groups = {u: f"p{i // 2}" for i, u in enumerate(ids)} + classes = {u: (i // 2) % 2 for i, u in enumerate(ids)} + folds = inner_folds(ids, groups=groups, classes=classes, k=3, seed=0) + assert len(folds) == 3 + for train, val in folds: + assert not {groups[u] for u in train} & {groups[u] for u in val} + assert {classes[u] for u in val} == {0, 1} + assert inner_folds(ids, groups=groups, classes=None, k=3, seed=0) + + +def test_load_folds_rejects_repeated_ids(tmp_path): + (tmp_path / "f.json").write_text(json.dumps([["a", "b"], ["b"]])) + with pytest.raises(ValueError): + load_folds(tmp_path / "f.json") + + +# ------------------------------------------------------------ datamodule +def test_datamodule_builds_reuses_and_refuses_stale_graphs(cohort, tmp_path): + dm = _datamodule(cohort, tmp_path) + dm.prepare_data() + dm.setup() + assert dm.input_spec.feature_dim == 1 + assert dm.input_spec.cell_types == ("A", "B", "C", "D") + assert dm.input_spec.edge_dim == 2 + sizes = [len(dm.dataset_train), len(dm.dataset_val), len(dm.dataset_test)] + assert sum(sizes) == len(cohort.regions) and min(sizes) > 0 + data = dm.dataset_train[0] + assert data.cell_type.dtype == torch.long and data.graph_y.shape == (1, 1) + _datamodule(cohort, tmp_path).prepare_data() # same settings: reused + with pytest.raises(ValueError, match="different settings"): + _datamodule(cohort, tmp_path, input={"features": None}).prepare_data() + + +def test_explicit_split_by_region_or_patientless_ids(cohort, tmp_path): + dm = _datamodule( + cohort, + tmp_path, + split={"train": cohort.regions[:6], "val": cohort.regions[6:8], "test": []}, + ) + dm.prepare_data() + dm.setup() + assert [e.unit_id for e in dm.dataset_train.entries] == cohort.regions[:6] + assert len(dm.dataset_test) == 0 + with pytest.raises(ValueError, match="no processed unit"): + bad = _datamodule(cohort, tmp_path, split={"train": ["nope"]}) + bad.setup() + + +def test_tile_units_inherit_region_labels(cohort, tmp_path): + dm = _datamodule( + cohort, tmp_path, unit="tile", tiling={"size": 300, "stride": 300}, min_cells=5 + ) + dm.prepare_data() + dm.setup() + entries = dm.all_units().entries + assert all("_tile" in e.unit_id for e in entries) + data = dm.all_units()[0] + assert data.graph_y.item() == cohort.label_of[entries[0].region_id] + + +def test_pca_reducer_end_to_end_and_reuse_of_its_basis(cohort, tmp_path): + reducer = {"name": "pca", "n_components": 1} + dm = _datamodule(cohort, tmp_path, feature_reducer=reducer) + dm.prepare_data() + index = ProcessedIndex.read(tmp_path / "processed") + assert index.feature_names == ["pc1"] + assert len(index.feature_reducer["state"]["mean"]) == 1 + other = _datamodule( + cohort, + tmp_path / "other", + feature_reducer={ + **reducer, + "state_from": str(tmp_path / "processed" / "processed_index.json"), + }, + ) + other.prepare_data() + other.setup() + a = dm.all_units()[0].x + b = other.all_units()[0].x + assert torch.allclose(a, b) + + +def test_train_fitted_reducer_is_refused_for_explicit_splits(cohort, tmp_path): + with pytest.raises(ValueError, match="fit_on=train"): + _datamodule( + cohort, + tmp_path, + feature_reducer={"name": "pca", "fit_on": "train"}, + split={"train": cohort.regions}, + ) diff --git a/tests/test_e2e_synthetic.py b/tests/test_e2e_synthetic.py deleted file mode 100644 index 6216c1a..0000000 --- a/tests/test_e2e_synthetic.py +++ /dev/null @@ -1,227 +0,0 @@ -"""End-to-end smoke test of the public analysis chain on the synthetic cohort. - -precompute -> BGRL A3 preset (a few steps) -> nested CV (regime=mil and the -probe regime) -> extract -> niches -> metaclusters -> spatial statistics -> -scores. Asserts that every artefact exists, AUROCs are finite, the attribution -identity holds and the planted niche (D cells packed around C cells in positive -regions) is the top enriched cluster of at least one fold model. -""" - -from __future__ import annotations - -import json -from pathlib import Path - -import numpy as np -import pandas as pd -import pytest - -from tests.helpers.synthetic_cohort import cohort_overrides, compose_cfg, make_synthetic_cohort - -pytest.importorskip("torch_geometric") -pytest.importorskip("torch_sparse") - - -@pytest.mark.slow -def test_end_to_end_synthetic_chain(tmp_path): - import hydra - - from grass_mil.evaluation.cli import run_scores - from grass_mil.extract import run_extract - from grass_mil.interpretability.metaclusters_cli import run_metaclusters - from grass_mil.interpretability.niches_cli import run_niches - from grass_mil.nested_cv import run_nested_cv - from grass_mil.spatial_stats.cli import run as run_spatial - from grass_mil.train import train - - root = tmp_path / "cohort" - c = make_synthetic_cohort( - root, n_patients=10, regions_per_patient=2, n_cells=120, n_outer_folds=2 - ) - hydra.utils.instantiate( - compose_cfg("train", cohort_overrides(c, root / "pre")).data - ).prepare_data() - - # 1. pretraining preset, a handful of steps - pre_out = root / "pretrain" - base = [ - o - for o in cohort_overrides(c, pre_out) - if not o.startswith(("regime=", "data.val_sampler", "model.attention", "task.")) - ] - train( - compose_cfg( - "train", - base - + [ - "experiment=grass_mil_pretrain_a3", - "task.total_steps=4", - "task.warmup_steps=1", - "task.lr_warmup.warmup_steps=1", - "task.augmentation.feature_noise_columns=[area]", - "callbacks.model_checkpoint.every_n_train_steps=2", - "trainer.precision=32", - "data.sampler.runtime.depth=2", - "data.sampler.runtime.subgraph_batch_size=16", - ], - ) - ) - bgrl = pre_out / "checkpoints" / "last.ckpt" - assert bgrl.exists() - - # 2. nested CV, default regime and the probe regime - summaries = {} - for regime in ("mil", "linear_probehead"): - out = root / f"nested_{regime}" - cfg = compose_cfg( - "nested_cv", - cohort_overrides(c, out) - + [ - f"regime={regime}", - f"model.init_from_ckpt={bgrl}", - f"nested_cv.folds_file={c.folds_path}", - "nested_cv.inner_k=2", - "nested_cv.probe.roots_per_unit=16", - "trainer.max_epochs=2", - ], - ) - summaries[regime] = run_nested_cv(cfg) - assert summaries[regime]["complete"] - assert all(np.isfinite(a) for a in summaries[regime]["auc_folds"]) - assert (out / "oof_predictions.csv").exists() and ( - out / "oof_patient_predictions.csv" - ).exists() - - # 3. scores on the OOF table - scores = run_scores( - compose_cfg( - "scores", - [ - f"input={root / 'nested_mil' / 'oof_predictions.csv'}", - f"output_dir={root / 'scores'}", - "n_boot=100", - ], - ) - ) - assert scores["n"] == 10 and np.isfinite(scores["scores"][0]["auroc"]) - - # 4. extract both fold refits - ext = run_extract( - compose_cfg( - "extract", - cohort_overrides(c, root / "extract") - + [ - f"extract.nested_cv_summary={root / 'nested_mil' / 'nested_cv_summary.json'}", - "extract.device=cpu", - ], - ) - ) - cells_dir = Path(ext["output_dir"]) - assert sorted(p.name for p in cells_dir.iterdir()) == ["fold0", "fold1"] - # a probe regime extracts from its refit checkpoints without a probe file - ext_lin = run_extract( - compose_cfg( - "extract", - cohort_overrides(c, root / "extract_lin") - + [ - "regime=linear_probehead", - f"extract.nested_cv_summary={root / 'nested_linear_probehead' / 'nested_cv_summary.json'}", - "extract.device=cpu", - ], - ) - ) - assert len(ext_lin["models"]) == 2 - - # 5. niches of one model (default backend with a small min_cluster_size) and a k-means override. - # The chain is a smoke test: a two-epoch model on 20 tiny regions is not expected to rank - # the planted niche first, so the checks here are structural; the planted niche is - # detected model-free by the spatial statistics in step 7. - for fold in ("fold0", "fold1"): - res = run_niches( - compose_cfg( - "niches", - [ - f"cells_dir={cells_dir / fold}", - f"output_dir={root / 'niches' / fold}", - "clustering.params.min_cluster_size=20", - "sample_size=null", - "n_boot=50", - "umap.enabled=false", - ], - ) - ) - assert res["n_clusters"] >= 1 - table = pd.read_csv(root / "niches" / fold / "clusters.csv") - cells = pd.read_csv(root / "niches" / fold / "cells.csv") - assert table["share"].sum() == pytest.approx(1.0) - assert ( - np.isfinite(table["relattn"]).all() - and (table["relattn_lo"] <= table["relattn_hi"]).all() - ) - assert set(cells["cluster"]) == set(table["cluster_id"]) - held_out = json.loads((cells_dir / fold / "index.json").read_text())["held_out_regions"] - assert set(cells["region_id"]) == set(held_out) - km = run_niches( - compose_cfg( - "niches", - [ - f"cells_dir={cells_dir / 'fold0'}", - f"output_dir={root / 'niches_km'}", - "clustering.method=kmeans", - "clustering.n_clusters=4", - "sample_size=null", - "n_boot=20", - "umap.enabled=false", - ], - ) - ) - assert km["n_clusters"] == 4 - - # 6. metaclusters, fixed threshold - meta = run_metaclusters( - compose_cfg( - "metaclusters", - [ - f"cells_dir={cells_dir}", - f"output_dir={root / 'meta'}", - f"labels.label_file={c.label_path}", - "labels.label_column=STAS", - "labels.patient_column=PATIENT", - "clustering.params.min_cluster_size=20", - "sample_size=null", - "n_boot=50", - "match.threshold=0.4", - "match.min_folds=2", - "transfer.min_cells=10", - "basis.k=4", - "basis.sample_size=2000", - "basis.n_init=2", - ], - ) - ) - assert ( - meta["match"]["n_metaclusters"] >= 1 - and (root / "meta" / "basis" / "basis_attribution.csv").exists() - ) - - # 7. spatial statistics - sp = run_spatial( - compose_cfg( - "spatial_stats", - [ - f"cells_dir={cells_dir / 'fold0'}", - f"output_dir={root / 'spatial'}", - f"labels.label_file={c.label_path}", - "labels.label_column=STAS", - "labels.patient_column=PATIENT", - "reference=[C]", - "partners=[D]", - "nhood.n_perms=30", - "cn.n_cn=3", - "cn.sample_size=500", - "min_per_group=3", - ], - ) - ) - nh = pd.read_csv(sp["nhood_tests"]).set_index("metric") - assert nh.loc["D|D", "diff"] > 0 diff --git a/tests/test_end_to_end.py b/tests/test_end_to_end.py new file mode 100644 index 0000000..5f18d7c --- /dev/null +++ b/tests/test_end_to_end.py @@ -0,0 +1,77 @@ +"""The documented workflow, through the installed commands, on synthetic data.""" + +from __future__ import annotations + +import os +import shutil +import subprocess + +import pandas as pd +import pytest + +from tests.helpers.cohort import cohort_overrides + +pytestmark = pytest.mark.slow + + +def _run(command: str, *args: str, cwd) -> None: + exe = shutil.which(command) + if exe is None: + pytest.skip("the package is not installed") + env = {k: v for k, v in os.environ.items() if k != "MKL_THREADING_LAYER"} + result = subprocess.run([exe, *args], cwd=cwd, capture_output=True, text=True, env=env) + assert result.returncode == 0, result.stderr[-3000:] + + +def test_workflow(cohort, tmp_path): + data = cohort_overrides(cohort, processed="processed_e2e", finetune=False) + small = ["finetune.hyperbatch_size=3", "finetune.probe.roots_per_unit=16"] + pre, cv, cells = tmp_path / "pretrain", tmp_path / "cv", tmp_path / "cells" + _run( + "grass-mil-train", + *data, + "regime=pretrain", + "pretrain.total_steps=3", + "pretrain.warmup_steps=1", + f"hydra.run.dir={pre}", + cwd=tmp_path, + ) + ckpt = pre / "checkpoints" / "last.ckpt" + _run( + "grass-mil-nested-cv", + *data, + *small, + f"model.init_from_ckpt={ckpt}", + f"nested_cv.folds_file={cohort.folds}", + "nested_cv.inner_k=2", + "trainer.max_epochs=1", + f"hydra.run.dir={cv}", + cwd=tmp_path, + ) + _run( + "grass-mil-extract", + *data, + f"extract.nested_cv_summary={cv / 'nested_cv_summary.json'}", + f"extract.output_dir={cells}", + f"hydra.run.dir={tmp_path / 'x'}", + cwd=tmp_path, + ) + _run( + "grass-mil-niches", + f"cells_dir={cells / 'fold1'}", + "clustering.min_cluster_size=20", + "descriptors.min_cells=10", + "descriptors.n_perms=5", + "n_boot=50", + f"hydra.run.dir={tmp_path / 'niches'}", + cwd=tmp_path, + ) + _run( + "grass-mil-predict", + *data, + f"ckpt_path={cv / 'outer_1' / 'refit' / 'checkpoints' / 'last.ckpt'}", + f"hydra.run.dir={tmp_path / 'predict'}", + cwd=tmp_path, + ) + assert len(pd.read_csv(tmp_path / "niches" / "niches.csv")) >= 1 + assert len(pd.read_csv(tmp_path / "predict" / "predictions.csv")) == len(cohort.regions) diff --git a/tests/test_eval.py b/tests/test_eval.py deleted file mode 100644 index f1c2d9d..0000000 --- a/tests/test_eval.py +++ /dev/null @@ -1,47 +0,0 @@ -import os -from pathlib import Path - -import pytest -import torch -from hydra.core.hydra_config import HydraConfig -from omegaconf import DictConfig, open_dict -from grass_mil.eval import evaluate -from grass_mil.train import train - - -def test_train_eval(tmp_path: Path, cfg_train: DictConfig, cfg_eval: DictConfig) -> None: - """Tests training and evaluation by training for 1 epoch with `train.py` then evaluating with - `eval.py`. - - :param tmp_path: The temporary logging path. - :param cfg_train: A DictConfig containing a valid training configuration. - :param cfg_eval: A DictConfig containing a valid evaluation configuration. - """ - assert str(tmp_path) == cfg_train.paths.output_dir == cfg_eval.paths.output_dir - - with open_dict(cfg_train): - cfg_train.trainer.max_epochs = 1 - cfg_train.test = True - - HydraConfig().set_config(cfg_train) - train_metric_dict, _ = train(cfg_train) - - assert "last.ckpt" in os.listdir(tmp_path / "checkpoints") - - with open_dict(cfg_eval): - cfg_eval.ckpt_path = str(tmp_path / "checkpoints" / "last.ckpt") - - HydraConfig().set_config(cfg_eval) - test_metric_dict, _ = evaluate(cfg_eval) - - assert "test/loss" in test_metric_dict - assert "test/loss" in train_metric_dict - assert torch.isfinite(test_metric_dict["test/loss"]) - assert torch.isfinite(train_metric_dict["test/loss"]) - - -def test_evaluate_raises_when_ckpt_path_missing(cfg_eval: DictConfig) -> None: - with open_dict(cfg_eval): - cfg_eval.ckpt_path = "" - with pytest.raises(ValueError, match="Missing required `ckpt_path` for evaluation."): - evaluate(cfg_eval) diff --git a/tests/test_inference.py b/tests/test_inference.py new file mode 100644 index 0000000..e2d5cb5 --- /dev/null +++ b/tests/test_inference.py @@ -0,0 +1,54 @@ +"""Prediction with a checkpoint and per-cell extraction.""" + +from __future__ import annotations + +import json + +import numpy as np +import pandas as pd +import pytest + +from grass_mil.inference.cell_cache import CellCache +from grass_mil.predict import run_predict +from tests.helpers.cohort import cohort_overrides, compose_cfg + + +@pytest.mark.parametrize("aggregation", ["mean", "merged_bag"]) +def test_predict_needs_only_the_checkpoint_and_reports_auroc( + cohort, nested_cv, tmp_path, aggregation +): + ckpt = nested_cv / "outer_0" / "refit" / "checkpoints" / "last.ckpt" + cfg = compose_cfg( + "predict", + cohort_overrides(cohort) + + [f"ckpt_path={ckpt}", f"predict.patient_aggregation={aggregation}"], + tmp_path, + ) + summary = run_predict(cfg) + assert np.isfinite(summary["unit_metrics"]["auroc"]) + assert "auroc" in summary["patient_metrics"] + preds = pd.read_csv(tmp_path / "predictions.csv") + assert sorted(preds["unit_id"]) == sorted(cohort.regions) + assert json.loads((tmp_path / "metrics.json").read_text())["n_units"] == len(cohort.regions) + assert len(pd.read_csv(tmp_path / "patient_predictions.csv")) == 6 + + +def test_extract_writes_every_cell_of_every_unit(cohort, cells_dir): + cache = CellCache.open(cells_dir / "fold0") + assert sorted(cache.units) == sorted(cohort.regions) + assert cache.held_out and set(cache.held_out) < set(cache.units) + data = cache.load(cache.units[0]) + n = data["pos"].shape[0] + for key in ( + "embedding", + "attention_space", + "instance_logits", + "attention_logits", + "cell_type", + "cell_ids", + ): + assert data[key].shape[0] == n + assert data["attention_space"].shape[1] == 8 + info = cache.unit_info(cache.units[0]) + assert info["patient_id"] == cohort.patient_of[cache.units[0]] + assert info["targets"] == [float(cohort.label_of[cache.units[0]])] diff --git a/tests/test_loocv.py b/tests/test_loocv.py deleted file mode 100644 index 86cf868..0000000 --- a/tests/test_loocv.py +++ /dev/null @@ -1,223 +0,0 @@ -from pathlib import Path - -import pytest -from hydra import compose, initialize_config_module -from hydra.core.global_hydra import GlobalHydra -from omegaconf import DictConfig, open_dict - -from grass_mil.loocv import _fold_slug, _resolve_selected_folds, _to_float_metrics, run_loocv - - -def _make_loocv_cfg(tmp_path: Path) -> DictConfig: - with initialize_config_module(version_base="1.3", config_module="grass_mil.configs"): - cfg = compose(config_name="loocv.yaml") - with open_dict(cfg): - cfg.paths.root_dir = str(tmp_path) - cfg.paths.output_dir = str(tmp_path / "outputs") - cfg.data.processed_dir = str(tmp_path / "processed") - cfg.data.raw_manifest_path = str(tmp_path / "manifest.csv") - cfg.data.split.loocv.enabled = True - cfg.data.split.loocv.fold_unit = "sample" - cfg.data.split.loocv.holdout_id = "s0" - cfg.data.split.loocv.validation_strategy = "heldout_fold_items" - return cfg - - -def _write_manifest(path: Path) -> None: - path.write_text( - "sample_id,input_path,input_type,region_id\n" - "s0,/tmp/f0.csv,csv,r0\n" - "s1,/tmp/f1.csv,csv,r1\n" - ) - - -def _write_manifest_for_sample_ids(path: Path, sample_ids: list[str]) -> None: - rows = ["sample_id,input_path,input_type,region_id"] - rows.extend( - f"{sample_id},/tmp/{idx}.csv,csv,r{idx}" for idx, sample_id in enumerate(sample_ids) - ) - path.write_text("\n".join(rows) + "\n") - - -def test_fold_slug_is_collision_safe_for_sanitized_aliases() -> None: - fold_a = "sampleA::region1" - fold_b = "sampleA//region1" - - slug_a = _fold_slug(fold_a) - slug_b = _fold_slug(fold_b) - - assert slug_a.startswith("sampleA_region1_") - assert slug_b.startswith("sampleA_region1_") - assert slug_a != slug_b - assert slug_a == _fold_slug(fold_a) - - -@pytest.mark.parametrize("mode", ["single", "all"]) -def test_loocv_rejects_unsafe_isolation_flags(tmp_path: Path, mode: str) -> None: - cfg = _make_loocv_cfg(tmp_path) - _write_manifest(Path(cfg.data.raw_manifest_path)) - with open_dict(cfg): - cfg.loocv.mode = mode - cfg.loocv.per_fold_processed_dir = False - cfg.loocv.force_precompute_per_fold = False - cfg.loocv.fold_index = 0 - - with pytest.raises(ValueError, match="Unsafe LOOCV configuration"): - run_loocv(cfg) - GlobalHydra.instance().clear() - - -def test_loocv_rejects_unsafe_isolation_flags_when_base_loocv_disabled( - tmp_path: Path, -) -> None: - cfg = _make_loocv_cfg(tmp_path) - _write_manifest(Path(cfg.data.raw_manifest_path)) - with open_dict(cfg): - cfg.data.split.loocv.enabled = False - cfg.loocv.mode = "single" - cfg.loocv.per_fold_processed_dir = False - cfg.loocv.force_precompute_per_fold = False - cfg.loocv.fold_index = 0 - - with pytest.raises(ValueError, match="Unsafe LOOCV configuration"): - run_loocv(cfg) - GlobalHydra.instance().clear() - - -def test_loocv_runs_with_safe_isolation_flags(tmp_path: Path, monkeypatch) -> None: - cfg = _make_loocv_cfg(tmp_path) - _write_manifest(Path(cfg.data.raw_manifest_path)) - with open_dict(cfg): - cfg.loocv.mode = "all" - cfg.loocv.per_fold_processed_dir = True - cfg.loocv.force_precompute_per_fold = False - cfg.loocv.summary_filename = "summary.json" - - def _fake_train(_cfg): - return {"val/loss": 1.23}, {} - - monkeypatch.setattr("grass_mil.loocv._discover_viable_fold_ids", lambda _cfg: ["s0", "s1"]) - monkeypatch.setattr("grass_mil.loocv.train", _fake_train) - summary = run_loocv(cfg) - summary_path = Path(cfg.paths.output_dir) / "summary.json" - - assert summary_path.exists() - assert len(summary["folds"]) == 2 - assert "aggregate_metrics" in summary - GlobalHydra.instance().clear() - - -def test_loocv_uses_unique_dirs_for_colliding_sanitized_fold_ids( - tmp_path: Path, monkeypatch -) -> None: - cfg = _make_loocv_cfg(tmp_path) - _write_manifest_for_sample_ids( - Path(cfg.data.raw_manifest_path), ["sampleA::region1", "sampleA//region1"] - ) - with open_dict(cfg): - cfg.loocv.mode = "all" - cfg.loocv.per_fold_processed_dir = True - cfg.loocv.force_precompute_per_fold = False - - seen_output_dirs: list[str] = [] - seen_processed_dirs: list[str] = [] - - def _fake_train(_cfg): - seen_output_dirs.append(str(_cfg.paths.output_dir)) - seen_processed_dirs.append(str(_cfg.data.processed_dir)) - return {"val/loss": 0.5}, {} - - monkeypatch.setattr( - "grass_mil.loocv._discover_viable_fold_ids", - lambda _cfg: ["sampleA::region1", "sampleA//region1"], - ) - monkeypatch.setattr("grass_mil.loocv.train", _fake_train) - summary = run_loocv(cfg) - - assert len(summary["folds"]) == 2 - assert len(set(seen_output_dirs)) == 2 - assert len(set(seen_processed_dirs)) == 2 - GlobalHydra.instance().clear() - - -def test_loocv_fails_fast_when_selected_fold_not_viable(tmp_path: Path, monkeypatch) -> None: - cfg = _make_loocv_cfg(tmp_path) - _write_manifest(Path(cfg.data.raw_manifest_path)) - with open_dict(cfg): - cfg.loocv.mode = "all" - cfg.loocv.per_fold_processed_dir = True - cfg.loocv.force_precompute_per_fold = False - - monkeypatch.setattr("grass_mil.loocv._discover_viable_fold_ids", lambda _cfg: ["s0"]) - monkeypatch.setattr( - "grass_mil.loocv.train", - lambda _cfg: pytest.fail("train() should not be called when preflight fails"), - ) - - with pytest.raises( - ValueError, - match="LOOCV preflight failed: selected folds are not viable after preprocessing filters", - ): - run_loocv(cfg) - GlobalHydra.instance().clear() - - -def test_loocv_preflight_passes_and_runs_all_selected_folds(tmp_path: Path, monkeypatch) -> None: - cfg = _make_loocv_cfg(tmp_path) - _write_manifest_for_sample_ids(Path(cfg.data.raw_manifest_path), ["s0", "s1", "s2"]) - with open_dict(cfg): - cfg.loocv.mode = "all" - cfg.loocv.per_fold_processed_dir = True - cfg.loocv.force_precompute_per_fold = False - - call_count = {"n": 0} - - def _fake_train(_cfg): - call_count["n"] += 1 - return {"val/loss": 0.75}, {} - - monkeypatch.setattr( - "grass_mil.loocv._discover_viable_fold_ids", lambda _cfg: ["s0", "s1", "s2"] - ) - monkeypatch.setattr("grass_mil.loocv.train", _fake_train) - - summary = run_loocv(cfg) - - assert call_count["n"] == 3 - assert len(summary["folds"]) == 3 - GlobalHydra.instance().clear() - - -def test_single_mode_warns_when_both_holdout_id_and_fold_index_are_set( - tmp_path: Path, -) -> None: - cfg = _make_loocv_cfg(tmp_path) - with open_dict(cfg): - cfg.loocv.mode = "single" - cfg.data.split.loocv.holdout_id = "s0" - cfg.loocv.fold_index = 1 - - discovered = [ - {"fold_id": "s0", "sample_id": "s0", "region_id": None}, - {"fold_id": "s1", "sample_id": "s1", "region_id": None}, - ] - with pytest.warns( - UserWarning, - match="loocv.fold_index takes priority", - ): - selected = _resolve_selected_folds(cfg, discovered) - - assert selected == ["s1"] - - -def test_to_float_metrics_warns_and_skips_non_convertible_values() -> None: - metrics = { - "train/loss": 0.5, - "val/ok": "1.25", - "val/bad": {"not": "numeric"}, - } - - with pytest.warns(UserWarning, match="Skipping non-numeric LOOCV metric 'val/bad'"): - converted = _to_float_metrics(metrics) - - assert converted == {"train/loss": 0.5, "val/ok": 1.25} diff --git a/tests/test_models.py b/tests/test_models.py new file mode 100644 index 0000000..b62518d --- /dev/null +++ b/tests/test_models.py @@ -0,0 +1,174 @@ +"""Encoder, attention pooling, objectives, probe, augmentations and optimisation.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +from torch_geometric.data import Batch, Data + +from grass_mil.data.datamodule import InputSpec +from grass_mil.models import SupervisedModule, attention_pool, pool_by_bag +from grass_mil.models.augmentations import drop_edges, drop_importance +from grass_mil.models.encoder import EncoderConfig, GINEEncoder +from grass_mil.models.objectives import ( + Objective, + class_weights_from_labels, + concordance_index, + cox_loss, +) +from grass_mil.models.probe import ProbeConfig, fit_linear_head + +SPEC = InputSpec(feature_dim=1, cell_types=("a", "b", "c"), edge_dim=2) + + +def _batch(n_graphs: int = 3, n: int = 6) -> Batch: + graphs = [] + for _ in range(n_graphs): + edge_index = torch.tensor( + [[i, i + 1] for i in range(n - 1)] + [[i + 1, i] for i in range(n - 1)] + ).t() + graphs.append( + Data( + x=torch.randn(n, 1), + cell_type=torch.randint(0, 3, (n,)), + edge_index=edge_index, + edge_attr=torch.rand(edge_index.size(1), 2), + ) + ) + return Batch.from_data_list(graphs) + + +def _module(**finetune) -> SupervisedModule: + module = SupervisedModule( + encoder={"hidden_dim": 8, "out_dim": 8, "num_layers": 2}, + head={"num_layers": 1}, + attention={"projection_dim": 4, "hidden_dim": 4}, + task={"target_type": "classification", "num_classes": 2, "class_weights": None}, + finetune={"lr_scales": {"encoder": 1.0, "head": 1.0, "attention": 1.0}, **finetune}, + optim={"_target_": "torch.optim.AdamW", "lr": 1e-3}, + scheduler=None, + ) + module.build(SPEC) + return module + + +def test_encoder_needs_some_input_and_embeds_cell_types_only(): + with pytest.raises(ValueError): + GINEEncoder(EncoderConfig(), feature_dim=0, cell_types=(), edge_dim=1) + enc = GINEEncoder( + EncoderConfig(hidden_dim=8, out_dim=4, num_layers=2), + feature_dim=0, + cell_types=("a", "b"), + edge_dim=1, + ) + batch = _batch() + batch.x = torch.zeros(batch.num_nodes, 0) + batch.cell_type = batch.cell_type % 2 + batch.edge_attr = batch.edge_attr[:, :1] + assert enc(batch).shape == (3, 4) + assert enc.unassigned_index == 2 # reserved extra row + + +def test_attention_pooling_is_an_exact_sum_of_contributions(): + logits, scores = torch.randn(10, 2), torch.randn(10, 2) + bag, weights = attention_pool(logits, scores) + assert torch.allclose(weights.sum(0), torch.ones(2)) + assert torch.allclose((weights * logits).sum(0), bag) + names, pooled = pool_by_bag( + torch.cat([logits, logits]), torch.cat([scores, scores]), ["x"] * 10 + ["y"] * 10 + ) + assert ( + names == ["x", "y"] and torch.allclose(pooled[0], bag) and torch.allclose(pooled[1], bag) + ) + + +def test_attention_needs_one_channel_per_class(): + with pytest.raises(ValueError, match="one channel per class"): + attention_pool(torch.randn(5, 2), torch.randn(5, 1)) + + +def test_class_weights_inverse_frequency_mean_one(): + w = class_weights_from_labels(np.array([0, 0, 0, 1]), 2) + assert torch.allclose(w.mean(), torch.tensor(1.0)) and w[1] > w[0] + + +def test_cox_loss_matches_the_partial_likelihood_and_ranks(): + risk = torch.tensor([0.5, -0.2, 1.0, 0.0], requires_grad=True) + time = torch.tensor([3.0, 5.0, 1.0, 4.0]) + event = torch.tensor([1.0, 0.0, 1.0, 1.0]) + loss = cox_loss(risk, time, event) + expected = 0.0 + for i in range(4): + if event[i]: + at_risk = time >= time[i] + expected -= float(risk[i] - torch.logsumexp(risk[at_risk], 0)) + assert loss.item() == pytest.approx(expected / 3, rel=1e-5) + loss.backward() + assert risk.grad is not None and torch.isfinite(risk.grad).all() + assert ( + concordance_index(np.array([3.0, 2, 1]), np.array([1.0, 2, 3]), np.array([1, 1, 1])) == 1.0 + ) + + +def test_survival_objective_needs_time_and_event(): + with pytest.raises(ValueError, match="time, event"): + Objective({"target_type": "survival", "target_columns": ["t"]}) + + +def test_probe_head_reproduces_the_logistic_decision(): + from sklearn.linear_model import LogisticRegression + from sklearn.preprocessing import StandardScaler + + rng = np.random.default_rng(0) + x = rng.normal(size=(200, 5)).astype(np.float32) + y = (x[:, 0] + 0.3 * rng.normal(size=200) > 0).astype(float) + weight, bias = fit_linear_head( + x, y[:, None], kind="classification", num_logits=2, cfg=ProbeConfig(class_weight=None) + ) + probs = torch.softmax(torch.from_numpy(x) @ weight.T + bias, 1)[:, 1].numpy() + reference = LogisticRegression(max_iter=3000).fit(StandardScaler().fit_transform(x), y) + expected = reference.predict_proba(StandardScaler().fit(x).transform(x))[:, 1] + assert np.allclose(probs, expected, atol=1e-4) + + +def test_augmentations_keep_graphs_undirected_and_inputs_intact(): + batch = _batch() + original = batch.edge_index.clone() + gen = torch.Generator().manual_seed(0) + out = drop_edges(batch, p=0.5, generator=gen) + pairs = set(map(tuple, out.edge_index.t().tolist())) + assert all((b, a) in pairs for a, b in pairs) + assert torch.equal(batch.edge_index, original) + masked = drop_importance(batch, mu=1.0, p_lambda=1.0, unassigned_index=3, generator=gen) + changed = masked.cell_type != batch.cell_type + assert bool((masked.cell_type[changed] == 3).all()) + assert bool((masked.x[changed] == 0).all()) + + +def test_gradient_clipping_happens_in_the_optimizer_hook(): + module = _module(grad_clip_norm=1.0) + for p in module.parameters(): + p.grad = torch.full_like(p, 10.0) + module.on_before_optimizer_step(None) + total = torch.sqrt(sum((p.grad**2).sum() for p in module.parameters())) + assert total.item() == pytest.approx(1.0, rel=1e-4) + + +def test_lr_scales_must_be_numbers(): + with pytest.raises(ValueError, match="lr_scales"): + SupervisedModule( + encoder={}, + head={}, + attention={}, + task={"target_type": "classification"}, + finetune={"lr_scales": {"encoder": None, "head": 1.0, "attention": 1.0}}, + ) + + +def test_zero_scale_freezes_the_encoder_in_eval_mode(): + module = _module(lr_scales={"encoder": 0.0, "head": 1.0, "attention": 2.0}) + module._apply_freezing() + module.train() + assert not any(p.requires_grad for p in module.encoder.parameters()) + assert not module.encoder.training and module.attention.training diff --git a/tests/test_nested_cv.py b/tests/test_nested_cv.py new file mode 100644 index 0000000..eae2523 --- /dev/null +++ b/tests/test_nested_cv.py @@ -0,0 +1,66 @@ +"""Nested cross-validation.""" + +from __future__ import annotations + +import json + +import pandas as pd +import pytest + +from grass_mil.nested_cv import best_epoch, run_nested_cv +from tests.helpers.cohort import cohort_overrides, compose_cfg + + +def test_best_epoch_prefers_the_metric_then_the_loss_then_the_last_epoch(): + rows = [ + {"epoch": 0, "val/auroc": 0.6, "val/loss": 0.7}, + {"epoch": 1, "val/auroc": 0.8, "val/loss": 0.8}, + {"epoch": 2, "val/auroc": 0.8, "val/loss": 0.5}, + ] + assert best_epoch(rows, "val/auroc", "max") == (2, "val/auroc") + nan = [ + {"epoch": 0, "val/auroc": float("nan"), "val/loss": 0.9}, + {"epoch": 1, "val/auroc": float("nan"), "val/loss": 0.4}, + ] + assert best_epoch(nan, "val/auroc", "max") == (2, "val/loss") + assert best_epoch([{"epoch": 0}, {"epoch": 3}], "val/auroc", "max") == (4, "last") + + +def test_nested_cv_summary_and_out_of_fold_predictions(cohort, nested_cv): + summary = json.loads((nested_cv / "nested_cv_summary.json").read_text()) + assert summary["regime"] == "finetune" and summary["metric"] == "auroc" + assert len(summary["fold_values"]) == 3 and summary["complete"] + oof = pd.read_csv(nested_cv / "oof_predictions.csv") + assert sorted(oof["unit_id"]) == sorted(cohort.regions) + patients = pd.read_csv(nested_cv / "oof_patient_predictions.csv") + assert len(patients) == len(set(cohort.patient_of.values())) + assert "auroc" in summary["patient"] + + +def test_refit_trains_without_validation(nested_cv): + refit = nested_cv / "outer_0" / "refit" + assert not (refit / "metrics_per_epoch.jsonl").exists() + assert (refit / "checkpoints" / "last.ckpt").exists() + done = json.loads((refit / "done.json").read_text()) + assert not any(k.startswith("val/") for k in done["metrics"]) + + +@pytest.mark.slow +def test_nested_cv_for_survival_selects_on_the_concordance(cohort, pretrained, tmp_path): + cfg = compose_cfg( + "nested_cv", + cohort_overrides(cohort) + + [ + "task=survival", + "task.target_columns=[time,event]", + f"nested_cv.folds_file={cohort.folds}", + "nested_cv.inner_k=2", + "nested_cv.outer_folds=[0]", + "trainer.max_epochs=1", + f"model.init_from_ckpt={pretrained}", + ], + tmp_path, + ) + summary = run_nested_cv(cfg) + assert summary["metric"] == "c_index" and summary["target_type"] == "survival" + assert "risk" in pd.read_csv(tmp_path / "oof_predictions.csv").columns diff --git a/tests/test_niches.py b/tests/test_niches.py new file mode 100644 index 0000000..445ee96 --- /dev/null +++ b/tests/test_niches.py @@ -0,0 +1,85 @@ +"""Niche discovery, attribution, descriptors, survival and the niches command.""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +from grass_mil.niches import cell_margins, fit_niches +from grass_mil.niches.cli import run_niches +from grass_mil.niches.discovery import ClusteringConfig +from grass_mil.niches.survival import cox_univariate, logrank +from tests.helpers.cohort import compose_cfg + + +def test_margins_sum_to_the_bag_margin(): + rng = np.random.default_rng(0) + logits, scores = rng.normal(size=(30, 2)), rng.normal(size=(30, 2)) + bias = np.array([0.2, -0.1]) + units = np.array(["a"] * 10 + ["b"] * 20) + margin = cell_margins(logits, scores, bias, units) + for unit, n in (("a", 10), ("b", 20)): + m = units == unit + w = np.exp(scores[m] - scores[m].max(0)) + w /= w.sum(0) + bag = (w * logits[m]).sum(0) - bias + assert margin[m].sum() / n == pytest.approx(bag[1] - bag[0], rel=1e-5) + + +def test_niches_are_numbered_by_margin_with_background_zero(): + rng = np.random.default_rng(0) + x = np.concatenate([rng.normal(c, 0.05, size=(100, 2)) for c in ((0, 1), (1, 0), (-1, 0))]) + margin = np.repeat([0.0, 2.0, -2.0], 100) + model, niches = fit_niches(x, margin, ClusteringConfig(min_cluster_size=20, sample_size=200)) + assert model.n_niches == 3 + assert np.mean(niches[100:200] == 1) > 0.9 and np.mean(niches[200:] == 3) > 0.9 + assert np.array_equal(model.predict(x[:5]), niches[:5]) + + +def test_cox_and_logrank_detect_a_strong_effect(): + rng = np.random.default_rng(0) + x = rng.normal(size=300) + time = rng.exponential(1.0 / np.exp(0.8 * x)) + event = np.ones(300) + hr, lo, hi, p = cox_univariate(x, time, event) + assert lo < np.exp(0.8) < hi and p < 1e-6 + groups = (x > 0).astype(int) + assert logrank(time, event, groups) < 1e-4 + assert logrank(time, event, rng.integers(0, 2, 300)) > 1e-3 + + +def test_niches_command_writes_the_primary_results(cohort, cells_dir, tmp_path): + (tmp_path / "outcomes.csv").write_text( + "patient_id,time,event\n" + "\n".join(f"P{p:02d},{10 + p},{p % 2}" for p in range(6)) + ) + cfg = compose_cfg( + "niches", + [ + f"cells_dir={cells_dir / 'fold0'}", + "units=all", + "clustering.min_cluster_size=20", + "descriptors.min_cells=10", + "descriptors.n_perms=10", + "descriptors.min_contacts=5", + "n_boot=100", + f"survival.file={tmp_path / 'outcomes.csv'}", + "survival.score=[Niche 1]", + ], + tmp_path / "out", + ) + out = run_niches(cfg) + table = pd.read_csv(out / "niches.csv") + assert {"niche", "share", "association", "margin", "margin_lo", "margin_hi", "frac_D"} <= set( + table.columns + ) + assert table["share"].sum() == pytest.approx(1.0) + for name in ("moran", "enrichment", "cross_l", "contacts"): + assert (out / "descriptors" / f"{name}.csv").exists() + cells = pd.read_csv(out / "cells.csv") + assert len(cells) == sum(120 for _ in cohort.regions) + for figure in ("niche_panel", "embedding", "unit_composition", "survival"): + assert (out / "figures" / f"{figure}.png").exists() + assert len(list((out / "figures" / "units").glob("*.png"))) == 4 + assert "| Niche |" in (out / "report.md").read_text() + assert len(pd.read_csv(out / "survival_cox.csv")) >= 1 diff --git a/tests/test_packaging.py b/tests/test_packaging.py deleted file mode 100644 index 5a0de9f..0000000 --- a/tests/test_packaging.py +++ /dev/null @@ -1,68 +0,0 @@ -"""Guards on the installable-package contract. - -Configs ship inside ``grass_mil`` so Hydra resolves them identically from a -source checkout and from an installed wheel. These tests fail loudly if a move -or a rename breaks that, which is otherwise only discovered when someone -installs the package. -""" - -from __future__ import annotations - -import importlib -import tomllib -from pathlib import Path - -import pytest -from hydra import compose, initialize_config_module -from hydra.core.global_hydra import GlobalHydra - -ROOT_CONFIGS = [ - "train.yaml", - "eval.yaml", - "loocv.yaml", - "inference/predict.yaml", - "interpretability/report.yaml", -] - - -def _pyproject() -> dict: - path = Path(__file__).resolve().parents[1] / "pyproject.toml" - with path.open("rb") as handle: - return tomllib.load(handle) - - -def test_package_is_importable_under_its_distribution_name() -> None: - module = importlib.import_module("grass_mil") - assert Path(module.__file__).parent.name == "grass_mil" - - -def test_configs_ship_inside_the_package() -> None: - configs = importlib.import_module("grass_mil.configs") - configs_dir = Path(configs.__file__).parent - for name in ROOT_CONFIGS: - assert (configs_dir / name).is_file(), f"missing packaged config: {name}" - - -@pytest.mark.parametrize("config_name", ROOT_CONFIGS) -def test_root_configs_compose_from_the_config_module(config_name: str) -> None: - """Catches a broken ``config_path`` after a layout change.""" - GlobalHydra.instance().clear() - with initialize_config_module(version_base="1.3", config_module="grass_mil.configs"): - cfg = compose(config_name=config_name, return_hydra_config=True, overrides=[]) - assert cfg is not None - GlobalHydra.instance().clear() - - -def test_console_script_targets_are_importable_and_callable() -> None: - scripts = _pyproject()["project"]["scripts"] - assert scripts, "no console scripts declared" - for name, target in scripts.items(): - module_path, _, attribute = target.partition(":") - module = importlib.import_module(module_path) - assert callable(getattr(module, attribute)), f"{name} -> {target} is not callable" - - -def test_wheel_packages_the_source_tree() -> None: - config = _pyproject() - assert config["build-system"]["build-backend"] == "hatchling.build" - assert config["tool"]["hatch"]["build"]["targets"]["wheel"]["packages"] == ["src/grass_mil"] diff --git a/tests/test_samplers.py b/tests/test_samplers.py new file mode 100644 index 0000000..67c88f6 --- /dev/null +++ b/tests/test_samplers.py @@ -0,0 +1,154 @@ +"""Ego-graph sampling.""" + +from __future__ import annotations + +import pytest +import torch +from torch_geometric.data import Data +from torch_geometric.utils import k_hop_subgraph + +from grass_mil.data.datamodule import GraphUnitDataset +from grass_mil.data.precompute import UnitEntry +from grass_mil.data.samplers import EgoGraphLoader, SamplerConfig + + +class _Units(GraphUnitDataset): + """In-memory units with the dataset interface the loader needs.""" + + def __init__(self, graphs): + super().__init__( + [UnitEntry(g.unit_id, g.region_id, g.sample_id, "", int(g.num_nodes)) for g in graphs] + ) + self.graphs = graphs + + def __getitem__(self, idx): + return self.graphs[idx] + + +def _grid_unit(name: str, side: int = 8, y: float = 1.0) -> Data: + xs, ys = torch.meshgrid(torch.arange(side), torch.arange(side), indexing="ij") + pos = torch.stack([xs.flatten(), ys.flatten()], 1).float() * 10 + edges = [] + for i in range(side * side): + r, c = divmod(i, side) + if c + 1 < side: + edges += [(i, i + 1), (i + 1, i)] + if r + 1 < side: + edges += [(i, i + side), (i + side, i)] + edge_index = torch.tensor(edges).t() + n = side * side + return Data( + x=torch.arange(n, dtype=torch.float).view(-1, 1), + pos=pos, + edge_index=edge_index, + edge_attr=torch.ones(edge_index.size(1), 1), + cell_type=torch.arange(n) % 3, + num_nodes=n, + unit_id=name, + region_id=name, + sample_id=name, + graph_y=torch.tensor([[y]]), + ) + + +def _loader(units, mode, **cfg): + return EgoGraphLoader( + _Units(units), + SamplerConfig.from_dict({"depth": 2, "batch_size": 5, **cfg}), + mode=mode, + roots_per_unit=cfg.get("roots_per_unit"), + ) + + +def test_bags_hold_one_unit_each_with_its_target(): + units = [_grid_unit("a", y=0.0), _grid_unit("b", y=1.0)] + batches = list(_loader(units, "bags", roots_per_unit=7)) + assert len(batches) == 2 + for batch in batches: + assert batch.num_graphs == 7 + assert len(set(batch.unit_id)) == 1 + assert batch.graph_y.shape == (7, 1) + + +def test_ego_graphs_are_the_k_hop_neighbourhoods_of_their_roots(): + unit = _grid_unit("a") + for batch in _loader([unit], "chunks"): + for i in range(batch.num_graphs): + root = int(batch.root_cell[i]) + nodes, _, _, _ = k_hop_subgraph(root, 2, unit.edge_index, num_nodes=unit.num_nodes) + sub = batch.x[batch.batch == i].view(-1).long() + assert sorted(sub.tolist()) == sorted(nodes.tolist()) + assert int(batch.x[batch.root_n_id[i]]) == root + + +def test_chunks_score_every_cell_once_and_len_matches(): + loader = _loader([_grid_unit("a"), _grid_unit("b")], "chunks") + batches = list(loader) + assert len(batches) == len(loader) + for name in "ab": + roots = torch.cat([b.root_cell for b in batches if b.unit_id[0] == name]) + assert sorted(roots.tolist()) == list(range(64)) + + +def test_fixed_seed_draws_the_same_roots_every_pass(): + units = [_grid_unit("a")] + loader = EgoGraphLoader( + _Units(units), + SamplerConfig.from_dict({"depth": 1}), + mode="chunks", + roots_per_unit=10, + seed=3, + ) + first = torch.cat([b.root_cell for b in loader]) + second = torch.cat([b.root_cell for b in loader]) + assert torch.equal(first, second) and first.numel() == 10 + + +def test_root_weighting_draws_distinct_roots(): + loader = EgoGraphLoader( + _Units([_grid_unit("a")]), + SamplerConfig.from_dict({"depth": 1, "root_weighting": "inverse"}), + mode="bags", + roots_per_unit=20, + ) + roots = next(iter(loader)).root_cell + assert roots.unique().numel() == 20 + + +def test_merged_mode_mixes_units(): + loader = _loader( + [_grid_unit("a"), _grid_unit("b")], "merged", roots_per_unit=None, merge_units=True + ) + batches = list(loader) + assert len(batches) == len(loader) == -(-128 // 5) + assert sum(b.num_graphs for b in batches) == 128 + + +def test_interior_roots_exclude_the_rim(): + loader = _loader( + [_grid_unit("a", side=10)], "chunks", interior_roots={"enabled": True, "n_hops": 1} + ) + roots = torch.cat([b.root_cell for b in loader]) + rows, cols = roots // 10, roots % 10 + assert roots.numel() < 100 + assert bool(((rows > 0) & (rows < 9) & (cols > 0) & (cols < 9)).all()) + + +def test_ego_radius_keeps_the_root_and_bounds_the_extent(): + loader = _loader( + [_grid_unit("a")], + "chunks", + ego_radius={"enabled": True, "radius_per_hop": 5.0, "radius_offset": 0.0}, + ) + for batch in loader: + for i in range(batch.num_graphs): + nodes = batch.batch == i + root = batch.root_n_id[i] + assert bool(nodes[root]) + dist = torch.linalg.norm(batch.pos[nodes] - batch.pos[root], dim=1) + assert float(dist.max()) <= 10.0 + 1e-6 + + +def test_unknown_sampler_keys_are_rejected(): + with pytest.raises(ValueError, match="Unknown"): + SamplerConfig.from_dict({"subsample_count": 3}) diff --git a/tests/test_spatial.py b/tests/test_spatial.py new file mode 100644 index 0000000..910c5bd --- /dev/null +++ b/tests/test_spatial.py @@ -0,0 +1,106 @@ +"""Spatial statistics, cross-fold replication and the model-free command.""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +from grass_mil.niches.replication.cli import run_metaclusters +from grass_mil.niches.replication.stages import constrained_linkage, hellinger +from grass_mil.spatial.cli import run_spatial_stats +from grass_mil.spatial.graph_stats import contact_counts, cross_l, enrichment_zscores, morans_i +from grass_mil.spatial.stats import bh, bootstrap_mean +from tests.helpers.cohort import compose_cfg + + +def _chain(n: int): + return np.stack([np.arange(n - 1), np.arange(1, n)]) + + +def test_contact_counts_count_edge_ends(): + labels = np.array([0, 0, 1, 1]) + counts = contact_counts(labels, _chain(4), 2) + assert counts.tolist() == [[2, 1], [1, 2]] + assert counts.sum() == 2 * 3 + + +def test_enrichment_is_positive_for_segregated_labels(): + labels = np.repeat([0, 1], 50) + z = enrichment_zscores(labels, _chain(100), 2, n_perms=200, rng=np.random.default_rng(0)) + assert z[0, 0] > 3 and z[0, 1] < -3 + + +def test_morans_i_of_clustered_and_alternating_indicators(): + clustered = np.repeat([1.0, 0.0], 50)[:, None] + alternating = (np.arange(100) % 2).astype(float)[:, None] + assert morans_i(clustered, _chain(100))[0] > 0.9 + assert morans_i(alternating, _chain(100))[0] < -0.9 + + +def test_cross_l_is_near_zero_under_complete_spatial_randomness(): + coords = np.random.default_rng(0).uniform(0, 100, size=(2000, 2)) + radii = np.array([5.0, 10.0]) + labels = np.arange(2000) % 2 + l_aa = cross_l(coords, labels == 0, labels == 0, radii) + l_ab = cross_l(coords, labels == 0, labels == 1, radii) + # Edge effects bias L slightly downwards; the self pairs must not bias it upwards. + assert np.all(np.abs(l_aa) < 0.6) and np.all(np.abs(l_ab) < 0.6) + + +def test_bh_and_bootstrap(): + assert np.allclose(bh(np.array([0.01, 0.04, np.nan])), [0.02, 0.04, np.nan], equal_nan=True) + mean, lo, hi = bootstrap_mean([1.0, 2.0, 3.0, np.nan], n_boot=500) + assert mean == 2.0 and lo <= 2.0 <= hi + + +def test_linkage_never_merges_two_niches_of_one_model(): + comp = np.array([[1.0, 0.0], [0.99, 0.01], [0.98, 0.02]]) + groups = constrained_linkage(hellinger(comp), ["m1", "m1", "m2"], threshold=0.5) + assert sorted(len(g) for g in groups) == [1, 2] + + +def test_metaclusters_command(cells_dir, tmp_path): + cfg = compose_cfg( + "metaclusters", + [ + f"cells_dir={cells_dir}", + "+clustering.min_cluster_size=20", + "match.min_models=2", + "basis.k=3", + "n_boot=50", + ], + tmp_path, + ) + summary = run_metaclusters(cfg) + assert summary["models"] == ["fold0", "fold1", "fold2"] + table = pd.read_csv(tmp_path / "metaclusters.csv") + assert table["metacluster"].str.startswith("M").all() + assert (tmp_path / "transfer" / "agreement.csv").exists() + assert len(pd.read_csv(tmp_path / "basis" / "groups.csv")) == 3 + + +def test_metaclusters_rejects_unknown_stages(cells_dir, tmp_path): + cfg = compose_cfg("metaclusters", [f"cells_dir={cells_dir}", "stage=everything"], tmp_path) + with pytest.raises(ValueError, match="stage"): + run_metaclusters(cfg) + + +def test_spatial_stats_command(cells_dir, tmp_path): + cfg = compose_cfg( + "spatial_stats", + [ + f"cells_dir={cells_dir / 'fold0'}", + "reference=[D]", + "n_perms=20", + "neighbourhoods.n=3", + "min_patients=2", + ], + tmp_path, + ) + run_spatial_stats(cfg) + tests = pd.read_csv(tmp_path / "cross_l_tests.csv").set_index("metric") + # The planted niche: D cells cluster in positive regions only. + assert tests.loc["D->D@20", "mean_pos"] > tests.loc["D->D@20", "mean_neg"] + for name in ("enrichment", "neighbourhoods", "contacts"): + assert (tmp_path / f"{name}_tests.csv").exists() diff --git a/tests/test_spatial_omics_datamodule.py b/tests/test_spatial_omics_datamodule.py deleted file mode 100644 index 5e2ce54..0000000 --- a/tests/test_spatial_omics_datamodule.py +++ /dev/null @@ -1,1170 +0,0 @@ -import json -from pathlib import Path - -import numpy as np -import pytest - - -def _write_polygon_wkt(path: Path, coords: np.ndarray) -> None: - pytest.importorskip("shapely") - from shapely.geometry import Polygon - - polygon = Polygon(coords) - path.write_text(polygon.wkt) - - -def test_spatial_omics_datamodule_precompute(tmp_path: Path) -> None: - pytest.importorskip("torch_geometric") - pytest.importorskip("scipy") - - data_dir = tmp_path / "data" - raw_dir = data_dir / "raw" - raw_dir.mkdir(parents=True) - - coords = np.array( - [ - [0.0, 0.0], - [0.0, 10.0], - [10.0, 0.0], - [10.0, 10.0], - ] - ) - csv_path = raw_dir / "sample.csv" - np.savetxt( - csv_path, - np.column_stack( - [ - np.arange(coords.shape[0]), - coords[:, 0], - coords[:, 1], - np.random.rand(coords.shape[0]), - np.random.rand(coords.shape[0]), - ["A", "B", "A", "B"], - ] - ), - delimiter=",", - fmt="%s", - header="cell_id,x,y,gene1,gene2,cell_type", - comments="", - ) - - polygon_path = raw_dir / "polygons.wkt" - _write_polygon_wkt( - polygon_path, - np.array([[-1, -1], [-1, 20], [20, 20], [20, -1], [-1, -1]]), - ) - - manifest_path = raw_dir / "manifest.csv" - manifest_path.write_text( - "sample_id,input_path,input_type,region_id,polygons_path\n" - f"sample_1,{csv_path},csv,region_1,{polygon_path}\n" - ) - - graph_labels_path = raw_dir / "graph_labels.csv" - graph_labels_path.write_text("region_id,outcome\nregion_1,1\n") - - from grass_mil.data.spatial_omics_datamodule import SpatialOmicsDataModule - - dm = SpatialOmicsDataModule( - data_dir=str(data_dir), - raw_manifest_path=str(manifest_path), - processed_dir=str(data_dir / "processed"), - coord_scale_um=1.0, - sample_unit="tile", - batch_size=1, - num_workers=0, - pin_memory=False, - reducer_scope="sample", - keep_raw_molecular=True, - force_precompute=True, - min_cells=1, - use_molecular_features=True, - categorical_features={ - "include_labels": ["cell_type"], - }, - split={"train_val_test_split": [1.0, 0.0, 0.0], "split_by": "patch"}, - csv={ - "sep": ",", - "coord_columns": ["x", "y"], - "cell_id_column": "cell_id", - "categorical_label_columns": ["cell_type"], - "molecular_columns": ["gene1", "gene2"], - }, - h5ad={ - "coord_columns": ["x", "y"], - "cell_id_column": None, - "categorical_label_columns": [], - "molecular_layer": None, - }, - sce={ - "assay_name": None, - "coord_source": "colData", - "coord_key": None, - "coord_columns": ["x", "y"], - "cell_id_column": None, - "categorical_label_columns": [], - "transpose_assay": True, - }, - graph_builder={ - "name": "delaunay", - "kwargs": { - "edge_features": ["distance", "neighbor"], - "neighbor_cutoff_um": 50.0, - }, - }, - feature_reducer={"name": "identity", "kwargs": {}}, - tiling={"tile_size_um": 50.0, "stride_um": 50.0, "min_cells": 1}, - manifest={ - "sample_id": "sample_id", - "input_path": "input_path", - "input_type": "input_type", - "region_id": "region_id", - "polygons_path": "polygons_path", - }, - graph_labels={ - "label_file": str(graph_labels_path), - "id_column": "region_id", - "tasks": ["outcome"], - "scope": "region", - }, - transforms=[], - ) - - dm.prepare_data() - dm.setup() - - dataset = dm.train_dataloader().dataset - assert len(dataset) > 0 - - metadata_path = data_dir / "processed" / "metadata.json" - metadata = json.loads(metadata_path.read_text()) - assert metadata["molecular_feature_dim"] == 2 - assert metadata["molecular_feature_names"] == ["gene1", "gene2"] - assert metadata["categorical_labels"] == ["cell_type"] - - data = dataset[0] - assert hasattr(data, "x") - assert hasattr(data, "pos") - assert data.patch_id.startswith("sample_1_region_1") - assert hasattr(data, "edge_attr") - assert hasattr(data, "graph_y") - assert data.x.shape[1] == 2 # gene1, gene2 - assert hasattr(data, "categorical_codes") - assert data.categorical_codes.shape[1] == 1 - assert data.categorical_slices["cell_type"] == 0 - - index_payload = json.loads((data_dir / "processed" / "processed_index.json").read_text()) - assert all( - entry.get("split") in {"train", "val", "test"} for entry in index_payload["entries"] - ) - - -def test_spatial_omics_datamodule_indices_no_molecular(tmp_path: Path) -> None: - pytest.importorskip("torch_geometric") - pytest.importorskip("scipy") - - data_dir = tmp_path / "data" - raw_dir = data_dir / "raw" - raw_dir.mkdir(parents=True) - - coords = np.array([[0.0, 0.0], [0.0, 10.0]]) - csv_path = raw_dir / "sample.csv" - np.savetxt( - csv_path, - np.column_stack( - [ - np.arange(coords.shape[0]), - coords[:, 0], - coords[:, 1], - ["A", "B"], - ] - ), - delimiter=",", - fmt="%s", - header="cell_id,x,y,cell_type", - comments="", - ) - - polygon_path = raw_dir / "polygons.wkt" - _write_polygon_wkt( - polygon_path, - np.array([[-1, -1], [-1, 20], [20, 20], [20, -1], [-1, -1]]), - ) - - manifest_path = raw_dir / "manifest.csv" - manifest_path.write_text( - "sample_id,input_path,input_type,region_id,polygons_path\n" - f"sample_1,{csv_path},csv,region_1,{polygon_path}\n" - ) - - from grass_mil.data.spatial_omics_datamodule import SpatialOmicsDataModule - - dm = SpatialOmicsDataModule( - data_dir=str(data_dir), - raw_manifest_path=str(manifest_path), - processed_dir=str(data_dir / "processed"), - coord_scale_um=1.0, - sample_unit="tile", - batch_size=1, - num_workers=0, - pin_memory=False, - reducer_scope="sample", - keep_raw_molecular=False, - force_precompute=True, - min_cells=1, - use_molecular_features=False, - categorical_features={ - "include_labels": ["cell_type"], - }, - split={"train_val_test_split": [1.0, 0.0, 0.0], "split_by": "patch"}, - csv={ - "sep": ",", - "coord_columns": ["x", "y"], - "cell_id_column": "cell_id", - "categorical_label_columns": ["cell_type"], - "molecular_columns": None, - }, - h5ad={ - "coord_columns": ["x", "y"], - "cell_id_column": None, - "categorical_label_columns": [], - "molecular_layer": None, - }, - sce={ - "assay_name": None, - "coord_source": "colData", - "coord_key": None, - "coord_columns": ["x", "y"], - "cell_id_column": None, - "categorical_label_columns": [], - "transpose_assay": True, - }, - graph_builder={ - "name": "delaunay", - "kwargs": { - "edge_features": ["distance", "neighbor"], - "neighbor_cutoff_um": 50.0, - }, - }, - feature_reducer={"name": "identity", "kwargs": {}}, - tiling={"tile_size_um": 50.0, "stride_um": 50.0, "min_cells": 1}, - manifest={ - "sample_id": "sample_id", - "input_path": "input_path", - "input_type": "input_type", - "region_id": "region_id", - "polygons_path": "polygons_path", - }, - graph_labels=None, - transforms=[], - ) - - dm.prepare_data() - dm.setup() - - data = dm.train_dataloader().dataset[0] - assert hasattr(data, "x") - assert data.x.shape[1] == 0 - assert hasattr(data, "categorical_codes") - assert data.categorical_codes.shape == (2, 1) - - -def test_h5ad_obsm_spatial_coords(tmp_path: Path) -> None: - pytest.importorskip("anndata") - - import anndata as ad - - coords = np.array([[1.0, 2.0], [3.0, 4.0]]) - x = np.random.rand(2, 3) - adata = ad.AnnData(X=x) - adata.var_names = ["0", "1", "2"] - adata.obsm["spatial"] = coords - adata.obs["cell_type"] = ["A", "B"] - - p = tmp_path / "test.h5ad" - adata.write_h5ad(p) - - from grass_mil.data.components.loaders import H5adConfig, load_h5ad_table - - table = load_h5ad_table( - p, - H5adConfig( - coord_source="obsm", - coord_key="spatial", - coord_columns=("x", "y"), - cell_id_column=None, - categorical_label_columns=("cell_type",), - molecular_layer=None, - molecular_features=["0", "2"], - use_molecular_features=True, - ), - ) - assert np.allclose(table.coords, coords) - assert table.molecular_features.shape[1] == 2 - assert list(table.molecular_feature_names) == ["0", "2"] - - -def test_datamodule_requires_persisted_splits(tmp_path: Path) -> None: - pytest.importorskip("lightning") - pytest.importorskip("torch_geometric") - - processed_dir = tmp_path / "processed" - processed_dir.mkdir(parents=True) - index_path = processed_dir / "processed_index.json" - - entries = [ - { - "path": str(tmp_path / f"graph_{i}.pt"), - "sample_id": f"sample_{i % 3}", - "region_id": f"region_{i % 2}", - "patch_id": f"patch_{i}", - } - for i in range(12) - ] - index_path.write_text( - json.dumps( - { - "entries": entries, - "label_maps": {}, - "graph_label_maps": {}, - "reducer_state": {}, - } - ) - ) - - from grass_mil.data.spatial_omics_datamodule import SpatialOmicsDataModule - - dm = SpatialOmicsDataModule( - data_dir=str(tmp_path), - raw_manifest_path=str(tmp_path / "unused_manifest.csv"), - processed_dir=str(processed_dir), - coord_scale_um=1.0, - sample_unit="tile", - batch_size=1, - num_workers=0, - pin_memory=False, - reducer_scope="sample", - keep_raw_molecular=False, - force_precompute=False, - min_cells=1, - use_molecular_features=False, - categorical_features={"include_labels": []}, - split={"train_val_test_split": [0.5, 0.25, 0.25], "split_by": "patch"}, - csv={ - "sep": ",", - "coord_columns": ["x", "y"], - "cell_id_column": "cell_id", - "categorical_label_columns": [], - "molecular_columns": None, - }, - h5ad={ - "coord_columns": ["x", "y"], - "cell_id_column": None, - "categorical_label_columns": [], - "molecular_layer": None, - }, - sce={ - "assay_name": None, - "coord_source": "colData", - "coord_key": None, - "coord_columns": ["x", "y"], - "cell_id_column": None, - "categorical_label_columns": [], - "transpose_assay": True, - }, - graph_builder={"name": "delaunay", "kwargs": {}}, - feature_reducer={"name": "identity", "kwargs": {}}, - tiling={"tile_size_um": 50.0, "stride_um": 50.0, "min_cells": 1}, - manifest=None, - graph_labels=None, - transforms=[], - split_seed=7, - ) - - with pytest.raises( - ValueError, - match="processed_index.json must contain valid per-entry split labels", - ): - dm.setup() - - -def test_datamodule_uses_persisted_splits(tmp_path: Path) -> None: - pytest.importorskip("lightning") - pytest.importorskip("torch_geometric") - - processed_dir = tmp_path / "processed" - processed_dir.mkdir(parents=True) - index_path = processed_dir / "processed_index.json" - entries = [] - for i, split in enumerate(["train", "train", "val", "test", "test"]): - entries.append( - { - "path": str(tmp_path / f"graph_{i}.pt"), - "sample_id": f"sample_{i}", - "region_id": f"region_{i}", - "patch_id": f"patch_{i}", - "split": split, - } - ) - index_path.write_text( - json.dumps( - { - "entries": entries, - "label_maps": {}, - "graph_label_maps": {}, - "reducer_state": {}, - } - ) - ) - - from grass_mil.data.spatial_omics_datamodule import SpatialOmicsDataModule - - dm = SpatialOmicsDataModule( - data_dir=str(tmp_path), - raw_manifest_path=str(tmp_path / "unused_manifest.csv"), - processed_dir=str(processed_dir), - coord_scale_um=1.0, - sample_unit="tile", - batch_size=1, - num_workers=0, - pin_memory=False, - reducer_scope="sample", - keep_raw_molecular=False, - force_precompute=False, - min_cells=1, - use_molecular_features=False, - categorical_features={"include_labels": []}, - split={"train_val_test_split": [1.0, 0.0, 0.0], "split_by": "patch"}, - csv={ - "sep": ",", - "coord_columns": ["x", "y"], - "cell_id_column": "cell_id", - "categorical_label_columns": [], - "molecular_columns": None, - }, - h5ad={ - "coord_columns": ["x", "y"], - "cell_id_column": None, - "categorical_label_columns": [], - "molecular_layer": None, - }, - sce={ - "assay_name": None, - "coord_source": "colData", - "coord_key": None, - "coord_columns": ["x", "y"], - "cell_id_column": None, - "categorical_label_columns": [], - "transpose_assay": True, - }, - graph_builder={"name": "delaunay", "kwargs": {}}, - feature_reducer={"name": "identity", "fit_mode": "train_only", "kwargs": {}}, - tiling={"tile_size_um": 50.0, "stride_um": 50.0, "min_cells": 1}, - manifest=None, - graph_labels=None, - transforms=[], - split_seed=999, - ) - dm.setup() - assert len(dm.dataset_train.entries) == 2 - assert len(dm.dataset_val.entries) == 1 - assert len(dm.dataset_test.entries) == 2 - - -def _make_unit( - table_idx: int, sample_id: str, region_id: str, patch_idx: int, split: str = "train" -): - from grass_mil.data.components.precompute import GraphUnitSpec - - return GraphUnitSpec( - table_idx=table_idx, - sample_id=sample_id, - region_id=region_id, - patch_idx=patch_idx, - patch_id=f"{sample_id}_{region_id}_patch_{patch_idx}", - indices=np.array([0, 1], dtype=int), - split=split, - ) - - -def test_region_split_groups_all_patches_per_region() -> None: - from grass_mil.data.components.precompute import _assign_split_labels - - units = [ - _make_unit(0, "s0", "r0", 0), - _make_unit(0, "s0", "r0", 1), - _make_unit(1, "s1", "r1", 0), - _make_unit(1, "s1", "r1", 1), - _make_unit(2, "s2", "r2", 0), - _make_unit(2, "s2", "r2", 1), - ] - labels = _assign_split_labels( - units=units, split_by="region", split_ratios=(0.5, 0.25, 0.25), seed=123 - ) - by_region = {} - for unit, label in zip(units, labels): - by_region.setdefault((unit.sample_id, unit.region_id), set()).add(label) - assert all(len(splits) == 1 for splits in by_region.values()) - - -def test_loocv_holdout_region_is_test_only_with_region_validation() -> None: - from grass_mil.data.components.precompute import LoocvConfig, _assign_split_labels - - units = [ - _make_unit(0, "s0", "r0", 0), - _make_unit(0, "s0", "r0", 1), - _make_unit(1, "s1", "r1", 0), - _make_unit(1, "s1", "r1", 1), - _make_unit(2, "s2", "r2", 0), - _make_unit(2, "s2", "r2", 1), - ] - labels = _assign_split_labels( - units=units, - split_by="region", - split_ratios=(0.8, 0.1, 0.1), - seed=3, - loocv_config=LoocvConfig( - enabled=True, - fold_unit="region", - holdout_id="s1::r1", - validation_strategy="heldout_fold_items", - val_ratio=0.5, - seed=11, - ), - ) - for unit, label in zip(units, labels): - if unit.sample_id == "s1" and unit.region_id == "r1": - assert label == "test" - else: - assert label in {"train", "val"} - - -def test_loocv_patch_validation_uses_only_train_fold_items() -> None: - from grass_mil.data.components.precompute import LoocvConfig, _assign_split_labels - - units = [ - _make_unit(0, "s0", "r0", 0), - _make_unit(0, "s0", "r0", 1), - _make_unit(1, "s1", "r1", 0), - _make_unit(1, "s1", "r1", 1), - _make_unit(2, "s2", "r2", 0), - _make_unit(2, "s2", "r2", 1), - ] - labels = _assign_split_labels( - units=units, - split_by="region", - split_ratios=(0.8, 0.1, 0.1), - seed=5, - loocv_config=LoocvConfig( - enabled=True, - fold_unit="region", - holdout_id="s2::r2", - validation_strategy="patches_from_train_items", - val_ratio=0.4, - seed=7, - ), - ) - # Held-out region must be test-only and never val. - for unit, label in zip(units, labels): - if unit.sample_id == "s2" and unit.region_id == "r2": - assert label == "test" - assert "val" in labels - - -def test_loocv_heldout_validation_allows_zero_val_ratio() -> None: - pytest.importorskip("lightning") - pytest.importorskip("torch_geometric") - from grass_mil.data.components.precompute import LoocvConfig, _assign_split_labels - - units = [ - _make_unit(0, "s0", "r0", 0), - _make_unit(1, "s1", "r1", 0), - _make_unit(2, "s2", "r2", 0), - ] - labels = _assign_split_labels( - units=units, - split_by="region", - split_ratios=(0.8, 0.1, 0.1), - seed=3, - loocv_config=LoocvConfig( - enabled=True, - fold_unit="region", - holdout_id="s1::r1", - validation_strategy="heldout_fold_items", - val_ratio=0.0, - seed=11, - ), - ) - assert "val" not in labels - for unit, label in zip(units, labels): - if unit.sample_id == "s1" and unit.region_id == "r1": - assert label == "test" - - -def test_loocv_heldout_validation_forces_one_for_nonzero_ratio() -> None: - pytest.importorskip("lightning") - pytest.importorskip("torch_geometric") - from grass_mil.data.components.precompute import LoocvConfig, _assign_split_labels - - units = [ - _make_unit(0, "s0", "r0", 0), - _make_unit(1, "s1", "r1", 0), - _make_unit(2, "s2", "r2", 0), - ] - labels = _assign_split_labels( - units=units, - split_by="region", - split_ratios=(0.8, 0.1, 0.1), - seed=3, - loocv_config=LoocvConfig( - enabled=True, - fold_unit="region", - holdout_id="s1::r1", - validation_strategy="heldout_fold_items", - val_ratio=0.01, - seed=11, - ), - ) - assert labels.count("val") == 1 - - -def test_loocv_patch_validation_allows_zero_val_ratio() -> None: - pytest.importorskip("lightning") - pytest.importorskip("torch_geometric") - from grass_mil.data.components.precompute import LoocvConfig, _assign_split_labels - - units = [ - _make_unit(0, "s0", "r0", 0), - _make_unit(0, "s0", "r0", 1), - _make_unit(1, "s1", "r1", 0), - _make_unit(1, "s1", "r1", 1), - _make_unit(2, "s2", "r2", 0), - _make_unit(2, "s2", "r2", 1), - ] - labels = _assign_split_labels( - units=units, - split_by="region", - split_ratios=(0.8, 0.1, 0.1), - seed=5, - loocv_config=LoocvConfig( - enabled=True, - fold_unit="region", - holdout_id="s2::r2", - validation_strategy="patches_from_train_items", - val_ratio=0.0, - seed=7, - ), - ) - assert "val" not in labels - - -def test_loocv_patch_validation_forces_one_for_nonzero_ratio() -> None: - pytest.importorskip("lightning") - pytest.importorskip("torch_geometric") - from grass_mil.data.components.precompute import LoocvConfig, _assign_split_labels - - units = [ - _make_unit(0, "s0", "r0", 0), - _make_unit(1, "s1", "r1", 0), - _make_unit(2, "s2", "r2", 0), - ] - labels = _assign_split_labels( - units=units, - split_by="region", - split_ratios=(0.8, 0.1, 0.1), - seed=5, - loocv_config=LoocvConfig( - enabled=True, - fold_unit="region", - holdout_id="s2::r2", - validation_strategy="patches_from_train_items", - val_ratio=0.01, - seed=7, - ), - ) - assert labels.count("val") == 1 - - -def test_loocv_supports_sample_fold_unit() -> None: - from grass_mil.data.components.precompute import LoocvConfig, _assign_split_labels - - units = [ - _make_unit(0, "s0", "r0", 0), - _make_unit(0, "s0", "r1", 0), - _make_unit(1, "s1", "r0", 0), - _make_unit(1, "s1", "r1", 0), - _make_unit(2, "s2", "r0", 0), - _make_unit(2, "s2", "r1", 0), - ] - labels = _assign_split_labels( - units=units, - split_by="sample", - split_ratios=(0.8, 0.1, 0.1), - seed=9, - loocv_config=LoocvConfig( - enabled=True, - fold_unit="sample", - holdout_id="s1", - validation_strategy="heldout_fold_items", - val_ratio=0.5, - seed=13, - ), - ) - s1_labels = {label for unit, label in zip(units, labels) if unit.sample_id == "s1"} - assert s1_labels == {"test"} - - -def _build_loocv_train_only_dm(tmp_path: Path, processed_dir: Path): - manifest_path = tmp_path / "manifest.csv" - manifest_path.write_text("sample_id,input_path,input_type\ns0,/tmp/fake.csv,csv\n") - - from grass_mil.data.spatial_omics_datamodule import SpatialOmicsDataModule - - return SpatialOmicsDataModule( - data_dir=str(tmp_path), - raw_manifest_path=str(manifest_path), - processed_dir=str(processed_dir), - coord_scale_um=1.0, - sample_unit="tile", - batch_size=1, - num_workers=0, - pin_memory=False, - reducer_scope="dataset", - keep_raw_molecular=False, - force_precompute=False, - min_cells=1, - use_molecular_features=False, - categorical_features={"include_labels": []}, - split={ - "train_val_test_split": [0.8, 0.1, 0.1], - "split_by": "sample", - "loocv": { - "enabled": True, - "fold_unit": "sample", - "holdout_id": "s0", - "validation_strategy": "patches_from_train_items", - "val_ratio": 0.2, - "seed": 42, - }, - }, - csv={ - "sep": ",", - "coord_columns": ["x", "y"], - "cell_id_column": "cell_id", - "categorical_label_columns": [], - "molecular_columns": None, - }, - h5ad={ - "coord_columns": ["x", "y"], - "cell_id_column": None, - "categorical_label_columns": [], - "molecular_layer": None, - }, - sce={ - "assay_name": None, - "coord_source": "colData", - "coord_key": None, - "coord_columns": ["x", "y"], - "cell_id_column": None, - "categorical_label_columns": [], - "transpose_assay": True, - }, - graph_builder={"name": "delaunay", "kwargs": {}}, - feature_reducer={ - "name": "identity", - "fit_mode": "train_only", - "kwargs": {}, - }, - tiling={"tile_size_um": 50.0, "stride_um": 50.0, "min_cells": 1}, - manifest=None, - graph_labels=None, - transforms=[], - ) - - -def test_loocv_train_only_reducer_rejects_existing_processed_index( - tmp_path: Path, -) -> None: - processed_dir = tmp_path / "processed" - processed_dir.mkdir(parents=True) - (processed_dir / "processed_index.json").write_text("{}") - dm = _build_loocv_train_only_dm(tmp_path=tmp_path, processed_dir=processed_dir) - - with pytest.raises(ValueError, match="cannot reuse an existing processed_index"): - dm.prepare_data() - - -def test_loocv_train_only_reducer_allows_fresh_processed_dir_without_force( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - processed_dir = tmp_path / "processed_fresh" - dm = _build_loocv_train_only_dm(tmp_path=tmp_path, processed_dir=processed_dir) - - calls = {"precompute": 0} - - class _StubPreprocessor: - def precompute(self) -> None: - calls["precompute"] += 1 - - monkeypatch.setattr(dm, "_build_preprocessor", lambda: _StubPreprocessor()) - dm.prepare_data() - assert calls["precompute"] == 1 - - -def _write_minimal_csv_and_manifest( - *, - raw_dir: Path, - coords: np.ndarray, - polygon_path: Path | None = None, -) -> Path: - csv_path = raw_dir / "sample.csv" - np.savetxt( - csv_path, - np.column_stack([np.arange(coords.shape[0]), coords[:, 0], coords[:, 1]]), - delimiter=",", - fmt="%s", - header="cell_id,x,y", - comments="", - ) - manifest_path = raw_dir / "manifest.csv" - polygons_field = str(polygon_path) if polygon_path is not None else "" - manifest_path.write_text( - "sample_id,input_path,input_type,region_id,polygons_path\n" - f"sample_1,{csv_path},csv,region_1,{polygons_field}\n" - ) - return manifest_path - - -def _build_minimal_datamodule( - *, - data_dir: Path, - manifest_path: Path, - processed_dir: Path, - coord_scale_um: float, - force_precompute: bool, - edge_features: list[str], -): - from grass_mil.data.spatial_omics_datamodule import SpatialOmicsDataModule - - return SpatialOmicsDataModule( - data_dir=str(data_dir), - raw_manifest_path=str(manifest_path), - processed_dir=str(processed_dir), - coord_scale_um=coord_scale_um, - sample_unit="full", - batch_size=1, - num_workers=0, - pin_memory=False, - reducer_scope="sample", - keep_raw_molecular=False, - force_precompute=force_precompute, - min_cells=1, - use_molecular_features=False, - categorical_features={"include_labels": []}, - split={"train_val_test_split": [1.0, 0.0, 0.0], "split_by": "patch"}, - csv={ - "sep": ",", - "coord_columns": ["x", "y"], - "cell_id_column": "cell_id", - "categorical_label_columns": [], - "molecular_columns": None, - }, - h5ad={ - "coord_columns": ["x", "y"], - "cell_id_column": None, - "categorical_label_columns": [], - "molecular_layer": None, - }, - sce={ - "assay_name": None, - "coord_source": "colData", - "coord_key": None, - "coord_columns": ["x", "y"], - "cell_id_column": None, - "categorical_label_columns": [], - "transpose_assay": True, - }, - graph_builder={"name": "delaunay", "kwargs": {"edge_features": edge_features}}, - feature_reducer={"name": "identity", "kwargs": {}}, - tiling={"tile_size_um": 50.0, "stride_um": 50.0, "min_cells": 1}, - manifest={ - "sample_id": "sample_id", - "input_path": "input_path", - "input_type": "input_type", - "region_id": "region_id", - "polygons_path": "polygons_path", - }, - graph_labels=None, - transforms=[], - ) - - -def test_coord_scale_um_scales_positions_and_distances(tmp_path: Path) -> None: - pytest.importorskip("torch_geometric") - pytest.importorskip("scipy") - - data_dir = tmp_path / "data" - raw_dir = data_dir / "raw" - raw_dir.mkdir(parents=True) - - coords = np.array([[0.0, 0.0], [3.0, 4.0]]) - manifest_path = _write_minimal_csv_and_manifest(raw_dir=raw_dir, coords=coords) - - dm = _build_minimal_datamodule( - data_dir=data_dir, - manifest_path=manifest_path, - processed_dir=data_dir / "processed", - coord_scale_um=2.0, - force_precompute=True, - edge_features=["distance"], - ) - dm.prepare_data() - dm.setup() - data = dm.train_dataloader().dataset[0] - - assert np.allclose(data.pos.detach().cpu().numpy(), np.array([[0.0, 0.0], [6.0, 8.0]])) - assert list(getattr(data, "edge_attr_names", [])) == ["distance"] - assert np.allclose(data.edge_attr.detach().cpu().numpy(), np.full((2, 1), 10.0)) - - -def test_coord_scale_um_scales_polygons_with_coordinates(tmp_path: Path) -> None: - pytest.importorskip("torch_geometric") - pytest.importorskip("scipy") - - data_dir = tmp_path / "data" - raw_dir = data_dir / "raw" - raw_dir.mkdir(parents=True) - - coords = np.array([[1.0, 1.0], [2.0, 2.0]]) - polygon_path = raw_dir / "polygons.wkt" - _write_polygon_wkt( - polygon_path, - np.array([[0, 0], [0, 3], [3, 3], [3, 0], [0, 0]], dtype=float), - ) - manifest_path = _write_minimal_csv_and_manifest( - raw_dir=raw_dir, - coords=coords, - polygon_path=polygon_path, - ) - - dm = _build_minimal_datamodule( - data_dir=data_dir, - manifest_path=manifest_path, - processed_dir=data_dir / "processed", - coord_scale_um=2.0, - force_precompute=True, - edge_features=["distance"], - ) - dm.prepare_data() - dm.setup() - data = dm.train_dataloader().dataset[0] - - assert int(data.num_nodes) == 2 - - -def test_precompute_reuse_rejects_coord_scale_um_mismatch(tmp_path: Path) -> None: - pytest.importorskip("torch_geometric") - pytest.importorskip("scipy") - - data_dir = tmp_path / "data" - raw_dir = data_dir / "raw" - raw_dir.mkdir(parents=True) - processed_dir = data_dir / "processed" - - coords = np.array([[0.0, 0.0], [3.0, 4.0]]) - manifest_path = _write_minimal_csv_and_manifest(raw_dir=raw_dir, coords=coords) - - dm_first = _build_minimal_datamodule( - data_dir=data_dir, - manifest_path=manifest_path, - processed_dir=processed_dir, - coord_scale_um=1.0, - force_precompute=True, - edge_features=["distance"], - ) - dm_first.prepare_data() - first_index = json.loads((processed_dir / "processed_index.json").read_text()) - assert float(first_index["preprocessing"]["coord_scale_um"]) == 1.0 - - dm_second = _build_minimal_datamodule( - data_dir=data_dir, - manifest_path=manifest_path, - processed_dir=processed_dir, - coord_scale_um=2.0, - force_precompute=False, - edge_features=["distance"], - ) - with pytest.raises(ValueError, match="coord_scale_um=1.0"): - dm_second.prepare_data() - - dm_third = _build_minimal_datamodule( - data_dir=data_dir, - manifest_path=manifest_path, - processed_dir=processed_dir, - coord_scale_um=2.0, - force_precompute=True, - edge_features=["distance"], - ) - dm_third.prepare_data() - third_index = json.loads((processed_dir / "processed_index.json").read_text()) - assert float(third_index["preprocessing"]["coord_scale_um"]) == 2.0 - - -def test_graph_dataset_migrates_legacy_categorical_index(tmp_path): - """Processed caches written before the rename must still load.""" - import json - - import torch - from torch_geometric.data import Data - - from grass_mil.data.components.datasets import SpatialOmicsGraphDataset - - graph_path = tmp_path / "legacy.pt" - data = Data(x=torch.randn(3, 2), edge_index=torch.empty(2, 0, dtype=torch.long)) - data.num_nodes = 3 - data.categorical_index = torch.tensor([[0], [1], [2]], dtype=torch.long) - torch.save(data, graph_path) - - index_path = tmp_path / "processed_index.json" - index_path.write_text( - json.dumps( - { - "entries": [ - { - "path": str(graph_path), - "sample_id": "s0", - "region_id": "r0", - "patch_id": "p0", - "split": "train", - } - ] - } - ) - ) - - loaded = SpatialOmicsGraphDataset(index_path)[0] - assert torch.equal(loaded.categorical_codes.view(-1), torch.tensor([0, 1, 2])) - assert getattr(loaded, "categorical_index", None) is None - - -def _packaged_data_defaults(): - """Reuse the shipped datamodule defaults so tests track config changes.""" - from omegaconf import OmegaConf - - from tests.helpers.config_paths import CONFIGS_DIR - - cfg = OmegaConf.load(CONFIGS_DIR / "data" / "spatial_omics.yaml") - return OmegaConf.to_container(cfg, resolve=False) - - -def _write_cell_type_csv_and_manifest(raw_dir: Path, *, n_samples: int = 4) -> Path: - """A cohort whose only node attribute is a categorical cell type.""" - rng = np.random.default_rng(0) - rows = ["sample_id,input_path,input_type,region_id"] - for sample in range(n_samples): - coords = rng.uniform(0.0, 100.0, size=(24, 2)) - cell_types = rng.choice(["tumor", "tcell", "stroma"], size=coords.shape[0]) - csv_path = raw_dir / f"sample_{sample}.csv" - lines = ["cell_id,x,y,cell_type"] - for idx, ((cx, cy), ctype) in enumerate(zip(coords, cell_types)): - lines.append(f"{idx},{cx},{cy},{ctype}") - csv_path.write_text("\n".join(lines) + "\n") - rows.append(f"sample_{sample},{csv_path},csv,region_{sample}") - manifest_path = raw_dir / "manifest.csv" - manifest_path.write_text("\n".join(rows) + "\n") - return manifest_path - - -def test_encoder_trains_on_cell_type_only_cohort(tmp_path): - """End-to-end guard for the published input scheme. - - With `use_molecular_features=false` the node feature matrix has zero - columns, so the learned cell-type embedding is the entire input - representation. This configuration was inexpressible before the encoder - gained a categorical embedding. - """ - import lightning as L - - from grass_mil.data.spatial_omics_datamodule import SpatialOmicsDataModule - from grass_mil.models.supervised_module import SupervisedModule - - raw_dir = tmp_path / "raw" - raw_dir.mkdir() - manifest_path = _write_cell_type_csv_and_manifest(raw_dir) - - targets_path = tmp_path / "targets.csv" - targets_path.write_text("id,label\n" + "".join(f"region_{i},{i % 2}\n" for i in range(4))) - - datamodule = SpatialOmicsDataModule( - data_dir=str(tmp_path), - raw_manifest_path=str(manifest_path), - processed_dir=str(tmp_path / "processed"), - sample_unit="full", - batch_size=1, - num_workers=0, - pin_memory=False, - min_cells=1, - use_molecular_features=False, - force_precompute=True, - coord_scale_um=1.0, - reducer_scope="sample", - keep_raw_molecular=False, - h5ad=_packaged_data_defaults()["h5ad"], - sce=_packaged_data_defaults()["sce"], - feature_reducer={"name": "identity", "fit_mode": "train_only", "kwargs": {}}, - tiling={"tile_size_um": 200.0, "stride_um": 200.0, "min_cells": 1}, - categorical_features={"include_labels": ["cell_type"]}, - csv={ - "sep": ",", - "coord_columns": ["x", "y"], - "cell_id_column": "cell_id", - "categorical_label_columns": ["cell_type"], - "molecular_columns": None, - }, - split={"train_val_test_split": [0.5, 0.25, 0.25], "split_by": "sample"}, - graph_builder={ - "name": "delaunay", - "kwargs": {"edge_features": ["distance", "neighbor"], "neighbor_cutoff_um": 20.0}, - }, - graph_labels={ - "label_file": str(targets_path), - "id_column": "id", - "tasks": ["label"], - "scope": "region", - }, - ) - datamodule.prepare_data() - datamodule.setup("fit") - - sample = datamodule.dataset_train[0] - assert sample.x.shape[1] == 0, "cell-type-only cohort must have zero continuous features" - assert sample.categorical_codes.shape[0] == sample.x.shape[0] - - model = SupervisedModule( - encoder={ - "input_dim": 0, - "hidden_dim": 8, - "out_dim": 8, - "num_layers": 2, - "conv_type": "gine", - "norm": "layernorm", - "pooling": "max", - "use_edge_attr": True, - "edge_attr_dim": 1, - "edge_feature_index": 0, - "categorical_embedding": {"label": "cell_type"}, - }, - graph_head={"input_dim": 8, "output_dim": 1, "hidden_dim": 8, "num_layers": 2}, - attention=None, - loss={"loss_type": "categorical_bce"}, - task={"aggregation": "mean", "target_type": "binary", "bag_key": "region_id"}, - optim={"_target_": "torch.optim.AdamW", "lr": 1e-3}, - ) - - trainer = L.Trainer( - fast_dev_run=True, accelerator="cpu", logger=False, enable_checkpointing=False - ) - trainer.fit(model, datamodule=datamodule) - - # The vocabulary was inferred from the data, plus the reserved unassigned row. - assert model.encoder.input_proj.embedding.num_embeddings == 4 diff --git a/tests/test_train.py b/tests/test_train.py deleted file mode 100644 index dfad5b1..0000000 --- a/tests/test_train.py +++ /dev/null @@ -1,202 +0,0 @@ -import os -import socket -from pathlib import Path - -import pytest -import torch -from hydra.core.hydra_config import HydraConfig -from omegaconf import DictConfig, OmegaConf, open_dict - -from grass_mil.train import ( - _filter_val_monitor_callbacks, - _requires_zero_validation, - train, -) -from tests.helpers.run_if import RunIf - - -def _can_bind_local_port() -> bool: - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - try: - sock.bind(("127.0.0.1", 0)) - except OSError: - return False - finally: - sock.close() - return True - - -def test_train_fast_dev_run(cfg_train: DictConfig) -> None: - """Run for 1 train, val and test step. - - :param cfg_train: A DictConfig containing a valid training configuration. - """ - HydraConfig().set_config(cfg_train) - with open_dict(cfg_train): - cfg_train.trainer.fast_dev_run = True - cfg_train.trainer.accelerator = "cpu" - train(cfg_train) - - -def test_train_test_mode_without_checkpoint_callback(cfg_train: DictConfig) -> None: - """`test=True` training should not fail without checkpoint callback.""" - HydraConfig().set_config(cfg_train) - with open_dict(cfg_train): - cfg_train.trainer.fast_dev_run = True - cfg_train.trainer.accelerator = "cpu" - cfg_train.callbacks = None - cfg_train.test = True - metric_dict, _ = train(cfg_train) - assert "test/loss" in metric_dict - - -def test_requires_zero_validation_detects_loocv_only_when_val_ratio_zero() -> None: - cfg = OmegaConf.create( - { - "data": { - "split": { - "loocv": { - "enabled": True, - "val_ratio": 0.0, - } - } - } - } - ) - assert _requires_zero_validation(cfg) - - cfg.data.split.loocv.val_ratio = 0.01 - assert not _requires_zero_validation(cfg) - - -def test_requires_zero_validation_detects_non_loocv_zero_val_split() -> None: - cfg = OmegaConf.create( - { - "data": { - "split": { - "train_val_test_split": [0.8, 0.0, 0.2], - } - } - } - ) - assert _requires_zero_validation(cfg) - - cfg.data.split.train_val_test_split = [0.8, 0.1, 0.1] - assert not _requires_zero_validation(cfg) - - -def test_filter_val_monitor_callbacks_removes_only_val_monitors() -> None: - callbacks_cfg = OmegaConf.create( - { - "model_checkpoint": { - "_target_": "lightning.pytorch.callbacks.ModelCheckpoint", - "monitor": "val/loss", - }, - "early_stopping": { - "_target_": "lightning.pytorch.callbacks.EarlyStopping", - "monitor": "val/loss", - }, - "model_summary": { - "_target_": "lightning.pytorch.callbacks.ModelSummary", - "max_depth": -1, - }, - } - ) - filtered, removed = _filter_val_monitor_callbacks(callbacks_cfg) - - assert set(removed) == {"model_checkpoint", "early_stopping"} - assert "model_summary" in filtered - assert "model_checkpoint" not in filtered - assert "early_stopping" not in filtered - - -@RunIf(min_gpus=1) -def test_train_fast_dev_run_gpu(cfg_train: DictConfig) -> None: - """Run for 1 train, val and test step on GPU. - - :param cfg_train: A DictConfig containing a valid training configuration. - """ - HydraConfig().set_config(cfg_train) - with open_dict(cfg_train): - cfg_train.trainer.fast_dev_run = True - cfg_train.trainer.accelerator = "gpu" - train(cfg_train) - - -@RunIf(min_gpus=1) -def test_train_epoch_gpu_amp(cfg_train: DictConfig) -> None: - """Train 1 epoch on GPU with mixed-precision. - - :param cfg_train: A DictConfig containing a valid training configuration. - """ - HydraConfig().set_config(cfg_train) - with open_dict(cfg_train): - cfg_train.trainer.max_epochs = 1 - cfg_train.trainer.accelerator = "gpu" - cfg_train.trainer.precision = 16 - train(cfg_train) - - -def test_train_epoch_double_val_loop(cfg_train: DictConfig) -> None: - """Train 1 epoch with validation loop twice per epoch. - - :param cfg_train: A DictConfig containing a valid training configuration. - """ - HydraConfig().set_config(cfg_train) - with open_dict(cfg_train): - cfg_train.trainer.max_epochs = 1 - cfg_train.trainer.val_check_interval = 0.5 - train(cfg_train) - - -def test_train_ddp_sim(cfg_train: DictConfig) -> None: - """Simulate DDP (Distributed Data Parallel) on 2 CPU processes. - - :param cfg_train: A DictConfig containing a valid training configuration. - """ - if not _can_bind_local_port(): - pytest.skip("Local port binding is not permitted in this environment") - HydraConfig().set_config(cfg_train) - with open_dict(cfg_train): - cfg_train.trainer.max_epochs = 2 - cfg_train.trainer.accelerator = "cpu" - cfg_train.trainer.devices = 2 - cfg_train.trainer.strategy = "ddp_spawn" - train(cfg_train) - - -def test_train_resume(tmp_path: Path, cfg_train: DictConfig) -> None: - """Run 1 epoch, finish, and resume for another epoch. - - :param tmp_path: The temporary logging path. - :param cfg_train: A DictConfig containing a valid training configuration. - """ - with open_dict(cfg_train): - cfg_train.trainer.max_epochs = 1 - - HydraConfig().set_config(cfg_train) - metric_dict_1, _ = train(cfg_train) - - files = os.listdir(tmp_path / "checkpoints") - assert "last.ckpt" in files - assert "epoch_000.ckpt" in files - - with open_dict(cfg_train): - cfg_train.ckpt_path = str(tmp_path / "checkpoints" / "last.ckpt") - cfg_train.trainer.max_epochs = 2 - - metric_dict_2, _ = train(cfg_train) - - files = os.listdir(tmp_path / "checkpoints") - assert any(name.startswith("epoch_") and name.endswith(".ckpt") for name in files) - assert "last.ckpt" in files - assert "epoch_002.ckpt" not in files - - assert "train/loss" in metric_dict_1 - assert "val/loss" in metric_dict_1 - assert "train/loss" in metric_dict_2 - assert "val/loss" in metric_dict_2 - assert torch.isfinite(metric_dict_1["train/loss"]) - assert torch.isfinite(metric_dict_1["val/loss"]) - assert torch.isfinite(metric_dict_2["train/loss"]) - assert torch.isfinite(metric_dict_2["val/loss"]) diff --git a/tests/test_train_regimes.py b/tests/test_train_regimes.py deleted file mode 100644 index ad66c5e..0000000 --- a/tests/test_train_regimes.py +++ /dev/null @@ -1,132 +0,0 @@ -from pathlib import Path - -import pytest -from hydra.core.hydra_config import HydraConfig -from omegaconf import OmegaConf, open_dict -from grass_mil.train import train -from tests.helpers.config_paths import CONFIGS_DIR - - -def _load_model_cfg(name: str): - # Model presets carry a defaults list (model/encoder, model/attention groups), - # so they must be composed through Hydra rather than loaded as plain YAML. - from hydra import compose, initialize_config_module - from hydra.core.global_hydra import GlobalHydra - - GlobalHydra.instance().clear() - with initialize_config_module(version_base="1.3", config_module="grass_mil.configs"): - cfg = compose(config_name="train", overrides=[f"model={name}"]) - GlobalHydra.instance().clear() - return cfg.model - - -def _load_task_cfg(name: str): - return OmegaConf.load(CONFIGS_DIR / "task" / f"{name}.yaml").task - - -def _load_task_model_overrides(name: str): - """Model overrides a task preset injects into the global package. - - Hydra applies these during composition; this test composes manually, so it - must apply them too or the head width and the task objective disagree. - """ - cfg = OmegaConf.load(CONFIGS_DIR / "task" / f"{name}.yaml") - return cfg.get("model", None) - - -def _make_pretrain_ckpt(path: Path) -> str: - from grass_mil.models.bgrl_module import BGRLModule - - module = BGRLModule( - encoder={ - "input_dim": 8, - "hidden_dim": 32, - "out_dim": 32, - "num_layers": 2, - "dropout": 0.1, - "conv_type": "gin", - "norm": "batchnorm", - "jk": "last", - "act": "relu", - "pooling": "mean", - }, - ssl={"method": "bgrl", "predictor": {"hidden_size": 64}}, - optim={"_target_": "torch.optim.AdamW", "lr": 1e-3, "weight_decay": 0.0}, - scheduler={ - "_target_": "torch.optim.lr_scheduler.CosineAnnealingLR", - "T_max": 10, - }, - task={"total_steps": 10}, - ) - module.setup("fit") - ckpt_path = path / "pretrain_mock.ckpt" - import torch - - torch.save({"state_dict": module.ssl_model.state_dict()}, ckpt_path) - return str(ckpt_path) - - -@pytest.mark.parametrize( - "task_name,model_name,use_pretrained,advanced_mil", - [ - ("pretrain_bgrl", "bgrl_module", False, False), - ("finetune_mean", "supervised_module", False, False), - ("finetune_mil", "supervised_module", False, False), - ("finetune_mil", "supervised_module", False, True), - ("finetune_mean", "supervised_module", True, False), - ("finetune_mil", "supervised_module", True, False), - ], -) -def test_training_regimes_fast_dev_run( - cfg_train, tmp_path, task_name, model_name, use_pretrained, advanced_mil -): - with open_dict(cfg_train): - cfg_train.data = OmegaConf.create( - { - "_target_": "tests.helpers.synthetic_datamodule.SyntheticBagDataModule", - "batch_size": 4, - "num_workers": 0, - "pin_memory": False, - "input_dim": 8, - } - ) - cfg_train.task = _load_task_cfg(task_name) - cfg_train.model = _load_model_cfg(model_name) - cfg_train.optim = OmegaConf.load(CONFIGS_DIR / "optim" / "adamw.yaml") - scheduler_name = ( - "cosine_step.yaml" if task_name == "pretrain_bgrl" else "cosine_epoch.yaml" - ) - cfg_train.scheduler = OmegaConf.load(CONFIGS_DIR / "scheduler" / scheduler_name) - task_model_overrides = _load_task_model_overrides(task_name) - if task_model_overrides is not None: - cfg_train.model = OmegaConf.merge(cfg_train.model, task_model_overrides) - cfg_train.model.encoder.input_dim = 8 - cfg_train.model.task = cfg_train.task - cfg_train.model.optim = cfg_train.optim - cfg_train.model.scheduler = cfg_train.scheduler - if advanced_mil: - # Region accumulation only; the auxiliary instance and entropy - # terms were removed with the region-CE-only objective. - cfg_train.model.task.region_accumulation.enabled = True - cfg_train.model.task.region_accumulation.hyperbatch_size = 2 - cfg_train.model.task.region_accumulation.flush_on_epoch_end = True - if use_pretrained: - cfg_train.model.init_from_ckpt = _make_pretrain_ckpt(tmp_path) - cfg_train.model.init_strict = False - # Align encoder dims with the mock checkpoint (hidden_dim=32, out_dim=32). - cfg_train.model.encoder.hidden_dim = 32 - cfg_train.model.encoder.out_dim = 32 - if "graph_head" in cfg_train.model: - cfg_train.model.graph_head.input_dim = cfg_train.model.encoder.out_dim - cfg_train.model.graph_head.hidden_dim = cfg_train.model.encoder.out_dim - if "attention" in cfg_train.model: - cfg_train.model.attention.input_dim = cfg_train.model.encoder.out_dim - cfg_train.trainer.fast_dev_run = True - cfg_train.trainer.accelerator = "cpu" - cfg_train.trainer.devices = 1 - cfg_train.train = True - cfg_train.test = False - - HydraConfig().set_config(cfg_train) - metric_dict, _ = train(cfg_train) - assert "train/loss" in metric_dict diff --git a/tests/test_training.py b/tests/test_training.py new file mode 100644 index 0000000..1db1d98 --- /dev/null +++ b/tests/test_training.py @@ -0,0 +1,113 @@ +"""Training runs: the default regime, the other regimes, other target types, pretraining.""" + +from __future__ import annotations + +import json + +import pandas as pd +import pytest +import torch + +from grass_mil.models import SupervisedModule +from grass_mil.train import train +from tests.helpers.cohort import cohort_overrides, compose_cfg + + +def _train(cohort, tmp_path, *overrides): + cfg = compose_cfg("train", cohort_overrides(cohort) + list(overrides), tmp_path) + return cfg, *train(cfg) + + +def test_default_regime_is_finetune_with_a_probe_head(cohort, pretrained, tmp_path): + cfg, metrics, objects = _train(cohort, tmp_path, f"model.init_from_ckpt={pretrained}") + assert cfg.regime_name == "finetune" + model = objects["model"] + assert model.head.is_linear and model.finetune.head_init == "probe" + assert {"val/auroc", "val/loss", "test/auroc"} <= set(metrics) + preds = pd.read_csv(tmp_path / "test_predictions.csv") + assert {"unit_id", "prob_0", "prob_1", "label", "patient_id"} <= set(preds.columns) + rows = [ + json.loads(line) + for line in (tmp_path / "metrics_per_epoch.jsonl").read_text().splitlines() + ] + assert len(rows) == 2 and "val/auroc" in rows[0] + + +def test_attention_only_trains_nothing_but_the_attention(cohort, pretrained, tmp_path): + _, _, objects = _train( + cohort, tmp_path, "regime=attention_only", f"model.init_from_ckpt={pretrained}" + ) + model = objects["model"] + ckpt = torch.load(pretrained, weights_only=False)["state_dict"] + for name, value in model.encoder.state_dict().items(): + assert torch.equal(value.cpu(), ckpt[f"encoder.{name}"]) + assert not any(p.requires_grad for p in model.head.parameters()) + + +@pytest.mark.parametrize("regime", ["finetune_reg", "mlp_head", "mlp_head_frozen"]) +def test_other_regimes_run(cohort, pretrained, tmp_path, regime): + cfg, metrics, objects = _train( + cohort, + tmp_path, + f"regime={regime}", + f"model.init_from_ckpt={pretrained}", + "trainer.max_epochs=1", + ) + assert cfg.regime_name == regime + assert "test/auroc" in metrics + assert objects["model"].head.is_linear == (regime == "finetune_reg") + + +def test_regression_and_survival_train_and_select_on_their_metric(cohort, tmp_path): + _, metrics, _ = _train( + cohort, + tmp_path / "reg", + "task=regression", + "task.target_columns=[score]", + "trainer.max_epochs=1", + ) + assert "val/rmse" in metrics and "test/rmse" in metrics + cfg, metrics, objects = _train( + cohort, + tmp_path / "surv", + "task=survival", + "task.target_columns=[time,event]", + "trainer.max_epochs=2", + ) + assert cfg.callbacks.model_checkpoint.monitor == "val/c_index" + assert "test/c_index" in metrics + trained = SupervisedModule.from_checkpoint( + str(tmp_path / "surv" / "checkpoints" / "last.ckpt") + ) + assert objects["model"].trainer.global_step > 0 + assert torch.isfinite(next(trained.head.parameters())).all() + preds = pd.read_csv(tmp_path / "surv" / "test_predictions.csv") + assert {"risk", "time", "event"} <= set(preds.columns) + + +def test_from_checkpoint_reproduces_the_trained_model(cohort, tmp_path): + _, _, objects = _train(cohort, tmp_path, "trainer.max_epochs=1", "finetune.head_init=random") + model = objects["model"].eval() + rebuilt = SupervisedModule.from_checkpoint(str(tmp_path / "checkpoints" / "last.ckpt")) + batch = next(iter(objects["datamodule"].test_dataloader())) + with torch.no_grad(): + assert torch.allclose(model.cpu()(batch), rebuilt(batch), atol=1e-5) + + +def test_training_without_validation_split(cohort, tmp_path): + _, metrics, _ = _train( + cohort, + tmp_path, + "data.split.ratios=[0.8,0.0,0.2]", + "trainer.max_epochs=1", + "finetune.head_init=random", + ) + assert not any(k.startswith("val/") for k in metrics) + assert (tmp_path / "checkpoints" / "last.ckpt").exists() + + +def test_pretraining_checkpoint_initialises_the_encoder(pretrained): + ckpt = torch.load(pretrained, weights_only=False) + assert ckpt["input_spec"]["cell_types"] == ["A", "B", "C", "D"] + assert any(k.startswith("target_encoder.") for k in ckpt["state_dict"]) + assert ckpt["hyper_parameters"]["pretrain"]["total_steps"] == 4 From 543c095af0e4980862cfd8fa3365f8fe54409833 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 15 Oct 2026 09:00:55 +0200 Subject: [PATCH 4/9] Tidy the niche figures: panel layout, ordering, marker sizes, r* in cross-L titles Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NqExFvaq8RhDYBj6RShCMk --- src/grass_mil/niches/figures.py | 66 ++++++++++++++++----------------- 1 file changed, 33 insertions(+), 33 deletions(-) diff --git a/src/grass_mil/niches/figures.py b/src/grass_mil/niches/figures.py index 67ea62b..b7ad2d6 100644 --- a/src/grass_mil/niches/figures.py +++ b/src/grass_mil/niches/figures.py @@ -52,49 +52,45 @@ def niche_panel( return path base = table.attrs.get("base_rate", np.nan) frac = [f"frac_{c}" for c in cell_types if f"frac_{c}" in t.columns] - y = np.arange(len(t))[::-1] - widths = [max(2.0, 0.28 * len(frac)), 1.1, 1.3, 0.9] if frac else [1.1, 1.3, 0.9] + y = np.arange(len(t)) + widths = ([max(2.4, 0.3 * len(frac))] if frac else []) + [1.6, 1.6, 1.1] fig, axes = plt.subplots( 1, len(widths), - figsize=(sum(widths) + 2.5, 0.42 * len(t) + 1.8), - gridspec_kw={"width_ratios": widths}, + figsize=(1.6 * sum(widths) + 1.5, 0.45 * len(t) + 2.0), + gridspec_kw={"width_ratios": widths, "wspace": 0.15}, + sharey=True, ) axes = list(np.atleast_1d(axes)) if frac: ax = axes.pop(0) im = ax.imshow(t[frac].to_numpy() * 100, aspect="auto", cmap="viridis", vmin=0) ax.set_xticks(range(len(frac)), [c[5:] for c in frac], rotation=90, fontsize=7) - ax.set_yticks(range(len(t)), t["niche"], fontsize=8) - plt.colorbar(im, ax=ax, fraction=0.04, label="% of niche cells") - ax.set_title("composition", fontsize=9) + cbar = fig.colorbar(im, ax=ax, location="top", shrink=0.8, pad=0.02) + cbar.set_label("composition (% of niche cells)", fontsize=8) + cbar.ax.tick_params(labelsize=7) ax_assoc, ax_margin, ax_share = axes - ax_assoc.barh( - y, - t["association"], - color=["#b2182b" if v > base else "#2166ac" for v in t["association"]], - height=0.6, - ) + colours = ["#b2182b" if v > base else "#2166ac" for v in t["association"]] + ax_assoc.barh(y, t["association"], color=colours, height=0.6) ax_assoc.axvline(base, color="k", ls=":", lw=1) ax_assoc.set_xlim(0, 1) - ax_assoc.set_xlabel(f"fraction from {positive_name} units") - ax_margin.errorbar( - t["margin"], - y, - xerr=[t["margin"] - t["margin_lo"], t["margin_hi"] - t["margin"]], - fmt="o", - color="#1b9e77", - ms=4, - capsize=2, - ) + ax_assoc.set_title(f"association\n(fraction from {positive_name} units)", fontsize=8) + xerr = [t["margin"] - t["margin_lo"], t["margin_hi"] - t["margin"]] + ax_margin.errorbar(t["margin"], y, xerr=xerr, fmt="o", color="#1b9e77", ms=4, capsize=2) ax_margin.axvline(0, color="k", lw=0.8) - ax_margin.set_xlabel(f"attention-weighted margin (+ = {positive_name})") + ax_margin.set_title(f"attention-weighted margin\n(95% CI; + = {positive_name})", fontsize=8) ax_share.barh(y, 100 * t["share"], color="#7f7f7f", height=0.6) - ax_share.set_xlabel("% of cells") + ax_share.set_title("share\n(% of cells)", fontsize=8) + axes_all = fig.axes + axes_all[0].set_yticks(y, t["niche"], fontsize=8) + if not frac: + axes_all[0].invert_yaxis() # Niche 1 on top (imshow already does this) for ax in (ax_assoc, ax_margin, ax_share): - ax.set_yticks(y, [] if frac else t["niche"]) + ax.tick_params(labelsize=7) ax.spines[["top", "right"]].set_visible(False) - return _save(fig, path) + fig.savefig(path, dpi=150, bbox_inches="tight") + plt.close(fig) + return path def embedding_panel( @@ -109,14 +105,15 @@ def embedding_panel( ) -> Path: """The clustering space (first two PCs): niches | margin | unit class.""" fig, axes = plt.subplots(1, 3, figsize=(15, 4.6)) + size = float(np.clip(20000.0 / max(len(coords), 1), 0.5, 20.0)) colours = niche_colours(sorted(set(niche_names), key=lambda n: (n == "Background", n))) for name, colour in colours.items(): m = niche_names == name - axes[0].scatter(coords[m, 0], coords[m, 1], s=2, color=colour, lw=0, label=name) + axes[0].scatter(coords[m, 0], coords[m, 1], s=size, color=colour, lw=0, label=name) axes[0].legend(markerscale=6, fontsize=6, ncol=2) lim = float(np.nanpercentile(np.abs(margin), 99)) or 1.0 sc = axes[1].scatter( - coords[:, 0], coords[:, 1], s=2, c=margin, cmap="RdBu_r", vmin=-lim, vmax=lim, lw=0 + coords[:, 0], coords[:, 1], s=size, c=margin, cmap="RdBu_r", vmin=-lim, vmax=lim, lw=0 ) plt.colorbar(sc, ax=axes[1], label=f"margin (+ = {positive_name})") hb = axes[2].hexbin( @@ -180,20 +177,21 @@ def unit_maps( """One unit in the tissue: cell types | cell margins | niches.""" fig, axes = plt.subplots(1, 3, figsize=(15, 5)) cmap = plt.get_cmap("tab20") + size = float(np.clip(4000.0 / max(len(pos), 1), 1.0, 30.0)) for j, name in enumerate(type_names): m = cell_types == j if m.any(): - axes[0].scatter(pos[m, 0], pos[m, 1], s=4, color=cmap(j % 20), lw=0, label=name) + axes[0].scatter(pos[m, 0], pos[m, 1], s=size, color=cmap(j % 20), lw=0, label=name) axes[0].legend(markerscale=3, fontsize=6, ncol=2) lim = float(np.nanpercentile(np.abs(margin), 99)) or 1.0 sc = axes[1].scatter( - pos[:, 0], pos[:, 1], s=4, c=margin, cmap="RdBu_r", vmin=-lim, vmax=lim, lw=0 + pos[:, 0], pos[:, 1], s=size, c=margin, cmap="RdBu_r", vmin=-lim, vmax=lim, lw=0 ) plt.colorbar(sc, ax=axes[1], label="margin") colours = niche_colours(sorted(set(niche_names), key=lambda n: (n == "Background", n))) for name, colour in colours.items(): m = niche_names == name - axes[2].scatter(pos[m, 0], pos[m, 1], s=4, color=colour, lw=0, label=name) + axes[2].scatter(pos[m, 0], pos[m, 1], s=size, color=colour, lw=0, label=name) axes[2].legend(markerscale=3, fontsize=6) for ax, sub in zip(axes, ("cell types", "attention-weighted margin", "niches")): ax.set_aspect("equal") @@ -222,7 +220,9 @@ def pair_heatmaps( a, b = index[r["type_a"]], index[r["type_b"]] m[a, b] = m[b, a] = r[value] im = ax.imshow(m, cmap="RdBu_r", vmin=-lim, vmax=lim) - ax.set_title(niche, fontsize=8) + rows_of = table[table["niche"] == niche] + suffix = f" (r* = {rows_of['r_star'].iloc[0]:.3g})" if "r_star" in table else "" + ax.set_title(niche + suffix, fontsize=8) ax.set_xticks(range(k), cell_types, rotation=90, fontsize=5) ax.set_yticks(range(k), cell_types, fontsize=5) for ax in list(axes.flat)[len(niches) :]: From edd1ac820104c535fed80a319e2a0f448cb7bdf9 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 15 Oct 2026 09:52:16 +0200 Subject: [PATCH 5/9] Docs, vignettes, packaging files and CI for 1.1.0 - README: the method, install, quick start, a workflow table, changing the defaults, documentation and citation; no preprint-reproduction section - docs/: data preparation, training and evaluation, niche analysis and a configuration reference replace the sixteen previous pages (including the method-correspondence page and the generated repository map) - vignettes: rewritten and executed on squidpy's MIBI-TOF data (tiles as units, a CD8-infiltration phenotype): pretraining, the default regime, niches, nested cross-validation, replication and model-free statistics; outputs carry no local paths - CI: one test workflow on Python 3.11/3.12 and macOS with torch-sparse, one pre-commit workflow; no release drafter; real coverage thresholds; weekly dependabot; a short PR template - environment.yaml installs the package; Makefile targets test/test-all; CHANGELOG 1.1.0 - spatial-stats compares patients or units (compare_by); niches are numbered by the margin statistic the table reports; random splits stratify only when every patient has one class Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NqExFvaq8RhDYBj6RShCMk --- .github/PULL_REQUEST_TEMPLATE.md | 25 +- .github/codecov.yml | 11 +- .github/dependabot.yml | 19 +- .github/release-drafter.yml | 44 - .github/workflows/code-quality-main.yaml | 27 - .github/workflows/code-quality-pr.yaml | 44 - .github/workflows/code-quality.yml | 18 + .github/workflows/release-drafter.yml | 27 - .github/workflows/test.yml | 149 +- .pre-commit-config.yaml | 2 +- CHANGELOG.md | 54 +- Makefile | 24 +- README.md | 221 +-- docs/capabilities_and_modes.md | 208 --- docs/configuration.md | 155 ++ docs/configuration_guide.md | 223 --- docs/data_preparation.md | 116 ++ docs/dataset_preparation.md | 167 -- docs/evaluation_stats.md | 34 - docs/getting_started.md | 189 --- docs/hyperparameter_reference.md | 907 ----------- docs/inference_utility_suite.md | 161 -- docs/interpretability_extension_guide.md | 200 --- docs/interpretability_suite.md | 760 --------- docs/method_correspondence.md | 52 - docs/model_component_contracts.md | 150 -- docs/nested_cv.md | 88 -- docs/niche_analysis.md | 133 ++ docs/niche_discovery.md | 89 -- docs/repository_map.md | 379 ----- docs/spatial_statistics.md | 25 - docs/training_and_evaluation.md | 116 ++ docs/workflows_train_infer.md | 692 --------- environment.yaml | 62 +- src/grass_mil/configs/spatial_stats.yaml | 5 +- src/grass_mil/data/datamodule.py | 42 +- src/grass_mil/niches/analysis.py | 4 +- src/grass_mil/niches/discovery.py | 19 +- src/grass_mil/niches/replication/stages.py | 2 +- src/grass_mil/spatial/cli.py | 12 +- tests/test_niches.py | 5 +- tools/update_repository_map.py | 152 -- vignettes/01_anndata_to_cell_graphs.ipynb | 492 ++++++ vignettes/01_anndata_to_spatial_graphs.ipynb | 589 ------- vignettes/02_pretraining_and_niches.ipynb | 1069 ------------- .../02_pretraining_training_and_niches.ipynb | 732 +++++++++ .../03_cross_validation_and_replication.ipynb | 937 +++++++++++ ...03_nested_cv_niches_and_metaclusters.ipynb | 1363 ----------------- vignettes/README.md | 53 +- 49 files changed, 2939 insertions(+), 8108 deletions(-) delete mode 100644 .github/release-drafter.yml delete mode 100644 .github/workflows/code-quality-main.yaml delete mode 100644 .github/workflows/code-quality-pr.yaml create mode 100644 .github/workflows/code-quality.yml delete mode 100644 .github/workflows/release-drafter.yml delete mode 100644 docs/capabilities_and_modes.md create mode 100644 docs/configuration.md delete mode 100644 docs/configuration_guide.md create mode 100644 docs/data_preparation.md delete mode 100644 docs/dataset_preparation.md delete mode 100644 docs/evaluation_stats.md delete mode 100644 docs/getting_started.md delete mode 100644 docs/hyperparameter_reference.md delete mode 100644 docs/inference_utility_suite.md delete mode 100644 docs/interpretability_extension_guide.md delete mode 100644 docs/interpretability_suite.md delete mode 100644 docs/method_correspondence.md delete mode 100644 docs/model_component_contracts.md delete mode 100644 docs/nested_cv.md create mode 100644 docs/niche_analysis.md delete mode 100644 docs/niche_discovery.md delete mode 100644 docs/repository_map.md delete mode 100644 docs/spatial_statistics.md create mode 100644 docs/training_and_evaluation.md delete mode 100644 docs/workflows_train_infer.md delete mode 100644 tools/update_repository_map.py create mode 100644 vignettes/01_anndata_to_cell_graphs.ipynb delete mode 100644 vignettes/01_anndata_to_spatial_graphs.ipynb delete mode 100644 vignettes/02_pretraining_and_niches.ipynb create mode 100644 vignettes/02_pretraining_training_and_niches.ipynb create mode 100644 vignettes/03_cross_validation_and_replication.ipynb delete mode 100644 vignettes/03_nested_cv_niches_and_metaclusters.ipynb diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index 410bcd8..9ca25ef 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -1,22 +1,9 @@ -## What does this PR do? +## What does this change? - + -Fixes #\ +## Checklist -## Before submitting - -- [ ] Did you make sure **title is self-explanatory** and **the description concisely explains the PR**? -- [ ] Did you make sure your **PR does only one thing**, instead of bundling different changes together? -- [ ] Did you list all the **breaking changes** introduced by this pull request? -- [ ] Did you **test your PR locally** with `pytest` command? -- [ ] Did you **run pre-commit hooks** with `pre-commit run -a` command? - -## Did you have fun? - -Make sure you had fun coding 🙃 +- [ ] `pytest` passes locally +- [ ] `pre-commit run -a` passes +- [ ] Documentation and `CHANGELOG.md` updated where behaviour changed diff --git a/.github/codecov.yml b/.github/codecov.yml index c66853c..3778238 100644 --- a/.github/codecov.yml +++ b/.github/codecov.yml @@ -1,15 +1,8 @@ coverage: status: - # measures overall project coverage project: default: - threshold: 100% # how much decrease in coverage is needed to not consider success - - # measures PR or single commit coverage + threshold: 2% # allowed drop in total coverage patch: default: - threshold: 100% # how much decrease in coverage is needed to not consider success - - - # project: off - # patch: off + target: 70% diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 5a861fd..ca8e216 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -1,16 +1,13 @@ -# To get started with Dependabot version updates, you'll need to specify which -# package ecosystems to update and where the package manifests are located. -# Please see the documentation for all configuration options: -# https://docs.github.com/github/administering-a-repository/configuration-options-for-dependency-updates - version: 2 updates: - - package-ecosystem: "pip" # See documentation for possible values - directory: "/" # Location of package manifests + - package-ecosystem: "pip" + directory: "/" schedule: - interval: "daily" + interval: "weekly" ignore: - - dependency-name: "pytorch-lightning" - update-types: ["version-update:semver-patch"] - - dependency-name: "torchmetrics" + - dependency-name: "lightning" update-types: ["version-update:semver-patch"] + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "monthly" diff --git a/.github/release-drafter.yml b/.github/release-drafter.yml deleted file mode 100644 index 59af159..0000000 --- a/.github/release-drafter.yml +++ /dev/null @@ -1,44 +0,0 @@ -name-template: "v$RESOLVED_VERSION" -tag-template: "v$RESOLVED_VERSION" - -categories: - - title: "🚀 Features" - labels: - - "feature" - - "enhancement" - - title: "🐛 Bug Fixes" - labels: - - "fix" - - "bugfix" - - "bug" - - title: "🧹 Maintenance" - labels: - - "maintenance" - - "dependencies" - - "refactoring" - - "cosmetic" - - "chore" - - title: "📝️ Documentation" - labels: - - "documentation" - - "docs" - -change-template: "- $TITLE @$AUTHOR (#$NUMBER)" -change-title-escapes: '\<*_&' # You can add # and @ to disable mentions - -version-resolver: - major: - labels: - - "major" - minor: - labels: - - "minor" - patch: - labels: - - "patch" - default: patch - -template: | - ## Changes - - $CHANGES diff --git a/.github/workflows/code-quality-main.yaml b/.github/workflows/code-quality-main.yaml deleted file mode 100644 index 925ef03..0000000 --- a/.github/workflows/code-quality-main.yaml +++ /dev/null @@ -1,27 +0,0 @@ -# Same as `code-quality-pr.yaml` but triggered on commit to main branch -# and runs on all files (instead of only the changed ones) - -name: Code Quality Main - -on: - push: - branches: [main] - -jobs: - code-quality: - runs-on: ubuntu-latest - - steps: - - name: Checkout - uses: actions/checkout@v5 - - - name: Set up Python - uses: actions/setup-python@v6 - with: - python-version: "3.12" - - - name: Install pre-commit - run: python -m pip install --upgrade pip pre-commit - - - name: Run pre-commit on all files - run: pre-commit run --all-files --show-diff-on-failure diff --git a/.github/workflows/code-quality-pr.yaml b/.github/workflows/code-quality-pr.yaml deleted file mode 100644 index 76161a6..0000000 --- a/.github/workflows/code-quality-pr.yaml +++ /dev/null @@ -1,44 +0,0 @@ -# This workflow finds which files were changed, prints them, -# and runs `pre-commit` on those files. - -# Inspired by the sktime library: -# https://github.com/alan-turing-institute/sktime/blob/main/.github/workflows/test.yml - -name: Code Quality PR - -on: - pull_request: - branches: [main, "release/*", "dev"] - -jobs: - code-quality: - runs-on: ubuntu-latest - - steps: - - name: Checkout - uses: actions/checkout@v5 - - - name: Set up Python - uses: actions/setup-python@v6 - with: - python-version: "3.12" - - - name: Find modified files - id: file_changes - uses: trilom/file-changes-action@v1.2.4 - with: - output: " " - - - name: List modified files - run: echo '${{ steps.file_changes.outputs.files}}' - - - name: Install pre-commit - run: python -m pip install --upgrade pip pre-commit - - - name: Run pre-commit on changed files - run: | - if [ -z "${{ steps.file_changes.outputs.files }}" ]; then - echo "No changed files detected by file-changes action; skipping pre-commit." - exit 0 - fi - pre-commit run --files ${{ steps.file_changes.outputs.files }} --show-diff-on-failure diff --git a/.github/workflows/code-quality.yml b/.github/workflows/code-quality.yml new file mode 100644 index 0000000..7f00ed4 --- /dev/null +++ b/.github/workflows/code-quality.yml @@ -0,0 +1,18 @@ +name: Code quality + +on: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + pre-commit: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v5 + - uses: actions/setup-python@v6 + with: + python-version: "3.12" + - run: python -m pip install --upgrade pip pre-commit + - run: pre-commit run --all-files --show-diff-on-failure diff --git a/.github/workflows/release-drafter.yml b/.github/workflows/release-drafter.yml deleted file mode 100644 index 6a45e15..0000000 --- a/.github/workflows/release-drafter.yml +++ /dev/null @@ -1,27 +0,0 @@ -name: Release Drafter - -on: - push: - # branches to consider in the event; optional, defaults to all - branches: - - main - -permissions: - contents: read - -jobs: - update_release_draft: - permissions: - # write permission is required to create a github release - contents: write - # write permission is required for autolabeler - # otherwise, read permission is required at least - pull-requests: write - - runs-on: ubuntu-latest - - steps: - # Drafts your next Release notes as Pull Requests are merged into "master" - - uses: release-drafter/release-drafter@v5 - env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index dda0175..ac8acf3 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -4,42 +4,32 @@ on: push: branches: [main] pull_request: - branches: [main, "release/*", "dev"] + branches: [main] env: - # Bump these together. torch-scatter/torch-sparse are only published for a - # subset of torch releases: https://data.pyg.org/whl/ has no wheels beyond - # torch 2.12, so raising TORCH_VERSION past it means those extensions would - # have to build from source, which fails on macOS (PyTorch's own - # c10/util/strong_type.h specializes std::is_arithmetic, which recent libc++ - # rejects). Check the index before changing this. - TORCH_VERSION: "2.12.0" - PYTHON_VERSION: "3.12" + # torch-sparse is published at https://data.pyg.org/whl/ for a subset of + # torch releases only; check the index before raising this. + TORCH_VERSION: "2.4.1" jobs: - run-tests: - name: Pytest (${{ matrix.os }}) + test: + name: Pytest (${{ matrix.os }}, Python ${{ matrix.python }}) runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: - os: [ubuntu-latest, macos-latest] + include: + - { os: ubuntu-latest, python: "3.11" } + - { os: ubuntu-latest, python: "3.12" } + - { os: macos-latest, python: "3.12" } timeout-minutes: 45 - steps: - - name: Checkout - uses: actions/checkout@v5 - - - name: Set up Python - uses: actions/setup-python@v6 + - uses: actions/checkout@v5 + - uses: actions/setup-python@v6 with: - python-version: ${{ env.PYTHON_VERSION }} + python-version: ${{ matrix.python }} cache: pip - - # Linux resolves CUDA builds by default, which are ~3GB and useless on a - # CPU runner. macOS wheels on PyPI are already CPU/MPS only, and the - # download.pytorch.org cpu index does not publish macOS builds. - - name: Install PyTorch (CPU) + - name: Install PyTorch (CPU) and torch-sparse run: | python -m pip install --upgrade pip if [ "$RUNNER_OS" = "Linux" ]; then @@ -47,111 +37,12 @@ jobs: else pip install "torch==${TORCH_VERSION}" fi - - # Prebuilt binaries. Never let these build from source; see the note on - # TORCH_VERSION above. - - name: Install PyG compiled extensions - run: | - pip install torch-scatter torch-sparse \ - --only-binary :all: \ + pip install torch-scatter torch-sparse --only-binary :all: \ -f "https://data.pyg.org/whl/torch-${TORCH_VERSION}+cpu.html" - - name: Install the package - run: pip install ".[interpretability,io,stats,plots]" pytest - - - name: Show key package versions - run: | - python --version - python - <<'PY' - import importlib - - checks = [ - "torch", - "torch_geometric", - "torch_scatter", - "torch_sparse", - "lightning", - "pytorch_lightning", - "torchmetrics", - "omegaconf", - "hydra", - "numpy", - "pandas", - "scipy", - "sklearn", - "plotly", - ] - for name in checks: - mod = importlib.import_module(name) - print(f"{name}=={getattr(mod, '__version__', 'unknown')}") - PY - - - name: Run pytest - run: pytest -v - - code-coverage: - runs-on: ubuntu-latest - timeout-minutes: 45 - - steps: - - name: Checkout - uses: actions/checkout@v5 - - - name: Set up Python - uses: actions/setup-python@v6 - with: - python-version: ${{ env.PYTHON_VERSION }} - cache: pip - - - name: Install PyTorch (CPU) and PyG extensions - run: | - python -m pip install --upgrade pip - pip install "torch==${TORCH_VERSION}" --index-url https://download.pytorch.org/whl/cpu - pip install torch-scatter torch-sparse \ - --only-binary :all: \ - -f "https://data.pyg.org/whl/torch-${TORCH_VERSION}+cpu.html" - - - name: Run tests and collect coverage - run: | - pip install ".[interpretability,io,stats,plots]" pytest "pytest-cov[toml]" - pytest --cov grass_mil - - - name: Upload coverage to Codecov + run: pip install ".[io,test]" + - name: Run the tests + run: pytest --cov grass_mil --cov-report xml + - name: Upload coverage + if: matrix.os == 'ubuntu-latest' && matrix.python == '3.12' uses: codecov/codecov-action@v5 - - packaging: - name: Install from source and run the CLI - runs-on: ubuntu-latest - timeout-minutes: 30 - - steps: - - name: Checkout - uses: actions/checkout@v5 - - # Deliberately the oldest supported Python, so `requires-python` is - # exercised at its floor rather than only at the version CI happens to use. - - name: Set up Python - uses: actions/setup-python@v6 - with: - python-version: "3.11" - cache: pip - - # A plain, unpinned pip install: this is the path a downstream user takes, - # and it is the one that silently rots when the package layout or the - # config tree moves. It deliberately does NOT pin torch, so a dependency - # that breaks on the latest release shows up here first. - - name: Install the package - run: pip install . - - - name: Entrypoints resolve packaged configs from an unrelated directory - run: | - mkdir -p "${RUNNER_TEMP}/elsewhere" - cd "${RUNNER_TEMP}/elsewhere" - grass-mil-train --help - grass-mil-report --help - python -c "import grass_mil, grass_mil.configs; print(grass_mil.__file__)" - - - name: Packaging contract tests - run: | - pip install pytest - pytest -v tests/test_packaging.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 72b1c6c..76b0314 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -14,7 +14,7 @@ repos: - id: check-added-large-files - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.6.9 + rev: v0.8.6 hooks: - id: ruff args: [--fix, --exit-non-zero-on-fix] diff --git a/CHANGELOG.md b/CHANGELOG.md index 3fa2e37..61cdf15 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,30 +1,32 @@ # Changelog -## 1.0.0 (2026-10) - -Public release accompanying the preprint -[GRASS-MIL: graph-based representation and discovery of phenotype-associated spatial structures with multiple instance learning](https://www.biorxiv.org/content/10.64898/2026.10.08.757722v1). - -### Method (defaults match the preprint; every setting stays configurable) - -- Delaunay cell graph with the bounded log edge length `clip((ln d - 2)/3, 0, 1)` as the edge feature (`log_distance`). -- GINE ego-graph encoder (4 x 256, learnable epsilon, layer norm, max pooling, xavier-initialised cell-type embedding) as `model=grass_mil`. -- Projected gated attention with one channel per class (32 / 8, kaiming) and logit-space pooling, so every bag logit is an exact sum of per-cell contributions. -- One region = one bag: all sampled ego-graphs of a region in one batch, one attention softmax, one loss; region-level validation metrics and selection. -- Training regimes as presets (`regime=mil|mil_frozen|linear_probehead|linear_ftall|linear_ftall_reg`) with per-component learning-rate scales, head pre-fit by a logistic probe, class weights from the training split, per-step cosine schedule, gradient accumulation over regions, AMP-safe scheduler stepping, elementwise feature drop with an unassigned token. -- BGRL "A3" pretraining preset (`regime=pretrain_a3`): roots drawn from the union of all regions, momentum 0.99 -> 1.0, DropImportance + feature noise, checkpoint every 5000 steps. -- Import of legacy SPACE-GM style BGRL checkpoints (`model.encoder_init_map=spacegm_bgrl`). +## 1.1.0 (2026-10) + +A package-level cleanup. Commands, configuration keys and file formats changed; +graphs and checkpoints of 1.0.0 need to be rebuilt. + +- **Regimes renamed**; `finetune` is the default: `linear_ftall` → `finetune`, + `linear_ftall_reg` → `finetune_reg`, `linear_probehead` → `attention_only`, + `mil` → `mlp_head`, `mil_frozen` → `mlp_head_frozen`, `pretrain_a3` → + `pretrain`. The probe head is fitted automatically; `grass-mil-probe` is gone. +- **Niche analysis** (`grass-mil-niches`) reports the primary results directly: + niches in the attention space, composition, association, margin with + intervals, per-niche Moran's I, neighbourhood enrichment, cross-L and contact + enrichment, and optional survival. `grass-mil-metaclusters` is optional. +- **Removed**: `grass-mil-report` and its plugin stack, `grass-mil-eval`, + `grass-mil-loocv`, `grass-mil-scores`, mean pooling, the GCN/GAT/GraphSAGE/GIN + encoders, the 1.0.0 checkpoint-import options and unused logger, debug and + multi-GPU presets. +- **Configuration**: one `data.input` block for all formats, `data.unit` + (`region`, `tile`, `polygon`), labels attached at load time (one set of + graphs for every task and fold), `task` (`classification`, `regression`, + `survival`) with its own selection metric. +- **Fixes**: gradient clipping under mixed precision now acts on unscaled + gradients; cross-validation refits no longer evaluate on the held-out fold; + `grass-mil-predict` reports AUROC and works without `regime=`; the PCA + feature reducer works; survival training optimises over several units per + step; numeric sample ids; stale processed graphs are detected. -### Analyses - -- `grass-mil-probe`, `grass-mil-nested-cv`: nested patient-grouped cross-validation with explicit fold files, inner-fold epoch selection, refit and out-of-fold predictions (region and patient level). -- `grass-mil-extract`: per-cell cache (embeddings, instance logits, attention) of trained models over every cell. -- `grass-mil-niches`: niche discovery on a model's cells (cosine HDBSCAN by default, any registry backend), attention-weighted margins, per-niche tables, co-localisation, figures. -- `grass-mil-metaclusters`: cross-fold niche replication (constrained matching on composition), whole-dataset transfer, shared label-free basis. -- `grass-mil-scores`: bootstrap AUROC intervals, DeLong tests, within-fold standardisation, operating points (sensitivity, specificity, PPV, NPV, likelihood ratios). -- `grass-mil-spatial-stats`: neighbourhood enrichment, cellular neighbourhoods, contact fractions and cross-K per region with patient-level group tests. -- Patient-level aggregation in `grass-mil-predict` (mean of region probabilities or one merged bag). - -### Housekeeping +## 1.0.0 (2026-10) -- Removed stale component configs, the dummy dataset, the embeddings shim and internal audit notes; added `CITATION.cff`; `hydra-core` 1.3.5+. +Public release accompanying the preprint. diff --git a/Makefile b/Makefile index 801692b..7ce5ae7 100644 --- a/Makefile +++ b/Makefile @@ -1,19 +1,15 @@ -help: ## Show help - @grep -E '^[.a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-30s\033[0m %s\n", $$1, $$2}' +help: ## Show the targets + @grep -E '^[a-z-]+:.*?## ' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-12s\033[0m %s\n", $$1, $$2}' -clean: ## Remove caches and build artefacts - rm -rf dist - find . -type f -name "*.DS_Store" -ls -delete - find . | grep -E "(__pycache__|\.pyc|\.pyo)" | xargs rm -rf - find . | grep -E ".pytest_cache" | xargs rm -rf - find . | grep -E ".ipynb_checkpoints" | xargs rm -rf - rm -f .coverage +test: ## Run the tests (fast ones only) + pytest -m "not slow" -clean-logs: ## Remove run outputs under logs/ - rm -rf logs/** +test-all: ## Run every test, including the end-to-end workflow + pytest -format: ## Run pre-commit hooks (ruff lint + format) +format: ## Lint and format (pre-commit hooks) pre-commit run -a -test: ## Run the test suite - pytest +clean: ## Remove caches and build artefacts + rm -rf dist build .pytest_cache .coverage + find . -name __pycache__ -type d -prune -exec rm -rf {} + diff --git a/README.md b/README.md index 65c0f5e..02556f4 100644 --- a/README.md +++ b/README.md @@ -7,40 +7,30 @@ GRASS-MIL predicts a patient- or region-level phenotype from cell-resolved spatial omics data (multiplexed imaging, spatial transcriptomics with cell -types) and tells you *which spatial structures* carry the prediction. It +types) and shows *which spatial structures* carry the prediction. It represents every cell by its spatial neighbourhood, learns those representations without labels, pools them with per-class attention into a -bag prediction whose every logit is an exact sum of per-cell contributions, -and turns the contributions into **niches**: recurrent cellular -neighbourhoods with an association to the phenotype and an attribution by -the model, checked for replication across cross-validation folds. - -This repository is the public implementation accompanying the preprint: - -> Rabuzin L, Yates J, Xu M, Boeva V. *GRASS-MIL: graph-based representation -> and discovery of phenotype-associated spatial structures with multiple -> instance learning.* bioRxiv (2026). -> doi: [10.64898/2026.10.08.757722](https://www.biorxiv.org/content/10.64898/2026.10.08.757722v1) +prediction whose every logit is an exact sum of per-cell contributions, and +turns the contributions into **niches**: recurrent cellular neighbourhoods, +each with its composition, its association with the phenotype, the model's +attribution to it and its spatial organisation. ## The method in one paragraph -Cells of a region form a Delaunay graph whose edges carry the bounded log -edge length. A GINE graph network encodes the *k*-hop ego-graph around every -cell (4 hops, width 256 by default) and is pretrained with BGRL on all regions, -label-free. For prediction, the ego-graphs of a region are the instances of -one bag: an instance head gives per-class logits, a projected gated attention -net gives per-class weights, and the bag logit is `L_c = sum_i A[i,c] * l[i,c]` -(logit-space pooling), so every cell's contribution `A[i,c] (l[i,c] - b_c)` is -exact rather than approximated. Fine-tuning regimes range from attention-only -on a frozen encoder with a probe-fitted linear head to end-to-end fine-tuning. -Evaluation is a nested, patient-grouped cross-validation. Interpretation -clusters the held-out cells of each fold model in its own embedding space -(cosine HDBSCAN by default), scores every cluster by its empirical association -and its attention-weighted margin, and matches clusters across folds into -metaclusters that replicate. Every setting above is a default of a Hydra -config, not a constant: depth, width, ego-graph hops, node and edge features, -attention sizes, learning rates, bag sizes, schedules, augmentations, fold -counts and clustering settings are all overridable from the command line. +Cells of a region form a Delaunay graph whose edges carry the bounded log edge +length. A GINE graph network encodes the *k*-hop ego-graph around every cell +(4 hops, width 256) and is pretrained with BGRL on all regions, label-free. For +prediction, the ego-graphs of a region are the instances of one bag: an +instance head gives per-class logits, a projected gated attention gives +per-class weights, and the bag logit is `L_c = sum_i A[i,c] * l[i,c]`, so every +cell's contribution `A[i,c] (l[i,c] - b_c)` is exact rather than approximated. +By default the pretrained encoder is fine-tuned end to end with a linear +instance head initialised by a probe on its frozen embeddings. Evaluation is a +nested, patient-grouped cross-validation. Interpretation clusters the held-out +cells of a model in its attention space (cosine HDBSCAN) and describes every +niche by its composition, empirical association, attention-weighted margin and +classical spatial statistics within the niche. Every setting is a key of the +[configuration](docs/configuration.md) and can be changed from the command line. ## Install @@ -48,151 +38,88 @@ counts and clustering settings are all overridable from the command line. git clone https://github.com/BoevaLab/grass-mil.git cd grass-mil python -m venv .venv && source .venv/bin/activate -pip install -e ".[interpretability,io,stats,plots]" +pip install -e ".[io]" +# torch-sparse, matching your PyTorch build (ego-graph sampling): +pip install torch-sparse -f https://data.pyg.org/whl/torch-$(python -c "import torch; print(torch.__version__.split('+')[0])").html ``` -`torch-geometric` needs the compiled `torch-sparse` / `torch-scatter` wheels -matching your PyTorch build for the ego-graph sampler (see -[`.github/workflows/test.yml`](.github/workflows/test.yml) for the exact pip -lines used in CI, or use [`environment.yaml`](environment.yaml) with conda). -Extras: `interpretability` (UMAP, plotly, matplotlib), `io` (h5ad, polygons), -`stats` (statsmodels), `plots`, `report` (PDF rendering), `spatial-stats` -(squidpy backend), `sce` (R inputs). +or `conda env create -f environment.yaml`. Extras: `io` (h5ad inputs, polygon +units), `sce` (SingleCellExperiment inputs; also needs R with Bioconductor), +`test`. -Installing exposes the CLI entry points below; the Hydra config tree ships -inside the package, so every command works from any directory. +## Quick start -## Quick start (public demo data, CPU, a few minutes) - -The vignettes run the whole chain on `squidpy.datasets.mibitof()` (three -MIBI-TOF images, ~3,300 cells): +The [vignettes](vignettes/) run the whole workflow on a public MIBI-TOF dataset +on a CPU in a few minutes: ```bash pip install squidpy jupyter jupyter lab vignettes/ ``` -1. [`01_anndata_to_spatial_graphs.ipynb`](vignettes/01_anndata_to_spatial_graphs.ipynb): AnnData -> manifest -> cellular graphs. -2. [`02_pretraining_and_niches.ipynb`](vignettes/02_pretraining_and_niches.ipynb): BGRL pretraining (`regime=pretrain_a3`), fine-tuning (`regime=mil`), instance export, interpretability report. -3. [`03_nested_cv_niches_and_metaclusters.ipynb`](vignettes/03_nested_cv_niches_and_metaclusters.ipynb): nested CV, cell extraction, niches, metaclusters, scores, spatial statistics. - -## Typical usage on your own cohort +## Workflow Describe the cohort once (copy [`configs/experiment/example_cohort.yaml`](src/grass_mil/configs/experiment/example_cohort.yaml) -to `my_cohort.yaml`, fill in paths and column names), then: +to `my_cohort.yaml` and fill in paths and column names; see +[data preparation](docs/data_preparation.md)), then: -```bash -grass-mil-train experiment=my_cohort +trainer.fast_dev_run=true # build the graphs -grass-mil-train experiment=my_cohort regime=pretrain_a3 \ - paths.output_dir=runs/pretrain_a3 trainer=gpu # pretrain, label-free -grass-mil-nested-cv experiment=my_cohort regime=linear_ftall_reg \ - model.init_from_ckpt=runs/pretrain_a3/checkpoints/last.ckpt \ - paths.output_dir=runs/nested trainer=gpu # nested CV, OOF scores -grass-mil-scores input=runs/nested/oof_predictions.csv output_dir=runs/scores # CIs, operating points -grass-mil-extract experiment=my_cohort regime=linear_ftall_reg \ - extract.nested_cv_summary=runs/nested/nested_cv_summary.json \ - extract.output_dir=runs/cells # every cell, every fold model -grass-mil-niches cells_dir=runs/cells/fold2 output_dir=runs/niches_fold2 # niches of one model -grass-mil-metaclusters cells_dir=runs/cells output_dir=runs/meta labels.label_file=... \ - labels.label_column=STAS labels.patient_column=PATIENT # replication across folds -grass-mil-spatial-stats cells_dir=runs/cells/fold0 output_dir=runs/spatial labels.label_file=... \ - labels.label_column=STAS labels.patient_column=PATIENT reference=[tumor] # model-free statistics -``` +| Step | Command | Output | +|---|---|---| +| 1. Pretrain the encoder, label-free (builds the cell graphs on first use) | `grass-mil-train experiment=my_cohort regime=pretrain` | `checkpoints/last.ckpt` | +| 2. Evaluate by nested cross-validation | `grass-mil-nested-cv experiment=my_cohort model.init_from_ckpt=` | per-fold metrics (mean ± s.d.), out-of-fold unit and patient predictions, one model per fold | +| 3. Score every cell with the fold models | `grass-mil-extract experiment=my_cohort extract.nested_cv_summary=/nested_cv_summary.json` | per-cell embeddings, attention and logits | +| 4. Discover and describe niches | `grass-mil-niches cells_dir=/cells/fold0` | niche table, spatial descriptors, figures, report | +| Optional: replication across folds | `grass-mil-metaclusters cells_dir=/cells` | metaclusters of niches found by several models | +| Optional: model-free statistics | `grass-mil-spatial-stats cells_dir=/cells/fold0` | classical spatial statistics compared between classes | +| Score new data | `grass-mil-predict experiment=new_cohort ckpt_path=` | unit and patient predictions | + +A single model on all labelled data (`grass-mil-train experiment=my_cohort +model.init_from_ckpt=`) uses a random patient-grouped split for +validation and testing. Add `trainer=gpu` to train on a GPU. -Step by step, with inputs, outputs and the variants (one split, tiles as -units, h5ad inputs, patient-level prediction): -[`docs/getting_started.md`](docs/getting_started.md). How configuration works -and how to change any setting: [`docs/configuration_guide.md`](docs/configuration_guide.md). - -## Workflows - -Inputs are a manifest CSV (`sample_id,input_path,input_type[,region_id,polygons_path]`), -per-sample cell tables (CSV/TSV, h5ad, SingleCellExperiment), a label CSV with a -unit-id column, a target column and a patient column, and a fold file (JSON list -of lists of unit ids). See [`docs/dataset_preparation.md`](docs/dataset_preparation.md). - -| Step | Command | Preset / keys | Outputs | Doc | -|---|---|---|---|---| -| Build graphs | `grass-mil-train data=... +trainer.fast_dev_run=true` (or any run) | `data.graph_builder.kwargs.edge_features=[log_distance,neighbor]`, `data.sample_unit=full` | `processed_index.json`, graph `.pt` files | [dataset_preparation](docs/dataset_preparation.md) | -| Pretrain the encoder | `grass-mil-train experiment=my_cohort regime=pretrain_a3` | `model=bgrl_a3`, `task=pretrain_bgrl_a3`, `data/sampler=ego_graph_pretrain` | `checkpoints/last.ckpt` | [workflows](docs/workflows_train_infer.md) | -| Fine-tune one split | `grass-mil-train regime=mil data=... model.init_from_ckpt=...` | `regime=mil|mil_frozen|linear_probehead|linear_ftall|linear_ftall_reg`, `data.split.explicit.*` | checkpoints, `metrics_per_epoch.jsonl`, `test_predictions.csv` | [nested_cv](docs/nested_cv.md) | -| Probe head | `grass-mil-probe regime=linear_probehead data=...` | `probe.*` | `probe_head.pt` | [nested_cv](docs/nested_cv.md) | -| Nested cross-validation | `grass-mil-nested-cv regime=... nested_cv.folds_file=folds.json` | `nested_cv.*` | `nested_cv_summary.json`, `oof_predictions.csv`, `oof_patient_predictions.csv` | [nested_cv](docs/nested_cv.md) | -| Predict / aggregate | `grass-mil-predict ckpt_path=... regime=...` | `aggregation.bag_scope=patient aggregation.mode=mean_probability|attention_weighted` | `predictions.csv`, `metrics.json`, instance tables | [inference](docs/inference_utility_suite.md) | -| Extract cells | `grass-mil-extract regime=... extract.nested_cv_summary=...` | `extract.*` | `cells//.npz` | [niche_discovery](docs/niche_discovery.md) | -| Niches of one model | `grass-mil-niches cells_dir=cells/fold2` | `reduction.*`, `clustering.*`, `sample_size`, `assign_k` | `clusters.csv`, `cells.csv`, `report.md`, figures | [niche_discovery](docs/niche_discovery.md) | -| Cross-fold replication | `grass-mil-metaclusters cells_dir=cells labels.label_file=...` | `stage`, `match.*`, `transfer.*`, `basis.*` | `metaclusters.csv`, patient shares and tests, transfer and basis reports | [niche_discovery](docs/niche_discovery.md) | -| Evaluation statistics | `grass-mil-scores input=oof_predictions.csv` | `scores`, `strata`, `operating_rules` | AUROC CIs, DeLong, operating points | [evaluation_stats](docs/evaluation_stats.md) | -| Classical spatial statistics | `grass-mil-spatial-stats cells_dir=... reference=[tumor]` | `analyses`, `group_map`, `nhood.*`, `cn.*`, `crossk.*` | neighbourhood enrichment, cellular neighbourhoods, contacts, cross-K tests | [spatial_statistics](docs/spatial_statistics.md) | -| Interpretability report | `grass-mil-report data.instance_table=...` | `clustering.*`, plugins | HTML / PDF report | [interpretability_suite](docs/interpretability_suite.md) | - -## Running the analyses of the preprint - -The preprint's numbers come from this chain with the default settings: - -1. graphs with `edge_features=[log_distance,neighbor]` and the cohort's native - coordinate units (`neighbor_cutoff_um` is the contact distance in those units); -2. `regime=pretrain_a3` on every region, label-free; -3. `grass-mil-nested-cv` with patient-grouped outer folds, inner `k=3`, the - regimes `linear_probehead` ("probehead"), `linear_ftall` ("ftALL"), - `linear_ftall_reg` ("ftALL, lr x0.3 + regularisation") and `mil`; - out-of-fold region scores, patient scores as the mean of a patient's regions - (merged bags via `grass-mil-predict`); -4. `grass-mil-scores` for bootstrap intervals, per-fold means, DeLong - comparisons and operating points; -5. `grass-mil-extract` of the fold refits, `grass-mil-niches` on each model's - held-out cells (PCA-10, L2, HDBSCAN mcs 300 / min_samples 1 / eom, 15-NN - assignment), `grass-mil-metaclusters` (Hellinger matching with cannot-link - within fold, transfer to all regions, shared k-means basis); -6. `grass-mil-spatial-stats` for the model-free comparison. - -[`docs/method_correspondence.md`](docs/method_correspondence.md) maps every -section of the preprint to the preset, module and test that implements it and -records the validation against the original runs. - -### Changing the defaults +## Changing the defaults ```bash -# shallower encoder, 3-hop ego-graphs, cell type only (no size feature) -grass-mil-nested-cv regime=mil model.encoder.num_layers=3 model.encoder.hidden_dim=128 \ - model.encoder.out_dim=128 data.sampler.runtime.depth=3 data.csv.molecular_columns=[] ... +# another fine-tuning regime: slower encoder with regularisation +grass-mil-nested-cv experiment=my_cohort regime=finetune_reg model.init_from_ckpt=... -# radius-capped ego-graphs and a different contact distance -grass-mil-train regime=mil data.sampler.runtime.ego_radius.enabled=true \ - data.sampler.runtime.ego_radius.radius_per_hop=30 data.graph_builder.kwargs.neighbor_cutoff_um=25 ... +# a shallower encoder on 3-hop ego-graphs, cell type only (no continuous feature) +grass-mil-train experiment=my_cohort regime=pretrain model.encoder.num_layers=3 \ + data.sampler.depth=3 model.encoder.hidden_dim=128 model.encoder.out_dim=128 data.input.features=null -# more roots per bag, longer schedule, five inner folds -grass-mil-nested-cv regime=linear_ftall data.sampler.runtime.subsample_count=512 \ +# radius-capped ego-graphs and another contact distance +grass-mil-train experiment=my_cohort data.sampler.ego_radius.enabled=true \ + data.sampler.ego_radius.radius_per_hop=30 data.graph.neighbor_cutoff=25 ... + +# more cells per bag, a longer schedule, five inner folds +grass-mil-nested-cv experiment=my_cohort data.sampler.roots_per_unit=512 \ trainer.max_epochs=40 nested_cv.inner_k=5 ... -# niches with k-means instead of HDBSCAN, 20 PCs, no sampling -grass-mil-niches cells_dir=... clustering.method=kmeans clustering.n_clusters=12 \ - reduction.n_components=20 sample_size=null -``` +# survival instead of classification +grass-mil-nested-cv experiment=my_cohort task=survival task.target_columns=[time,event] ... -Every key is documented in [`docs/hyperparameter_reference.md`](docs/hyperparameter_reference.md). +# niches in the encoder embedding with k-means +grass-mil-niches cells_dir=... clustering.space=embedding clustering.pca_components=10 \ + clustering.method=kmeans clustering.n_clusters=12 +``` ## Documentation -- [Getting started: typical usage on your own cohort](docs/getting_started.md) -- [Configuration guide: groups, presets, precedence, every knob by concern](docs/configuration_guide.md) -- [Dataset preparation](docs/dataset_preparation.md): manifest, inputs, labels, patient column, fold files -- [Training and inference workflows](docs/workflows_train_infer.md) -- [Nested cross-validation and regimes](docs/nested_cv.md) -- [Niche discovery and cross-fold replication](docs/niche_discovery.md) -- [Evaluation statistics](docs/evaluation_stats.md) -- [Classical spatial statistics](docs/spatial_statistics.md) -- [Inference utility suite](docs/inference_utility_suite.md) -- [Interpretability report suite](docs/interpretability_suite.md) and [extension guide](docs/interpretability_extension_guide.md) -- [Hyperparameter reference](docs/hyperparameter_reference.md) -- [Capabilities and modes](docs/capabilities_and_modes.md) -- [Model component contracts](docs/model_component_contracts.md) -- [Method correspondence with the preprint](docs/method_correspondence.md) -- [Repository map](docs/repository_map.md) +- [Data preparation](docs/data_preparation.md): cell tables, manifest, units, labels, folds, graphs +- [Training and evaluation](docs/training_and_evaluation.md): the model, pretraining, regimes, tasks, nested cross-validation, prediction +- [Niche analysis](docs/niche_analysis.md): per-cell outputs, niches, descriptors, survival, replication, model-free statistics +- [Configuration](docs/configuration.md): groups, precedence and every key ## Citation +If you use GRASS-MIL, please cite: + +> Rabuzin L, Yates J, Xu M, Boeva V. *GRASS-MIL: graph-based representation +> and discovery of phenotype-associated spatial structures with multiple +> instance learning.* bioRxiv (2026). +> doi: [10.64898/2026.10.08.757722](https://www.biorxiv.org/content/10.64898/2026.10.08.757722v1) + ```bibtex @article{rabuzin2026grassmil, title = {GRASS-MIL: graph-based representation and discovery of phenotype-associated spatial structures with multiple instance learning}, @@ -204,8 +131,6 @@ Every key is documented in [`docs/hyperparameter_reference.md`](docs/hyperparame } ``` -A machine-readable citation is in [`CITATION.cff`](CITATION.cff). - ## License and contact MIT license. Questions and issues: open a GitHub issue, or contact the Boeva diff --git a/docs/capabilities_and_modes.md b/docs/capabilities_and_modes.md deleted file mode 100644 index a890ac2..0000000 --- a/docs/capabilities_and_modes.md +++ /dev/null @@ -1,208 +0,0 @@ -# Capabilities And Modes - -This document catalogs operational modes and non-default capabilities in `grass-mil`. - -## 1) Training Regimes - -| Regime | Task | Model | Primary Use | -|---|---|---|---| -| Supervised mean | `finetune_mean` | `supervised_module` | Stable baseline graph-level classification/regression/survival | -| Supervised MIL | `finetune_mil` | `supervised_module` | Per-class attention aggregation; region cross-entropy is the sole objective | -| Cox survival MIL | `finetune_survival` | `supervised_module` | Scalar log-hazard per bag, scored by c-index | -| SSL pretrain | `pretrain_bgrl` | `bgrl_module` | Representation learning before supervised fine-tuning | -| GRASS-MIL pretrain (A3) | `regime=pretrain_a3` | `bgrl_a3` | The preprint's pretraining recipe | -| GRASS-MIL fine-tuning | `regime=mil` (`mil_frozen`, `linear_probehead`, `linear_ftall`, `linear_ftall_reg`) | `grass_mil` | One region = one bag, scaled learning rates, probe pre-fit, region-level selection | - -## 1b) Analysis Entry Points - -| CLI | Purpose | Doc | -|---|---|---| -| `grass-mil-probe` | logistic probe on the frozen encoder, folded into a linear head | `nested_cv.md` | -| `grass-mil-nested-cv` | nested patient-grouped cross-validation with out-of-fold predictions | `nested_cv.md` | -| `grass-mil-extract` | per-cell cache (embedding, instance logits, attention) of trained models | `niche_discovery.md` | -| `grass-mil-niches` | niche discovery and attention-weighted attribution on a model's cells | `niche_discovery.md` | -| `grass-mil-metaclusters` | cross-fold niche replication, transfer, shared basis | `niche_discovery.md` | -| `grass-mil-scores` | AUROC intervals, DeLong, operating points | `evaluation_stats.md` | -| `grass-mil-spatial-stats` | neighbourhood enrichment, cellular neighbourhoods, contacts, cross-K | `spatial_statistics.md` | -| `grass-mil-predict` | predictions, region/patient aggregation (mean or merged bag), embeddings, instance tables | `inference_utility_suite.md` | -| `grass-mil-report` | HTML/PDF interpretability report on instance tables | `interpretability_suite.md` | - -## 2) Aggregation And MIL Feature Modes - -| Mode | Config | Behavior | -|---|---|---| -| Mean bag aggregation | `task.aggregation=mean` | Averages instance logits per bag | -| MIL attention aggregation | `task.aggregation=mil_attention` | Learns attention weights over instance embeddings | -| Region accumulation | `task.region_accumulation.enabled=true` | Manual optimizer steps over region hyperbatches | -| One region per bag | `data.sampler.runtime.unit_as_bag=true` (`data/sampler=ego_graph_mil`) | All sampled ego-graphs of a region in one batch: one attention softmax, one loss | -| Patient aggregation | `aggregation.bag_scope=patient` with `mode=mean_probability` or `attention_weighted` | Mean of region probabilities, or one merged bag per patient | - -## 3) Backbone/Encoder Modes - -Supported `model.encoder.conv_type` values: - -- `gin` -- `gcn` -- `gat` -- `graphsage` -- `gine` - -Notable behavior: - -1. `gcn` and `gine` support edge-aware message passing. -2. `gine` requires edge feature dimensionality (`model.encoder.edge_attr_dim`) and runtime `edge_attr`. -3. `gat` requires `hidden_dim` divisible by `gat_heads`. -4. JK modes: `last`, `concat`, `max`, `sum`. -5. Pooling modes: `sum`, `mean`, `max`. Learned pooling was removed: instance - selection is the MIL head's job, pooling across ego-graphs. - -## 4) Data Sampling Modes - -Sampler strategies (`data.sampler.name`): - -- `identity`: standard graph batching. -- `shadow_native`: native PyG ShaDow sampler. -- `shadow_custom`: custom ShaDow implementation supporting per-subgraph transforms. - -Runtime controls (`data.sampler.runtime.*`): - -- depth and neighbor count (`num_neighbors=-1` = unlimited) -- subgraph batch size -- replacement and shuffling controls -- proportional root sampling by categorical property -- class-weight mode (`inverse`, `sqrt_inverse`, `proportional`) -- roots per unit per epoch (`subsample_count` / `subsample_fraction`), fixed per-unit seeds -- one batch per unit (`unit_as_bag`), roots from the union of units (`merge_units`) -- separate validation and test samplers (`data.val_sampler`, `data.test_sampler`) - -## 5) Loss/Task Objective Modes - -Supported `task.target_type`: - -- `binary` -- `categorical` (per-class attention; the GRASS-MIL default for binary tasks) -- `regression` -- `survival` - -Supported `task.loss` values (via model loss presets): - -- `categorical_bce` -- `categorical_ce` -- `regression_mse` -- `regression_huber` -- `survival_coxsgd` - -## 6) Runtime Device/Distribution Modes - -Trainer presets: - -- `trainer=cpu` -- `trainer=gpu` -- `trainer=mps` -- `trainer=ddp` -- `trainer=ddp_sim` - -Additional runtime toggles can be overridden directly via Hydra: - -- precision/mixed precision -- deterministic mode -- max/min epochs -- validation frequency - -## 7) Logging/Experiment Tracking Modes - -Logger presets: - -- `logger=csv` -- `logger=tensorboard` -- `logger=wandb` -- `logger=mlflow` -- `logger=comet` -- `logger=neptune` -- `logger=aim` -- `logger=many_loggers` (composition) - -Callbacks can be selected via `callbacks=...` group or disabled with `callbacks=none`. - -## 8) Debug/Development Modes - -Debug presets: - -- `debug=default` -- `debug=fdr` -- `debug=limit` -- `debug=overfit` -- `debug=profiler` - -Common usage: - -```bash -grass-mil-train debug=fdr -grass-mil-train debug=limit -``` - -## 9) Checkpoint Initialization And Resume Modes - -| Mode | Config | Behavior | -|---|---|---| -| Resume training | `ckpt_path=/abs/last.ckpt` | Resume optimizer/trainer state through Lightning fit | -| Initialize weights from checkpoint | `model.init_from_ckpt=/abs/ckpt` | Load state dict before training | -| BGRL-to-supervised key remap | `model.encoder_init_map=auto_bgrl_or_identity` (`spacegm_bgrl` for legacy checkpoints) | Remaps known encoder key prefixes | -| Probe head init | `task.head_init_path=` | Linear head from `grass-mil-probe`; `freeze_head` keeps it fixed | -| Strict load | `model.init_strict=true` | Fail on missing/unexpected keys | -| Freeze encoder | `model.freeze_encoder=true` | Train head/attention while encoder weights stay fixed | - -## 10) Inference Modes - -1. CLI evaluation mode: `src/grass_mil/eval.py` with required `ckpt_path`. -2. Advanced prediction mode: `Trainer.predict(...)` using `SupervisedModule.predict_step`. -3. Inference utility CLI mode: `src/grass_mil/inference/predict.py` for prediction export, metrics, and embeddings. - -Predict outputs include: - -- `bag_ids` -- `bag_logits` -- `row_region_ids` -- `row_sample_ids` -- optional `bag_targets` -- optional `bag_attention` -- optional `instance_*` fields -- optional embedding fields (`embedding_bag_ids`, `graph_embeddings`, `embedding_bag_counts`) -- optional node embedding fields (`node_embeddings`, `node_bag_ids`) - -Metrics available through the inference utility suite: - -- categorical/binary: accuracy, precision, recall, F1, ROC-AUC -- regression: R2, MAE, RMSE -- survival: concordance index (c-index) - -Validation (during `src/grass_mil/train.py`) includes: - -- step-level `val/loss` logging -- epoch-end region-level aggregation of buffered validation outputs -- task-specific epoch metrics: - - `val/acc`, `val/loss_region`, `val/auroc`, `val/bal_acc` for binary/categorical (region level) - - `val/mae`, `val/rmse`, `val/r2` for regression - - `val/c_index` for survival -- optional dedicated validation sampler via `data.val_sampler.*` - -## 11) Notebook Validation Modes - -Notebooks support interactive smoke validation of: - -- ingesting a public AnnData cohort and building cellular graphs - (`vignettes/01_anndata_to_spatial_graphs.ipynb`) -- self-supervised pretraining and spatial-niche discovery - (`vignettes/02_pretraining_and_niches.ipynb`) -- nested cross-validation, niche attribution, cross-fold metaclusters and scores - (`vignettes/03_nested_cv_niches_and_metaclusters.ipynb`) - -## 12) CI And Quality Modes - -Quality and validation tooling: - -- `pre-commit run -a` -- `pytest` -- GitHub Actions workflows for tests and code quality - -These operational modes are versioned in `.github/workflows/*.yaml` and `.pre-commit-config.yaml`. diff --git a/docs/configuration.md b/docs/configuration.md new file mode 100644 index 0000000..0d4cb2e --- /dev/null +++ b/docs/configuration.md @@ -0,0 +1,155 @@ +# Configuration + +Every command is a [Hydra](https://hydra.cc) application whose configuration +ships with the package (`src/grass_mil/configs`). Any key can be set on the +command line (`key=value`, lists as `[a,b]`), and `--cfg job` prints the +composed configuration without running anything: + +```bash +grass-mil-train experiment=my_cohort regime=finetune_reg --cfg job +``` + +## Groups + +`grass-mil-train` (and `nested-cv`, `extract`, `predict`, which start from it) +composes these groups: + +| Group | Options | Selects | +|---|---|---| +| `data` | `spatial_omics` | inputs, units, graphs, labels, splits, sampling | +| `data/sampler` | `ego_graph`, `ego_graph_pretrain` | ego-graph sampling (set by the regime) | +| `model` | `grass_mil`, `bgrl` | the module (`bgrl` is set by `regime=pretrain`) | +| `task` | `classification`, `regression`, `survival` | targets, objective, selection metric | +| `regime` | `finetune`, `finetune_reg`, `attention_only`, `mlp_head`, `mlp_head_frozen`, `pretrain` | the training procedure | +| `experiment` | your cohort file | paths and columns of one cohort | +| `trainer` | `default` (CPU), `cpu`, `gpu` (fp16 mixed precision) | the Lightning trainer | +| `logger` | `csv`, `tensorboard` | optional metric loggers | +| `optim`, `scheduler`, `callbacks` | one option each | AdamW, cosine schedule, selection callbacks | + +Later entries win: a cohort file sets data and labels, the regime sets the +procedure, and command-line keys override both. + +## Top level + +| Key | Default | Meaning | +|---|---|---| +| `seed` | `0` | seeds every random number generator | +| `train`, `test` | `true` | fit; score the test split with the selected checkpoint | +| `ckpt_path` | `null` | resume an interrupted run (prediction: the model to score) | +| `print_config` | `false` | log the composed configuration | +| `paths.data_dir`, `paths.log_dir` | `./data`, `./logs` | relative to the working directory | +| `paths.output_dir` | per run | `logs//runs/` | + +## `data` + +| Key | Default | Meaning | +|---|---|---| +| `manifest_path` | `${paths.data_dir}/manifest.csv` | see [data preparation](data_preparation.md) | +| `processed_dir` | `${paths.data_dir}/processed` | built graphs | +| `force_precompute` | `false` | rebuild the graphs | +| `unit` | `region` | `region`, `tile` or `polygon` | +| `coord_scale_um` | `1.0` | coordinate multiplier | +| `min_cells` | `10` | smaller units are dropped | +| `tiling.size`, `tiling.stride` | `200`, `200` | tiles, coordinate units | +| `input.coord_columns` | `[x, y]` | | +| `input.obsm_key` | `spatial` | h5ad coordinates (`null`: `obs` columns) | +| `input.reduced_dim` | `null` | SCE coordinates (`null`: `colData` columns) | +| `input.cell_id_column` | `null` | | +| `input.cell_type_column` | `cell_type` | `null`: no cell types | +| `input.features` | `all` | `all`, a list of names, or `null` | +| `input.h5ad_layer`, `input.sce_assay`, `input.sep` | `null`, `null`, `","` | | +| `graph.edge_features` | `[log_distance, neighbor]` | also `distance` | +| `graph.neighbor_cutoff` | `20.0` | contact distance | +| `graph.log_distance_bounds` | `[2.0, 5.0]` | ln-length range mapped to [0, 1] | +| `feature_reducer.name` | `identity` | `pca` | +| `feature_reducer.n_components`, `.fit_on`, `.state_from` | `32`, `all`, `null` | see data preparation | +| `labels.file`, `.id_column`, `.patient_column` | `null`, `region_id`, `null` | | +| `labels.columns` | `task.target_columns` | | +| `labelled_only` | `true` | skip unlabelled units (pretraining: `false`) | +| `split.ratios`, `split.seed` | `[0.8, 0.1, 0.1]`, `0` | random split, grouped by patient | +| `split.train`, `.val`, `.test` | `null` | explicit ids (lists or files) | +| `val_roots_per_unit` | `1024` | cells scored per validation unit | +| `num_workers`, `pin_memory` | `0`, `false` | | + +### `data.sampler` + +| Key | `ego_graph` | Meaning | +|---|---|---| +| `depth` | `4` | hops (keep equal to `model.encoder.num_layers`) | +| `num_neighbors` | `-1` | neighbours per node and hop (`-1`: all) | +| `roots_per_unit` | `256` | cells per bag and epoch (`null`: all) | +| `batch_size` | `64` | ego-graphs per batch when scoring and pretraining | +| `merge_units` | `false` | pool the cells of all units (pretraining) | +| `root_weighting` | `null` | `inverse` / `sqrt_inverse` cell-type frequency | +| `interior_roots.enabled` | `false` | only cells at least `n_hops` hops from the convex hull | +| `ego_radius.enabled` | `false` | drop nodes farther than `radius_per_hop * depth + radius_offset` | + +## `model` + +| Key | Default | Meaning | +|---|---|---| +| `encoder.hidden_dim`, `.out_dim` | `256`, `256` | | +| `encoder.num_layers` | `4` | GINE layers | +| `encoder.dropout` | `0.0` (`0.1` in `bgrl`) | | +| `encoder.norm`, `.jk`, `.act`, `.pooling` | `layernorm`, `last`, `relu`, `max` | | +| `encoder.mlp_ratio`, `.train_eps` | `2`, `true` | GINE MLP width, learnable epsilon | +| `encoder.edge_feature` | `0` | edge-feature column the convolutions see (`null`: all) | +| `encoder.input_init` | `xavier` | | +| `encoder.unassigned_cell_type` | `null` | masked-node token (`null`: an extra row) | +| `head.num_layers` | `1` | `1` linear, `2` MLP | +| `head.hidden_dim`, `.dropout`, `.act` | `null`, `0.0`, `leaky_relu` | | +| `attention.projection_dim`, `.hidden_dim` | `32`, `8` | | +| `attention.dropout`, `.init` | `0.25`, `kaiming` | | +| `init_from_ckpt` | `null` | encoder weights (a pretraining checkpoint) | + +## `task` + +| Key | Default | Meaning | +|---|---|---| +| `target_type` | `classification` | `regression`, `survival` | +| `target_columns` | `null` | label columns (`null`: the first one) | +| `num_classes` | `2` | classification | +| `class_weights` | `auto` | inverse training frequency, a list, or `null` | +| `loss`, `huber_delta` | `mse`, `1.0` | regression | +| `cox_regularization` | `0.0` | survival | +| `selection.metric`, `.mode`, `.patience` | `val/auroc`, `max`, `8` | checkpoint selection and early stopping | + +## `finetune` (fine-tuning regimes) + +| Key | `finetune` | Meaning | +|---|---|---| +| `lr_scales.encoder/head/attention` | `1.0`, `0.01`, `2.0` | multipliers of `optim.lr`; `0` freezes | +| `head_weight_decay` | `1e-3` | (`null`: `optim.weight_decay`) | +| `head_init` | `probe` | or `random` | +| `freeze_head` | `false` | | +| `hyperbatch_size` | `12` | units per optimizer step | +| `grad_clip_norm` | `1.0` | | +| `warmup_steps` | `0` | linear learning-rate warmup | +| `augmentation.drop_edge_p`, `.mask_node_p` | `0.0`, `0.0` | | +| `probe.roots_per_unit`, `.C`, `.class_weight` | `96`, `1.0`, `balanced` | | +| `optim.lr`, `optim.weight_decay` | `1e-4`, `1e-3` | | +| `trainer.max_epochs`, `.min_epochs` | `25`, `10` | | + +## `pretrain` (`regime=pretrain`) + +| Key | Default | Meaning | +|---|---|---| +| `total_steps`, `warmup_steps` | `200000`, `1000` | optimizer steps | +| `momentum_start`, `momentum_end` | `0.99`, `1.0` | target-encoder EMA, cosine schedule | +| `predictor_hidden_dim` | `512` | | +| `augmentation.mu`, `.p_lambda` | `0.25`, `0.45` | node/edge drop scale and ceiling | +| `augmentation.noise_std`, `.noise_features` | `0.2`, `[]` | Gaussian noise on named features | +| `optim.lr`, `optim.weight_decay` | `5e-5`, `1e-5` | | + +## Other commands + +* `grass-mil-nested-cv`: `nested_cv.folds_file`, `inner_k` (3), `inner_seed`, + `epoch_reduce` (`median`), `min_refit_epochs`, `outer_folds`, `resume`. +* `grass-mil-extract`: `extract.nested_cv_summary`, `models`, `output_dir`, + `units`, `batch_size` (512), `device` (`auto`), `resume`. +* `grass-mil-predict`: `ckpt_path`, `predict.units`, `patient_aggregation` + (`mean`, `merged_bag`), `batch_size`, `device`. +* `grass-mil-niches`, `grass-mil-metaclusters`, `grass-mil-spatial-stats`: see + [niche analysis](niche_analysis.md); their configs + (`configs/niches.yaml`, `metaclusters.yaml`, `spatial_stats.yaml`) document + every key. diff --git a/docs/configuration_guide.md b/docs/configuration_guide.md deleted file mode 100644 index 29e8825..0000000 --- a/docs/configuration_guide.md +++ /dev/null @@ -1,223 +0,0 @@ -# Configuration Guide - -Everything in `grass-mil` is configured through Hydra: a command composes a -tree of YAML files from `src/grass_mil/configs/`, and every leaf can be -overridden on the command line as `key.path=value`. This page explains how the -tree is composed, what the groups are, and how to change the things people -usually change. The exhaustive key-by-key tables are in the -[hyperparameter reference](hyperparameter_reference.md). - -## 1. How a command is composed - -Each CLI has a root config (`train.yaml`, `nested_cv.yaml`, `niches.yaml`, ...). -A root config lists *groups* in its `defaults`, in order; later entries override -earlier ones, and command-line overrides win over everything: - -``` -train.yaml defaults: data -> model -> optim -> scheduler -> task -> callbacks -> trainer -> paths - -> experiment (null) -> regime (null) -> debug (null) -``` - -A *preset* is a file in a group. `regime=mil` is the file -`configs/regime/mil.yaml`; it is "`@package _global_`", so it may set keys -anywhere in the tree and may swap other groups (`override /model: grass_mil`, -`override /task: finetune_mil`, `override /scheduler: cosine_mil_step`, -`override /callbacks: mil_selection`, `override /data/sampler: ego_graph_mil`). -That is how one word selects a whole training recipe. An `experiment=` file is -the same mechanism, meant for *your* settings (the cohort template in -`configs/experiment/example_cohort.yaml` is one). - -Precedence, lowest to highest: group file defaults -> `experiment=` -> `regime=` --> command-line `key=value`. Two consequences: - -- a regime's values (e.g. `data.sampler.runtime.subsample_count=256`) beat the - same key in your experiment file; override it on the command line if needed; -- a key that does not exist yet needs `+key=value` (e.g. - `+trainer.fast_dev_run=true`); an existing key is set with `key=value`. - -Useful Hydra flags: `--cfg job` prints the fully composed config and exits; -`--info defaults` shows which files were merged; `-m key=a,b` runs a sweep. - -The composed config is also written to the run directory -(`/.hydra/config.yaml`) so every run is reproducible from its -folder. - -## 2. The groups and their presets - -| Group | Presets (default first) | What it decides | -|---|---|---| -| `data` | `spatial_omics` | inputs, graph building, labels, splits, tiling (section 3) | -| `data/sampler` | `shadow_custom`, `ego_graph_mil` (regimes), `ego_graph_pretrain` (pretraining), `identity`, `shadow_native` | how ego-graphs are drawn (section 4) | -| `model` | `supervised_module`, `grass_mil` (regimes), `bgrl_module`, `bgrl_a3` (pretraining) | the module: encoder + head + attention, or BGRL | -| `model/encoder` | `gin`, `gine_length` (GRASS-MIL), `gine`, `gcn`, `gat`, `graphsage` | encoder architecture (section 5) | -| `model/attention` | `gated_projected`, `gated` | attention net (section 6) | -| `task` | `finetune_mean`, `finetune_mil` (regimes), `finetune_survival`, `pretrain_bgrl`, `pretrain_bgrl_a3` | objective, bag definition, optimisation knobs, augmentation (section 7) | -| `regime` | `mil`, `mil_frozen`, `linear_probehead`, `linear_ftall`, `linear_ftall_reg`, `pretrain_a3` | the preprint recipes: fine-tuning (section 7) and pretraining (section 9) | -| `optim` | `adamw` | optimiser | -| `scheduler` | `cosine_epoch`, `cosine_mil_step` (regimes), `cosine_step` (pretraining) | learning-rate schedule | -| `callbacks` | `default`, `mil_selection` (regimes), `refit`, `pretrain_checkpoint`, `none` | checkpointing, early stopping (section 8) | -| `trainer` | `default` (cpu), `gpu`, `gpu_amp`, `mps`, `ddp`, `ddp_sim` | device, epochs, precision (section 11) | -| `experiment` | `example_cohort` (template), `grass_mil_pretrain_a3` (alias of `regime=pretrain_a3`) | your cohort file goes here; one experiment file per command | -| `logger` | `csv`, `tensorboard`, `wandb`, `mlflow`, ... | experiment tracking (off by default) | -| `debug` | `fdr`, `limit`, `overfit`, `profiler` | quick runs | - -The analysis CLIs have flat root configs rather than groups: `probe.yaml`, -`nested_cv.yaml`, `extract.yaml`, `scores.yaml`, `niches.yaml`, -`metaclusters.yaml`, `spatial_stats.yaml`, `inference/predict.yaml`. The ones -that build a model (`probe`, `nested_cv`, `extract`, `predict`, `eval`) -include the `train.yaml` tree, so `experiment=`, `regime=` and every `data.*` -/ `model.*` key work there too. - -## 3. Inputs, features and graph (`data.*`) - -| What | Key | Notes | -|---|---|---| -| where the regions are | `data.raw_manifest_path`, `data.processed_dir` | manifest columns: `sample_id,input_path,input_type[,region_id,polygons_path]`; relative input paths resolve against the manifest | -| what a unit is | `data.sample_unit=full|tile` (+ `data.tiling.{tile_size_um,stride_um,min_cells}`) | `full` = one graph per region | -| coordinate unit | `data.coord_scale_um` | multiplies the input coordinates; all distances below are in the resulting unit | -| cell-type column | `data.csv.categorical_label_columns=[cell_type]`, `data.categorical_features.include_labels=[cell_type]` | the vocabulary is learnt from the data and stored in `processed/metadata.json` | -| continuous node features | `data.csv.molecular_columns=[SIZE]` (`[]` for none; `null` = every non-coordinate numeric column) | h5ad: `data.h5ad.molecular_features`, `data.h5ad.molecular_layer` | -| feature reduction | `data.feature_reducer.name=identity|pca|...`, `fit_mode` | identity by default | -| edges | `data.graph_builder.name=delaunay|knn`, `kwargs.edge_features=[log_distance,neighbor]`, `kwargs.neighbor_cutoff_um` | `log_distance` = `clip((ln d - 2)/3, 0, 1)` (`log_lower`, `log_upper` adjustable); `distance` = raw length | -| labels | `data.graph_labels.{label_file,id_column,tasks,scope,patient_column}` | `scope` = `region|patch|sample`; `patient_column` groups, never a target | -| splits | `data.split.explicit.{train,val,test,id_kind}` (lists or files) or the persisted `data.split.{split_by,train_val_test_split}` | nested CV sets explicit splits per fold | -| which cells may be roots | `data.sampler.runtime.interior_seeds.*` | exclude the tissue rim | - -Which node-feature layout reaches the encoder: `x` = continuous columns -(after the reducer), `categorical_codes` = the cell-type code; the encoder sums -a learnt cell-type embedding and a linear projection of `x`. A cohort with cell -types only has `x` of width 0 and works unchanged. - -## 4. Ego-graphs (`data.sampler.runtime.*`) - -| What | Key (default in `ego_graph_mil`) | -|---|---| -| hops | `depth` (4) | -| neighbours kept per hop | `num_neighbors` (-1 = all) | -| radius cap | `ego_radius.enabled` (false), `radius_per_hop`, `radius_offset` | -| roots per region per epoch (bag size) | `subsample_count` (256; `subsample_fraction` as an alternative) | -| uniform or cell-type-weighted roots | `proportional_root_sampling` (false), `property_name`, `weight_mode` | -| one bag per region | `unit_as_bag` (true) | -| validation subset | `data.val_sampler.runtime.subsample_count` (1024) with `fixed_per_unit_seed` (0) | -| test / prediction | `data.test_sampler` (null = every cell once) | -| pretraining batches | `ego_graph_pretrain`: `merge_units=true`, `subgraph_batch_size=64` | - -Example, 3-hop ego-graphs capped at 150 units of distance, 128 roots per bag: -`data.sampler.runtime.depth=3 data.sampler.runtime.ego_radius.enabled=true data.sampler.runtime.ego_radius.radius_per_hop=50 data.sampler.runtime.subsample_count=128`. - -## 5. Encoder (`model.encoder.*`) - -Selected by `model/encoder=gine_length` (the GRASS-MIL default inside -`model=grass_mil` and `model=bgrl_a3`). - -| What | Key (default) | -|---|---| -| depth and width | `num_layers` (4), `hidden_dim` (256), `out_dim` (256) | -| convolution | `conv_type` (`gine`; `gin`, `gcn`, `gat`, `graphsage`), `train_eps` (true) | -| edge conditioning | `use_edge_attr` (true), `edge_feature_index` (0 = the `log_distance` column), `edge_attr_dim` (1) | -| normalisation, activation, dropout | `norm` (`layernorm`), `act` (`relu`), `dropout` (0 when fine-tuning, 0.1 in `bgrl_a3`) | -| readout | `pooling` (`max`), `jk` (`last`) | -| cell-type embedding | `categorical_embedding.{label,reserve_unassigned,unassigned_label}`; `input_init` (`xavier`) | - -`model.encoder.num_layers=3 model.encoder.hidden_dim=128 model.encoder.out_dim=128` -gives a smaller encoder; the head and attention widths follow `out_dim` -automatically. A pretrained checkpoint must have been trained with the same -encoder settings. - -## 6. Attention and instance head - -| What | Key (default) | -|---|---| -| attention net | `model/attention=gated_projected` (`gated` = no projection); `model.attention.projection_dim` (32), `hidden_dim` (8), `dropout` (true), `init` (`kaiming`) | -| channels | `model.attention.n_classes` (2, one per class; set with `model.graph_head.output_dim` for more classes) | -| instance head | `model.graph_head.num_layers` (2 = MLP; 1 = linear, required by the `linear_*` regimes), `hidden_dim`, `dropout` | -| pooling | always logit-space: `L_c = sum_i A[i,c] * l[i,c]` (`task.aggregation=mil_attention`); `task.aggregation=mean` averages instance logits instead | - -## 7. Training regimes and optimisation (`regime=`, `task.*`, `optim.*`) - -Pick a regime, then adjust: - -| What | Key (where the regimes set it) | -|---|---| -| learning rate, weight decay | `optim.lr` (1e-4), `optim.weight_decay` (1e-3; 1e-2 in `linear_ftall_reg`) | -| per-component rates | `task.optimization.encoder_lr_scale` (1 / 0.3 / 0 = frozen), `head_lr_scale`, `attention_lr_scale` (2), `head_weight_decay` | -| head pre-fit | `task.head_init_path` (`auto` in `linear_*`: fitted by nested CV; or a `grass-mil-probe` file), `task.freeze_head` | -| class weights | `task.class_weights` (`auto_train`, a list, or null) | -| augmentation while fine-tuning | `task.finetune_augmentation.{drop_edge_p,drop_feat_p}` (0.05 in `linear_ftall_reg`) | -| accumulation and clipping | `task.region_accumulation.hyperbatch_size` (12), `task.grad_clip_norm` (1.0) | -| schedule | `scheduler=cosine_mil_step` with `scheduler.T_max=auto` (whole run) and `task.scheduler_interval=step`; `scheduler.eta_min` | -| epochs | `trainer.max_epochs` (25), `trainer.min_epochs` (10, before early stopping may act) | -| what a bag is | `task.bag_key` (`region_id`; `patch_id` for tiles) | -| probe settings | `probe.{roots_per_unit,seed,C,class_weight,max_iter}` and `nested_cv.probe.*` | - -To define your own recipe, copy `configs/regime/mil.yaml` to -`configs/regime/my_regime.yaml`, edit, and run `regime=my_regime`. - -## 8. Selection and checkpoints (`callbacks=`) - -`mil_selection` (regimes): checkpoint on `val/auroc`, early stopping on -`val/loss_region` with `callbacks.early_stopping.patience` (8). `refit`: no -selection, last weights. `pretrain_checkpoint`: every -`callbacks.model_checkpoint.every_n_train_steps` (5000). Any Lightning -callback argument is reachable, e.g. -`callbacks.model_checkpoint.save_top_k=3`. - -Logged per epoch and written to `metrics_per_epoch.jsonl`: `val/loss_region`, -`val/auroc`, `val/bal_acc`, `val/acc`, `train/loss`. - -## 9. Pretraining (`regime=pretrain_a3`) - -| What | Key (default) | -|---|---| -| length | `task.total_steps` (200000), `task.warmup_steps` and `task.lr_warmup.warmup_steps` (1000) | -| optimiser | `optim.lr` (5e-5), `optim.weight_decay` (1e-5) | -| target network | `task.momentum` (0.99) -> `task.momentum_final` (1.0) | -| views | `task.augmentation.{mode,mu,p_lambda,feature_noise_std,feature_noise_columns}`; `augmentation=uniform` for plain edge/feature dropout | -| encoder dropout | `model.encoder.dropout` (0.1) | -| batches | `data.sampler.runtime.subgraph_batch_size` (64), `depth` (4) | -| data | the regime sets `data.split.train_val_test_split=[1,0,0]`: every region, no labels needed | - -Any `model.encoder.*` change here must be repeated when fine-tuning. - -## 10. Nested cross-validation (`nested_cv.*`) - -| What | Key (default) | -|---|---| -| folds and labels | `nested_cv.folds_file`, `nested_cv.{label_file,id_column,label_column,patient_column}` (taken from `data.graph_labels` when unset), `nested_cv.id_kind=region|sample|patch` | -| inner selection | `nested_cv.inner_k` (3), `inner_seed` (0), `selection_metric` (`val/auroc`), `selection_mode`, `epoch_reduce` (`median`), `min_refit_epochs` | -| which outer folds | `nested_cv.outer_folds=[2]` (null = all), `nested_cv.resume` (true) | -| probes | `nested_cv.probe.{enabled,cache_dir,roots_per_unit,...}` | -| patient score | `nested_cv.patient_aggregation` (`mean`) | - -## 11. Devices, precision, runtime - -`trainer=gpu` (one GPU), `trainer.devices=2 trainer.strategy=ddp`, -`trainer.precision=16-mixed` (the regimes' default; `32` on CPU), -`data.num_workers` (0 is usually fastest for ego-graph sampling, which creates a -loader per region), `seed=0` for reproducible runs, `paths.output_dir=...` to -choose where a run writes (default `logs//runs/`). - -## 12. Extraction, niches, metaclusters, statistics - -These read a cell cache, not the training config, so their keys are flat: - -| CLI | Keys people change | -|---|---| -| `grass-mil-extract` | `extract.nested_cv_summary` or `+extract.models.=`, `extract.unit_id`, `extract.regions`, `extract.batch_size`, `extract.output_dir` (model composed from the same `experiment=`/`regime=` as training) | -| `grass-mil-niches` | `regions=held_out|all|`, `space=gnn|attention`, `reduction.{method,n_components}`, `clustering.{method,n_clusters,params.min_cluster_size,params.min_samples,params.cluster_selection_method}`, `sample_size`, `l2_normalise`, `assign_k`, `present_threshold`, `n_boot`, `umap.*`, `figures` | -| `grass-mil-metaclusters` | the same clustering keys, `match.{distance,threshold,scan,min_folds}`, `transfer.{min_cells,present_threshold}`, `basis.{k,space_model,sample_size}`, `stage` | -| `grass-mil-scores` | `scores`, `fold_column`, `patient_column`, `patient_level`, `standardise`, `strata`, `operating_rules`, `n_boot` | -| `grass-mil-spatial-stats` | `group_map`, `reference`, `partners`, `infiltrating`, `analyses`, `nhood.{n_perms,backend}`, `cn.{window_k,n_cn}`, `crossk.radii`, `min_per_group` | - -Lists on the command line are written `key=[a,b]`; dictionaries -`key='{a: b, c: d}'`; `null` clears a value. - -## 13. Reading the defaults - -`src/grass_mil/configs/` is the source of truth; every file is commented. -Start from `regime/mil.yaml` (what a fine-tuning run composes), -`task/finetune_mil.yaml` (every MIL knob), `data/sampler/ego_graph_mil.yaml` -(ego-graphs), `model/encoder/gine_length.yaml` (encoder) and -`regime/pretrain_a3.yaml` (pretraining). The -[hyperparameter reference](hyperparameter_reference.md) lists every key with -its type, valid values and an example override. diff --git a/docs/data_preparation.md b/docs/data_preparation.md new file mode 100644 index 0000000..a2d1d88 --- /dev/null +++ b/docs/data_preparation.md @@ -0,0 +1,116 @@ +# Data preparation + +GRASS-MIL needs three files describing a cohort, plus the per-region cell +tables: + +| File | Content | +|---|---| +| manifest (CSV) | one row per region: where its cell table is | +| labels (CSV) | one row per labelled region (or tile, or sample): targets and patient | +| folds (JSON) | outer cross-validation folds, for `grass-mil-nested-cv` | + +Describe them once in a cohort file (copy +[`configs/experiment/example_cohort.yaml`](../src/grass_mil/configs/experiment/example_cohort.yaml)) +and pass `experiment=my_cohort` to every command. + +## Cell tables + +One file per region with, for every cell, its coordinates and (normally) its +cell type, plus any continuous features: + +| Format | Coordinates | Cell type and features | +|---|---|---| +| CSV / TSV | `data.input.coord_columns` (default `x`, `y`) | columns | +| h5ad | `obsm[data.input.obsm_key]` (default `spatial`), or `obs` columns with `obsm_key: null` | `obs` column; features from `X` or `data.input.h5ad_layer` | +| SingleCellExperiment (`.rds`, needs `pip install 'grass-mil[sce]'` and R) | `colData` columns, or `reducedDim(data.input.reduced_dim)` | `colData` column; features from the assay | + +`data.input.cell_type_column` names the cell-type column (`null` if there is +none; the model then uses only the features). `data.input.features` selects the +continuous node features: `all` (every remaining CSV column, every h5ad/SCE +feature), a list of names (for example `[area]`), or `null` for cell type only. + +Coordinates are multiplied by `data.coord_scale_um`; every distance setting +(contact distance, tile size, radii) is in the resulting unit. + +## Manifest + +```csv +sample_id,input_path,region_id +S01,cells/S01_roi1.csv,S01_roi1 +S01,cells/S01_roi2.csv,S01_roi2 +S02,cells/S02.h5ad, +``` + +* `sample_id` and `input_path` are required; relative paths are resolved + against the manifest's directory. +* `region_id` defaults to `sample_id` and must be unique. +* `input_type` (`csv`, `tsv`, `h5ad`, `sce`) is optional when the file suffix + says it. +* `polygons_path` (optional) points to GeoJSON or WKT polygons of the region: + cells outside them are dropped, or with `data.unit=polygon` every polygon + becomes its own unit. + +## Units + +A unit is what one bag, one label and one fold member is: + +| `data.unit` | Unit | Id | +|---|---|---| +| `region` (default) | the whole region | `region_id` | +| `tile` | squares of `data.tiling.size`, every `data.tiling.stride` | `_tile` | +| `polygon` | every polygon of the region's `polygons_path` | `_polygon` | + +Units with fewer than `data.min_cells` cells are dropped. Tiles and polygons +inherit the label of their region unless the label file lists them. + +## Labels + +```csv +region_id,label,patient_id +S01_roi1,1,P01 +S01_roi2,1,P01 +S02,0,P02 +``` + +* `data.labels.id_column` holds unit, region or sample ids (a unit takes the + label of its own id, else of its region, else of its sample). +* `task.target_columns` selects the target column(s): one class column for + classification (numbers, or text encoded in sorted order), one or more + numeric columns for regression, `[time, event]` for survival. +* `data.labels.patient_column` groups units: random splits, cross-validation + folds and patient-level scores never separate a patient's units. + +Units without a label are skipped for supervised training and used for +pretraining. + +## Folds + +A JSON list of lists of unit (or region, or sample) ids, one list per outer +fold, grouped by patient: + +```json +[["S01_roi1", "S01_roi2", "S07"], ["S02", "S05_roi1"], ["S03", "S04"]] +``` + +Keep the file with the data so that every model is evaluated on the same folds. + +## Cell graphs + +The first command that needs the graphs builds them into +`data.processed_dir`: a Delaunay triangulation of every unit, with edge +features `log_distance` (the bounded log length +`clip((ln d - 2) / 3, 0, 1)`, which the encoder conditions on) and `neighbor` +(length at most `data.graph.neighbor_cutoff`). The directory is reused while +the settings that shaped it are unchanged; after a change, a command refuses +to reuse it and asks for `data.force_precompute=true` or a new +`data.processed_dir`. Labels and splits are not part of the graphs, so one +directory serves pretraining, every fold and every task. + +### Reducing many features + +For high-dimensional features (for example gene expression in `X`), +`data.feature_reducer.name=pca` projects them onto `n_components` principal +components fitted on all cells (`fit_on: all`), the training split +(`fit_on: train`, single random split only) or each sample (`fit_on: sample`). +To score a new cohort in the basis of a trained one, set +`data.feature_reducer.state_from=/processed/processed_index.json`. diff --git a/docs/dataset_preparation.md b/docs/dataset_preparation.md deleted file mode 100644 index 1809984..0000000 --- a/docs/dataset_preparation.md +++ /dev/null @@ -1,167 +0,0 @@ -# Dataset Preparation Guide - -This guide describes how to prepare a brand-new dataset so it can run through the `grass-mil` training pipeline. - -## 1) Prepare Raw Input Files - -`grass-mil` supports per-sample inputs in: - -- `csv` -- `tsv` -- `h5ad` -- `sce` / `rds` (SingleCellExperiment via `rpy2` + R) - -Each manifest row points to one input file. - -## 2) Create The Manifest CSV - -Default location: `data/raw/manifest.csv` (override with `data.raw_manifest_path=...`). - -Required columns: - -- `sample_id` -- `input_path` -- `input_type` (`csv`, `tsv`, `h5ad`, `sce`, `rds`) - -Optional columns: - -- `region_id` -- `polygons_path` - -Example: - -```csv -sample_id,input_path,input_type,region_id,polygons_path -sample_001,data/raw/sample_001.csv,csv,region_001, -sample_002,data/raw/sample_002.h5ad,h5ad,region_002,data/raw/sample_002_polygons.geojson -``` - -Important manifest rules: - -- `input_path` and `polygons_path` can be absolute or relative. -- Relative paths are resolved relative to the manifest file directory. -- If `region_id` is used, values must be globally unique. -- If a `sample_id` appears multiple times, each row for that sample must have a non-empty, unique `region_id`. - -## 3) Match Input Schema To Config - -Default parsing config lives in [`src/grass_mil/configs/data/spatial_omics.yaml`](../src/grass_mil/configs/data/spatial_omics.yaml). - -### CSV / TSV - -Default expectation: - -- coordinate columns: `x`, `y` (`data.csv.coord_columns`) -- optional cell id column: `data.csv.cell_id_column` -- optional categorical columns: `data.csv.categorical_label_columns` - -If `data.use_molecular_features=true` (default), molecular columns are required: - -- if `data.csv.molecular_columns=null`, all non-coordinate/non-id/non-categorical columns are treated as molecular features -- if that inferred set is empty, preprocessing fails - -### H5AD - -Configure under `data.h5ad.*`: - -- `coord_source=obsm` with `coord_key` (default `spatial`), or -- `coord_source=obs` with `coord_columns` -- optional `cell_id_column`, `categorical_label_columns`, `molecular_layer`, `molecular_features` - -### SCE / RDS - -Configure under `data.sce.*`: - -- coordinate source from `colData` (`coord_columns`) or `reducedDims` (`coord_key`) -- optional assay/feature selection controls - -## 4) Set Coordinate Units - -Spatial distances are normalized to micrometers during preprocessing. - -- `data.coord_scale_um=1.0` means input coordinates are already in micrometers -- set another scale if your inputs are in a different unit - -If you change `data.coord_scale_um`, regenerate processed artifacts (`data.force_precompute=true` or a new `data.processed_dir`). - -## 5) Optional Polygon And Graph Labels - -### Polygon files (`polygons_path`) - -Supported formats: - -- `.pkl` / `.pickle` -- `.wkt` / `.txt` (line-delimited WKT) -- `.json` / `.geojson` - -Behavior: - -- coordinates are filtered by polygon union -- when `data.sample_unit != full`, patches can be built from polygons - -### Graph labels (`data.graph_labels.*`) - -Use a CSV file to attach supervision targets at `sample`, `region`, or `patch` scope. - -```csv -REGION_ID,STAS,PATIENT -P001_core1,1,P001 -P001_core2,1,P001 -P002_core1,0,P002 -``` - -```bash -data.graph_labels.label_file=labels.csv data.graph_labels.id_column=REGION_ID \ -data.graph_labels.tasks=[STAS] data.graph_labels.patient_column=PATIENT -``` - -`patient_column` is never a target: it is attached to every graph as -`data.patient_id`, recorded in `processed_index.json`, and used by -`grass-mil-nested-cv` to keep a patient's regions in one fold and by -`grass-mil-predict` / `grass-mil-scores` for patient-level scores. - -### Fold files - -Cross-validation folds are a JSON list of lists of unit ids (regions by default): - -```json -[["P001_core1", "P001_core2", "P007_core1"], ["P002_core1", "..."], ["..."]] -``` - -Build them once, patient-grouped, and keep them with the data -(`grass_mil.data.components.splits.write_fold_json`); every regime is then -evaluated on identical folds. - -## 6) Run First Precompute/Training Pass - -Use a fresh processed directory when onboarding new data: - -```bash -grass-mil-train \ - task=finetune_mean \ - model=supervised_module \ - data=spatial_omics \ - data.raw_manifest_path=/absolute/path/to/manifest.csv \ - data.processed_dir=/absolute/path/to/processed/spatial_omics \ - data.force_precompute=true -``` - -Generated artifacts: - -- `${data.processed_dir}/processed_index.json` -- `${data.processed_dir}/metadata.json` -- `${data.processed_dir}/graphs/.../*.pt` - -## 7) When You Must Re-Precompute - -Rebuild processed artifacts if you change: - -- `data.split.*` (splits are persisted in `processed_index.json`) -- `data.coord_scale_um` -- feature reducer settings that affect embedding generation -- tiling / polygon behavior (`data.sample_unit`, `data.tiling.*`, `data.min_cells`) - -Use either: - -- `data.force_precompute=true`, or -- a new `data.processed_dir` diff --git a/docs/evaluation_stats.md b/docs/evaluation_stats.md deleted file mode 100644 index 1bdfefc..0000000 --- a/docs/evaluation_stats.md +++ /dev/null @@ -1,34 +0,0 @@ -# Evaluation Statistics (`grass-mil-scores`) - -Held-out scores of a cross-validation come from different fold models with -different score scales. `grass-mil-scores` standardises them within fold before -pooling, reports both the per-fold mean and the pooled AUROC with bootstrap -intervals, compares scores with the DeLong test and tabulates operating points -at pre-specified rules. - -```bash -grass-mil-scores input=runs/nested_ftall_reg/oof_predictions.csv output_dir=runs/scores -grass-mil-scores input=patients.csv scores=[grass,clinical] strata=[subtype] patient_column=null -``` - -Input: a CSV with one row per unit, a `label_column` (default `label`), one or -more score columns (`scores`, default `[prob_class_1]`), optionally a fold column -(`fold_column`, default `outer`), a patient column (`patient_column`, default -`patient_id`; with `patient_level: true` the scores are averaged per patient -first) and strata columns. - -| output | content | -|---|---| -| `scores_summary.csv` | per score: pooled AUROC on within-fold standardised scores with a percentile bootstrap CI (`n_boot`, `seed`), per-fold AUROC mean and sd, raw pooled AUROC | -| `delong.csv` | every pair of scores: delta AUROC, z, p (Sun & Xu structural components; same labels) | -| `strata.csv` | AUROC with CI within each level of each stratum (first score) | -| `operating_points.csv` | per score and rule: sensitivity, specificity, PPV, NPV with bootstrap CIs, LR+, LR-, confusion counts; rules = flag the top `q` within each fold (`operating_rules`) plus Youden / rule-in / rule-out thresholds chosen on all data (flagged optimistic) | -| `ppv_npv_by_prevalence.csv` | PPV and NPV of the middle rule at other prevalences | -| `report_scores.md` | all of the above as tables | - -`standardise=z|rank|none` selects the within-fold transform. The functions are -importable from `grass_mil.evaluation` (`auroc`, `bootstrap_auroc_ci`, -`delong_test`, `z_within_fold`, `rank_within_fold`, `per_fold_auroc`, -`stratified_auroc`, `flag_top_within_fold`, `operating_point`, -`operating_point_bootstrap`, `likelihood_ratios`, `ppv_npv_at_prevalence`, -`threshold_points`). diff --git a/docs/getting_started.md b/docs/getting_started.md deleted file mode 100644 index 5435c63..0000000 --- a/docs/getting_started.md +++ /dev/null @@ -1,189 +0,0 @@ -# Getting Started: Typical Usage On Your Own Cohort - -This page is the end-to-end path from cell tables to niches. Each step is one -command; the data description is written once in a cohort config and reused. -For the meaning of every setting see the -[configuration guide](configuration_guide.md); for a runnable demo on public -data see the [vignettes](../vignettes/README.md). - -## 0. What you need - -| Input | Form | Example | -|---|---|---| -| one cell table per region | CSV (or h5ad / SingleCellExperiment): one row per cell with coordinates, a cell-type column and optional continuous features such as cell size | `raw/P001_core1.csv` with `cell_id,x,y,cell_type,SIZE` | -| manifest | CSV listing the regions: `sample_id,input_path,input_type,region_id` | `raw/manifest.csv` | -| labels | CSV with a region id, the target and a patient column | `raw/labels.csv` with `REGION_ID,STAS,PATIENT` | -| folds | JSON list of lists of region ids, patient-grouped | `raw/folds.json`: `[["P001_core1","P001_core2",...], [...], ...]` | - -Coordinates can stay in pixels: set `coord_scale_um` to the factor that brings -them to the unit you want to express distances in, and give `neighbor_cutoff_um` -in that unit. Regions are whole cores, ROIs or slides (`sample_unit: full`); a -cohort of a few large images is cut into tiles instead (see the end of this page). - -Write the folds once and keep them with the data: - -```python -from grass_mil.data.components.splits import write_fold_json -# folds = [[...region ids...], ...] built from your patient ids, e.g. with -# sklearn.model_selection.StratifiedGroupKFold over regions grouped by patient -write_fold_json(folds, "raw/folds.json") -``` - -Install (`pip install -e ".[interpretability,io,stats,plots]"`), then check the -entry points: `grass-mil-train --help`. - -## 1. Describe the cohort once - -Copy the template and fill in paths and column names: - -```bash -cp src/grass_mil/configs/experiment/example_cohort.yaml src/grass_mil/configs/experiment/my_cohort.yaml -``` - -(With an installed wheel, put the file in any directory `my_configs/experiment/` -and add `--config-dir my_configs` to every command.) From here on every command -starts with `experiment=my_cohort`. Confirm the composed configuration before -running anything: - -```bash -grass-mil-train experiment=my_cohort --cfg job | less -``` - -## 2. Build the graphs - -The first run precomputes the Delaunay graphs of every region into -`processed_dir` (`processed_index.json` + one `.pt` per region) and reuses them -afterwards. A dry run does only that: - -```bash -grass-mil-train experiment=my_cohort +trainer.fast_dev_run=true -``` - -Check `processed/metadata.json`: number of regions, cell-type vocabulary, -continuous feature names. Re-run with `data.force_precompute=true` whenever an -input, a graph setting or `coord_scale_um` changes. - -## 3. Pretrain the encoder (label-free) - -```bash -grass-mil-train experiment=my_cohort regime=pretrain_a3 paths.output_dir=runs/pretrain_a3 trainer=gpu -``` - -Every region is used, no labels. Defaults: 200k steps, 64 ego-graphs per step, -4-hop ego-graphs, DropImportance + noise on `SIZE`. Shorten with -`task.total_steps=50000`; without a cell-size feature set -`task.augmentation.feature_noise_columns=[]`. Output: -`runs/pretrain_a3/checkpoints/last.ckpt` (plus one checkpoint every 5000 steps). - -## 4. Nested cross-validation - -```bash -grass-mil-nested-cv experiment=my_cohort regime=linear_ftall_reg \ - model.init_from_ckpt=runs/pretrain_a3/checkpoints/last.ckpt \ - paths.output_dir=runs/nested_ftall_reg trainer=gpu -``` - -Per outer fold this fits the probe head on the inner training regions, trains -three inner runs with region-level selection, refits on the whole development -set for the median best epoch and scores the held-out fold on every cell. -Expect about one minute per epoch for 300 regions of 3,000 cells on one GPU. -Resumable: re-run the same command after an interruption. - -Outputs in `runs/nested_ftall_reg/`: - -| file | content | -|---|---| -| `nested_cv_summary.json` | `auc_mean`, `auc_sd`, `auc_folds`, refit epochs, run directories | -| `oof_predictions.csv` | one held-out score per region (`outer, region_id, patient_id, label, prob_class_1`) | -| `oof_patient_predictions.csv` | mean of the region scores per patient | -| `outer_/refit/checkpoints/last.ckpt` | the fold models | - -Run the other regimes the same way (`regime=mil`, `linear_probehead`, -`linear_ftall`) into their own output directories; the probes are cached and -shared through `nested_cv.probe.cache_dir`. - -## 5. Evaluation statistics - -```bash -grass-mil-scores input=runs/nested_ftall_reg/oof_predictions.csv output_dir=runs/scores_ftall_reg -``` - -Patient-level (mean of regions) bootstrap AUROC interval, per-fold mean, -operating points at "flag the top q" rules; add a second table's scores with -`scores=[...]` to get DeLong comparisons. See -[evaluation_stats.md](evaluation_stats.md). - -## 6. Extract every cell through the fold models - -```bash -grass-mil-extract experiment=my_cohort regime=linear_ftall_reg \ - extract.nested_cv_summary=runs/nested_ftall_reg/nested_cv_summary.json \ - extract.output_dir=runs/cells trainer=gpu -``` - -Writes `runs/cells/fold/.npz` with the embedding, instance logits -and attention of every cell, plus the cell types, positions and graph. The -held-out regions of every fold model are recorded. - -## 7. Niches of one model, and replication across models - -```bash -grass-mil-niches cells_dir=runs/cells/fold2 output_dir=runs/niches_fold2 positive_name=STAS -grass-mil-metaclusters cells_dir=runs/cells output_dir=runs/meta positive_name=STAS \ - labels.label_file=/path/to/my_cohort/raw/labels.csv labels.label_column=STAS labels.patient_column=PATIENT -``` - -The first clusters the held-out cells of one fold model (PCA-10, L2, HDBSCAN, -15-NN assignment) and writes `clusters.csv`, `cells.csv`, `report.md` and the -figures. The second does it for every fold model, matches the clusters across -folds into metaclusters, tests their patient-level shares, transfers every -model to every region and fits a shared label-free basis. Both accept any -clustering backend (`clustering.method=kmeans clustering.n_clusters=12`). See -[niche_discovery.md](niche_discovery.md). - -## 8. Model-free spatial statistics - -```bash -grass-mil-spatial-stats cells_dir=runs/cells/fold0 output_dir=runs/spatial positive_name=STAS \ - labels.label_file=/path/to/my_cohort/raw/labels.csv labels.label_column=STAS labels.patient_column=PATIENT \ - reference=[tumor] partners=[myeloid,CD8,Treg] infiltrating=[myeloid,CD8] \ - group_map='{normal: tumor, ki67_cd8: CD8, cd8: CD8}' -``` - -Neighbourhood enrichment, cellular neighbourhoods, contact fractions and -cross-K per region, compared between classes at the patient level. See -[spatial_statistics.md](spatial_statistics.md). - -## Variants - -**One split instead of nested CV.** Give explicit id lists: - -```bash -grass-mil-train experiment=my_cohort regime=mil model.init_from_ckpt=runs/pretrain_a3/checkpoints/last.ckpt \ - data.split.explicit.train=raw/train.json data.split.explicit.val=raw/val.json data.split.explicit.test=raw/test.json -``` - -For a `linear_*` regime, first `grass-mil-probe` with the same splits and pass -`task.head_init_path=`. - -**Patient-level prediction from one model** (`mean_probability` = mean of the -patient's regions, `attention_weighted` = one merged bag): - -```bash -grass-mil-predict experiment=my_cohort regime=linear_ftall_reg ckpt_path=runs/nested_ftall_reg/outer_2/refit/checkpoints/last.ckpt \ - aggregation.bag_scope=patient aggregation.mode=mean_probability \ - aggregation.patient_map.label_file=/path/to/my_cohort/raw/labels.csv -``` - -**Tiles as units** (a few large images): `data.sample_unit=tile`, -`data.tiling.tile_size_um=...`, labels keyed by `patch_id` with -`data.graph_labels.scope=patch`, folds of patch ids, `nested_cv.id_kind=patch`. -Vignette 3 does exactly this. - -**h5ad inputs**: replace the `csv:` block by `h5ad: {coord_source: obsm, -coord_key: spatial, categorical_label_columns: [cell_type]}` and use -`input_type: h5ad` in the manifest; the marker matrix becomes the continuous -features (`molecular_features` selects a subset). - -**Cell type only** (no continuous feature): `data.csv.molecular_columns=[]` -and `task.augmentation.feature_noise_columns=[]` for pretraining. diff --git a/docs/hyperparameter_reference.md b/docs/hyperparameter_reference.md deleted file mode 100644 index 5f6c4e5..0000000 --- a/docs/hyperparameter_reference.md +++ /dev/null @@ -1,907 +0,0 @@ -# Hyperparameter Reference - -This reference covers repo-defined configuration keys and behavior for: - -- `configs/train.yaml` -- `configs/eval.yaml` -- `configs/loocv.yaml` -- `configs/probe.yaml`, `configs/nested_cv.yaml`, `configs/extract.yaml`, `configs/scores.yaml`, `configs/niches.yaml`, `configs/metaclusters.yaml`, `configs/spatial_stats.yaml` (analysis CLIs, section 12b) -- `configs/regime/*.yaml` (training regime presets, section 12a) -- `configs/experiment/grass_mil_pretrain_a3.yaml` -- `configs/data/spatial_omics.yaml` -- `configs/task/*.yaml` -- `configs/model/**/*.yaml` (presets `grass_mil`, `bgrl_a3`, generic modules, `encoder/`, `attention/` groups) -- `configs/augmentation/*.yaml` -- `configs/data/sampler/*.yaml` -- `configs/optim/*.yaml` -- `configs/scheduler/*.yaml` -- `configs/trainer/*.yaml` -- `configs/callbacks/*.yaml` -- `configs/logger/*.yaml` -- `configs/debug/*.yaml` -- `configs/hydra/default.yaml` -- `configs/extras/default.yaml` - -> Throughout this document `configs/...` is shorthand for -> `src/grass_mil/configs/...`. The config tree lives inside the package so it -> ships with the wheel; Hydra resolves it relative to the entrypoint module. - -Override syntax examples use: - -```bash -grass-mil-train key.path=value -``` - -## 1) Top-Level Train/Eval Composition Keys - -## 1.1 `configs/train.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `defaults` | composition list | list | Hydra config groups | always | `grass-mil-train trainer=gpu` | -| `task_name` | `train` | str | any string | always | `grass-mil-train task_name=my_run` | -| `tags` | `[dev]` | list[str] | any list | always | `grass-mil-train tags="[paper,ablation]"` | -| `train` | `True` | bool | `true/false` | always | `grass-mil-train train=false` | -| `test` | `True` | bool | `true/false` | after fit | `grass-mil-train test=false` | -| `ckpt_path` | `null` | str/null | path or null | resume training when set | `grass-mil-train ckpt_path=/abs/last.ckpt` | -| `seed` | `null` | int/null | integer or null | random seeding | `grass-mil-train seed=42` | - -## 1.2 `configs/eval.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `defaults` | composition list | list | Hydra config groups | always | `grass-mil-eval trainer=gpu` | -| `task_name` | `eval` | str | any string | always | `grass-mil-eval task_name=holdout_eval` | -| `tags` | `[dev]` | list[str] | any list | always | `grass-mil-eval tags="[final_eval]"` | -| `ckpt_path` | `???` | str | required absolute/relative checkpoint path | always | `grass-mil-eval ckpt_path=/abs/model.ckpt` | - -## 1.3 `configs/loocv.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `task_name` | `loocv` | str | any string | LOOCV entrypoint | `grass-mil-loocv task_name=loocv_run` | -| `loocv.mode` | `single` | str | `single`, `all` | LOOCV entrypoint | `grass-mil-loocv loocv.mode=all` | -| `loocv.fold_index` | `null` | int/null | `>=0` or null | single mode fold selection | `grass-mil-loocv loocv.mode=single loocv.fold_index=3` | -| `loocv.per_fold_processed_dir` | `true` | bool | `true/false` | fold isolation | `grass-mil-loocv loocv.per_fold_processed_dir=false` | -| `loocv.force_precompute_per_fold` | `true` | bool | `true/false` | leakage-safe fold preprocessing | `grass-mil-loocv loocv.force_precompute_per_fold=true` | -| `loocv.summary_filename` | `loocv_summary.json` | str | filename | fold summary artifact naming | `grass-mil-loocv loocv.summary_filename=my_summary.json` | - -## 2) Data Config (`configs/data/spatial_omics.yaml`) - -## 2.1 Datamodule Core - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `data._target_` | `grass_mil.data.spatial_omics_datamodule.SpatialOmicsDataModule` | str | import path | always | `grass-mil-train data._target_=...` | -| `data.data_dir` | `${paths.data_dir}` | str | path | always | `grass-mil-train data.data_dir=/tmp/data` | -| `data.raw_manifest_path` | `${paths.data_dir}/raw/manifest.csv` | str | path | always | `grass-mil-train data.raw_manifest_path=/abs/manifest.csv` | -| `data.processed_dir` | `${paths.data_dir}/processed/spatial_omics` | str | path | always | `grass-mil-train data.processed_dir=/abs/processed` | -| `data.batch_size` | `1` | int | `>=1` | loader batching | `grass-mil-train data.batch_size=4` | -| `data.num_workers` | `0` | int | `>=0` | dataloader workers | `grass-mil-train data.num_workers=8` | -| `data.pin_memory` | `False` | bool | `true/false` | dataloader | `grass-mil-train data.pin_memory=true` | -| `data.coord_scale_um` | `1.0` | float | `>0` | preprocessing (`1.0` means input coords already in `um`) | `grass-mil-train data.coord_scale_um=0.5` | -| `data.sample_unit` | `tile` | str | `tile`/`full` | patch generation | `grass-mil-train data.sample_unit=full` | -| `data.reducer_scope` | `sample` | str | `sample`/`dataset` | reducer fitting | `grass-mil-train data.reducer_scope=dataset` | -| `data.keep_raw_molecular` | `false` | bool | `true/false` | precompute outputs | `grass-mil-train data.keep_raw_molecular=true` | -| `data.force_precompute` | `false` | bool | `true/false` | precompute trigger | `grass-mil-train data.force_precompute=true` | -| `data.min_cells` | `10` | int | `>=0` | patch filtering | `grass-mil-train data.min_cells=25` | -| `data.use_molecular_features` | `true` | bool | `true/false` | loader/reducer behavior | `grass-mil-train data.use_molecular_features=false` | - -## 2.2 Categorical Feature Encoding - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `data.categorical_features.include_labels` | `[cell_type]` | list[str] | label names present in loaded data | attach `categorical_codes` | `grass-mil-train data.categorical_features.include_labels="[cell_type,state]"` | - -## 2.3 Split Config - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `data.split.train_val_test_split` | `[0.8,0.1,0.1]` | list[float] len=3 | must sum to `1.0` | precompute split assignment | `grass-mil-train data.split.train_val_test_split="[0.7,0.15,0.15]"` | -| `data.split.split_by` | `sample` | str | `sample`/`region`/`patch` | split semantics | `grass-mil-train data.split.split_by=region` | -| `data.split.loocv.enabled` | `false` | bool | `true/false` | LOOCV override path | `grass-mil-train data.split.loocv.enabled=true` | -| `data.split.loocv.fold_unit` | `region` | str | `region`/`sample` | LOOCV override path | `grass-mil-train data.split.loocv.fold_unit=sample` | -| `data.split.loocv.holdout_id` | `null` | str/null | canonical fold id (`sample` or `sample::region`) | LOOCV override path | `grass-mil-train data.split.loocv.holdout_id=s1::r3` | -| `data.split.loocv.validation_strategy` | `heldout_fold_items` | str | `heldout_fold_items`/`patches_from_train_items` | LOOCV override path | `grass-mil-train data.split.loocv.validation_strategy=patches_from_train_items` | -| `data.split.loocv.val_ratio` | `0.1` | float | `[0.0,1.0)` | LOOCV validation split sizing | `grass-mil-train data.split.loocv.val_ratio=0.2` | -| `data.split.loocv.seed` | `42` | int | integer | LOOCV validation split determinism | `grass-mil-train data.split.loocv.seed=7` | - -Notes: - -- When LOOCV is disabled, behavior is identical to current non-LOOCV splitting. -- `data.coord_scale_um` is applied at preprocessing time; if changed, processed artifacts must be regenerated (`data.force_precompute=true` or a fresh `data.processed_dir`). -- Split labels are persisted at precompute time in `processed_index.json`. - Changing `data.split.*` requires regenerating processed artifacts - (`data.force_precompute=true` or a new `data.processed_dir`). -- For `reducer_scope=dataset` + `feature_reducer.fit_mode=train_only`, LOOCV requires fold-specific precompute (`force_precompute=true` or fold-specific `processed_dir`). -- For safe fold isolation, at least one of `loocv.force_precompute_per_fold` or `loocv.per_fold_processed_dir` must be `true`. -- In `loocv.mode=single`, fold selection can use either - `data.split.loocv.holdout_id` or `loocv.fold_index`. If both are set, - `loocv.fold_index` takes priority and a warning is emitted. - -### Explicit splits (`data.split.explicit.*`) - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `data.split.explicit.train` | `null` | list/str | unit ids, or a `.json` / `.txt` file of ids | re-partition an existing precompute | `grass-mil-train data.split.explicit.train=train.json` | -| `data.split.explicit.val` | `null` | list/str | as above (may equal `test` for refits) | | `data.split.explicit.val=[r1,r2]` | -| `data.split.explicit.test` | `null` | list/str | as above | | `data.split.explicit.test=test.json` | -| `data.split.explicit.id_kind` | `region` | str | `region`, `sample`, `patch` | explicit split set | `data.split.explicit.id_kind=patch` | - -Explicit lists override the persisted splits without re-processing the graphs -(ids not listed are dropped); they need `feature_reducer.name=identity` or -`fit_mode=global`. `grass-mil-nested-cv` and `grass-mil-probe` set them per fold. - -## 2.4 Manifest Field Mapping - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `data.manifest.sample_id` | `sample_id` | str | manifest column name | manifest parse | `grass-mil-train data.manifest.sample_id=sid` | -| `data.manifest.input_path` | `input_path` | str | manifest column name | manifest parse | `grass-mil-train data.manifest.input_path=file_path` | -| `data.manifest.input_type` | `input_type` | str | manifest column name | manifest parse | `grass-mil-train data.manifest.input_type=source_type` | -| `data.manifest.region_id` | `region_id` | str | manifest column name | optional region grouping | `grass-mil-train data.manifest.region_id=roi` | -| `data.manifest.polygons_path` | `polygons_path` | str | manifest column name | optional polygons | `grass-mil-train data.manifest.polygons_path=poly_path` | - -## 2.5 CSV Loader Keys - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `data.csv.sep` | `,` | str | delimiter string | `input_type=csv/tsv` | `grass-mil-train data.csv.sep=';'` | -| `data.csv.coord_columns` | `[x,y]` | list[str] len=2 | existing columns | `input_type=csv/tsv` | `grass-mil-train data.csv.coord_columns="[x_um,y_um]"` | -| `data.csv.cell_id_column` | `cell_id` | str/null | existing column or null | `input_type=csv/tsv` | `grass-mil-train data.csv.cell_id_column=id` | -| `data.csv.categorical_label_columns` | `[]` | list[str] | existing columns | `input_type=csv/tsv` | `grass-mil-train data.csv.categorical_label_columns="[cell_type]"` | -| `data.csv.molecular_columns` | `null` | list[str]/null | existing columns or null infer | `input_type=csv/tsv` and `use_molecular_features=true` | `grass-mil-train data.csv.molecular_columns="[gene1,gene2]"` | - -## 2.6 H5AD Loader Keys - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `data.h5ad.coord_source` | `obsm` | str | `obs`/`obsm` | `input_type=h5ad` | `grass-mil-train data.h5ad.coord_source=obs` | -| `data.h5ad.coord_key` | `spatial` | str/null | obsm key | `coord_source=obsm` | `grass-mil-train data.h5ad.coord_key=X_spatial` | -| `data.h5ad.coord_columns` | `[x,y]` | list[str] len=2 | obs columns | `coord_source=obs` | `grass-mil-train data.h5ad.coord_columns="[x_um,y_um]"` | -| `data.h5ad.cell_id_column` | `null` | str/null | obs column or null | `input_type=h5ad` | `grass-mil-train data.h5ad.cell_id_column=cell_id` | -| `data.h5ad.categorical_label_columns` | `[]` | list[str] | obs columns | `input_type=h5ad` | `grass-mil-train data.h5ad.categorical_label_columns="[cell_type]"` | -| `data.h5ad.molecular_layer` | `null` | str/null | layer key or null (`X`) | `use_molecular_features=true` | `grass-mil-train data.h5ad.molecular_layer=counts` | -| `data.h5ad.molecular_features` | `null` | list[str]/null | var names subset | `use_molecular_features=true` | `grass-mil-train data.h5ad.molecular_features="[CD3,CD8]"` | - -## 2.7 SCE Loader Keys - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `data.sce.assay_name` | `null` | str/null | assay name or default assay | `input_type=sce/rds` | `grass-mil-train data.sce.assay_name=logcounts` | -| `data.sce.coord_source` | `colData` | str | `colData`/`reducedDims` | `input_type=sce/rds` | `grass-mil-train data.sce.coord_source=reducedDims` | -| `data.sce.coord_key` | `null` | str/null | reducedDims key | `coord_source=reducedDims` | `grass-mil-train data.sce.coord_key=spatial` | -| `data.sce.coord_columns` | `[x,y]` | list[str]/null | colData columns | `coord_source=colData` | `grass-mil-train data.sce.coord_columns="[x_um,y_um]"` | -| `data.sce.cell_id_column` | `null` | str/null | colData column or colnames fallback | `input_type=sce/rds` | `grass-mil-train data.sce.cell_id_column=cell_id` | -| `data.sce.categorical_label_columns` | `[]` | list[str] | colData columns | `input_type=sce/rds` | `grass-mil-train data.sce.categorical_label_columns="[cell_type]"` | -| `data.sce.transpose_assay` | `true` | bool | `true/false` | molecular matrix orientation | `grass-mil-train data.sce.transpose_assay=false` | -| `data.sce.molecular_features` | `null` | list[str]/null | rowname subset | `use_molecular_features=true` | `grass-mil-train data.sce.molecular_features="[geneA,geneB]"` | - -## 2.8 Graph Builder Keys - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `data.graph_builder.name` | `delaunay` | str | registered builder name/import path | graph construction | `grass-mil-train data.graph_builder.name=delaunay` | -| `data.graph_builder.kwargs.edge_features` | `[distance,neighbor]` | list[str] | `distance`, `neighbor` | `delaunay` builder | `grass-mil-train data.graph_builder.kwargs.edge_features="[distance]"` | -| `data.graph_builder.kwargs.neighbor_cutoff_um` | `20.0` | float | `>0` | when `neighbor` edge feature enabled | `grass-mil-train data.graph_builder.kwargs.neighbor_cutoff_um=50.0` | - -## 2.9 Feature Reducer Keys - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `data.feature_reducer.name` | `identity` | str | `identity`, `pca`, custom import path | molecular reduction | `grass-mil-train data.feature_reducer.name=pca` | -| `data.feature_reducer.fit_mode` | `train_only` | str | `train_only`/`global` | dataset-scope reducer | `grass-mil-train data.feature_reducer.fit_mode=global` | -| `data.feature_reducer.kwargs` | `{}` | dict | reducer kwargs | reducer-specific | `grass-mil-train data.feature_reducer.kwargs.n_components=64` | - -## 2.10 Sampler Keys - -`data.sampler` is a config group (`data/sampler=shadow_custom` by default, -`ego_graph_mil` for MIL fine-tuning, `ego_graph_pretrain` for BGRL). Single -keys are overridden as `data.sampler.runtime.=...`. - -| Key | Default (`shadow_custom`) | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `data.sampler.name` | `shadow_custom` | str | `identity`, `shadow_native`, `shadow_custom` | dataloader build | `grass-mil-train data.sampler.name=identity` | -| `data.sampler.kwargs` | `{}` | dict | strategy kwargs | strategy-specific | `grass-mil-train data.sampler.kwargs='{}'` | -| `data.sampler.runtime.enabled` | `true` | bool | `true/false` | shadow strategies | `grass-mil-train data.sampler.runtime.enabled=false` | -| `data.sampler.runtime.depth` | `2` (`4` in `ego_graph_*`) | int | `>=1` | ego-graph hops | `grass-mil-train data.sampler.runtime.depth=3` | -| `data.sampler.runtime.num_neighbors` | `8` (`-1` in `ego_graph_*`) | int | `>=1`, or `-1` = every neighbour at every hop | shadow runtime | `grass-mil-train data.sampler.runtime.num_neighbors=-1` | -| `data.sampler.runtime.subgraph_batch_size` | `32` | int | `>=1` | ego-graphs per chunk (upper bound when `unit_as_bag`) | `grass-mil-train data.sampler.runtime.subgraph_batch_size=64` | -| `data.sampler.runtime.replace` | `false` | bool | `true/false` | root sampling | `grass-mil-train data.sampler.runtime.replace=true` | -| `data.sampler.runtime.shuffle_subgraphs` | `true` | bool | `true/false` | training shuffle behavior | `grass-mil-train data.sampler.runtime.shuffle_subgraphs=false` | -| `data.sampler.runtime.proportional_root_sampling` | `true` (`false` in `ego_graph_*`) | bool | `true/false` | cell-type weighted roots | `grass-mil-train data.sampler.runtime.proportional_root_sampling=false` | -| `data.sampler.runtime.property_name` | `cell_type` | str | categorical label present in `categorical_slices` | weighted root sampling | `grass-mil-train data.sampler.runtime.property_name=cell_type` | -| `data.sampler.runtime.weight_mode` | `inverse` | str | `inverse`, `sqrt_inverse`, `proportional` | weighted root sampling | `grass-mil-train data.sampler.runtime.weight_mode=sqrt_inverse` | -| `data.sampler.runtime.min_weight` | `1.0e-6` | float | `>0` | weighted root sampling | `grass-mil-train data.sampler.runtime.min_weight=1.0e-4` | -| `data.sampler.runtime.subsample_fraction` | `1.0` | float | `(0, 1]` | fraction of candidate roots per unit per epoch | `grass-mil-train data.sampler.runtime.subsample_fraction=0.5` | -| `data.sampler.runtime.subsample_count` | `null` (`256` in `ego_graph_mil`) | int/null | `>=1` or null | absolute roots per unit per epoch (uniform, without replacement; overrides the fraction) | `grass-mil-train data.sampler.runtime.subsample_count=512` | -| `data.sampler.runtime.subsample_seed` | `null` | int/null | | seed of the root subsampling generator | `grass-mil-train data.sampler.runtime.subsample_seed=0` | -| `data.sampler.runtime.fixed_per_unit_seed` | `null` | int/null | | identical root subset per unit on every pass (validation) | `grass-mil-train data.val_sampler.runtime.fixed_per_unit_seed=0` | -| `data.sampler.runtime.unit_as_bag` | `false` (`true` in `ego_graph_mil`) | bool | `true/false` | one batch per unit holding all its selected roots (one bag, one softmax, one loss) | `grass-mil-train data.sampler.runtime.unit_as_bag=true` | -| `data.sampler.runtime.merge_units` | `false` (`true` in `ego_graph_pretrain`) | bool | `true/false` | pool all units into one graph and draw roots from the union (BGRL) | `grass-mil-train data.sampler.runtime.merge_units=true` | -| `data.sampler.runtime.interior_seeds.enabled` | `false` | bool | `true/false` | shadow runtime | `grass-mil-train data.sampler.runtime.interior_seeds.enabled=true` | -| `data.sampler.runtime.interior_seeds.n_hops` | `null` | int/null | `>=1`, `null` uses `depth` | `interior_seeds.enabled=true` | `grass-mil-train data.sampler.runtime.interior_seeds.n_hops=3` | -| `data.sampler.runtime.interior_seeds.seed_tolerance_multiplier` | `1.0` | float | `>=0` | `interior_seeds.enabled=true` | `grass-mil-train data.sampler.runtime.interior_seeds.seed_tolerance_multiplier=1.5` | -| `data.sampler.runtime.interior_seeds.keep_largest_component` | `true` | bool | `true/false` | `interior_seeds.enabled=true` | `grass-mil-train data.sampler.runtime.interior_seeds.keep_largest_component=false` | -| `data.sampler.runtime.interior_seeds.cache_dir` | `null` | str/null | path, `null` disables caching | `interior_seeds.enabled=true` | `grass-mil-train data.sampler.runtime.interior_seeds.cache_dir=.cache/seeds` | -| `data.sampler.runtime.ego_radius.enabled` | `false` | bool | `true/false` | shadow runtime | `grass-mil-train data.sampler.runtime.ego_radius.enabled=true` | -| `data.sampler.runtime.ego_radius.radius_per_hop` | `75.0` | float | `>0`, datamodule units | `ego_radius.enabled=true` | `grass-mil-train data.sampler.runtime.ego_radius.radius_per_hop=30.0` | -| `data.sampler.runtime.ego_radius.radius_offset` | `55.0` | float | `>=0`, datamodule units | `ego_radius.enabled=true` | `grass-mil-train data.sampler.runtime.ego_radius.radius_offset=22.0` | -| `data.sampler.runtime.ego_radius.keep_root_component` | `true` | bool | `true/false` | `ego_radius.enabled=true` | `grass-mil-train data.sampler.runtime.ego_radius.keep_root_component=false` | -| `data.val_sampler` | `null` (fixed 1024-root subset in `regime=mil`) | dict/null | full sampler config | validation loader | see `configs/regime/mil.yaml` | -| `data.test_sampler` | `null` | dict/null | full sampler config; `null` derives an all-roots, unweighted, no-bagging variant of `sampler` | test / predict loaders | `grass-mil-train data.test_sampler.runtime.subsample_count=1024` | - -Two notes on these, both of which cause silent misconfiguration rather than errors: - -- **`interior_seeds`** restricts ego-graph roots to the tissue interior. - Ego-graphs rooted at the section edge are truncated by the slide boundary - rather than by biology, so they are not comparable to interior ones. Only the - `n_hops` rim around the convex hull is excluded; internal holes are not - treated as boundary. The interior mask is intersected into the candidate pool - *before* class-proportional weighting. -- **`ego_radius` distances are in datamodule coordinate units** (µm after - `coord_scale_um`), **not pixels**. The `75.0`/`55.0` defaults come from the - NSCLC cohort at its native pixel scale; convert your cohort's values before - setting them, or the cap will be wrong by whatever the scale factor is. - -## 2.11 Tiling And Graph Label Keys - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `data.tiling.tile_size_um` | `200.0` | float | `>0` | `sample_unit=tile` | `grass-mil-train data.tiling.tile_size_um=300.0` | -| `data.tiling.stride_um` | `200.0` | float | `>0` | `sample_unit=tile` | `grass-mil-train data.tiling.stride_um=150.0` | -| `data.tiling.min_cells` | `10` | int | `>=0` | tile retention | `grass-mil-train data.tiling.min_cells=20` | -| `data.graph_labels.label_file` | `null` | str/null | CSV path or null | optional graph labels | `grass-mil-train data.graph_labels.label_file=/abs/labels.csv` | -| `data.graph_labels.id_column` | `id` | str | column in label file | graph labels enabled | `grass-mil-train data.graph_labels.id_column=region_id` | -| `data.graph_labels.tasks` | `null` | list[str]/null | task columns or null for all | graph labels enabled | `grass-mil-train data.graph_labels.tasks="[outcome]"` | -| `data.graph_labels.scope` | `region` | str | `region`, `patch`, `sample` | graph labels enabled | `grass-mil-train data.graph_labels.scope=sample` | -| `data.graph_labels.patient_column` | `null` | str/null | column of the label file | graph labels enabled | `grass-mil-train data.graph_labels.patient_column=PATIENT` (attached as `data.patient_id`, recorded in the index; never a target) | -| `data.transforms` | `[]` | list[dict] | importable transform specs | transform pipeline | `grass-mil-train data.transforms='[{name: grass_mil.data.components.transforms.CompositionVector, kwargs: {label_attr: label_cell_type}}]'` | - -## 2.12 Sampler Preset Files (`configs/data/sampler/*.yaml`) - -| File | Keys In File | Defaults | -|---|---|---| -| `configs/data/sampler/identity.yaml` | `name`, `kwargs` | `name=identity`, `kwargs={}` | -| `configs/data/sampler/shadow_native.yaml` | `name`, `kwargs`, `runtime.*` | `name=shadow_native`, runtime defaults mirror `data.sampler.runtime.*` | -| `configs/data/sampler/shadow_custom.yaml` | `name`, `kwargs`, `runtime.*` | `name=shadow_custom`, runtime defaults mirror `data.sampler.runtime.*` | -| `configs/data/sampler/ego_graph_mil.yaml` | `name`, `kwargs`, `runtime.*` | GRASS-MIL fine-tuning: depth 4, unlimited neighbours, uniform roots, `subsample_count=256`, `unit_as_bag=true` | -| `configs/data/sampler/ego_graph_pretrain.yaml` | `name`, `kwargs`, `runtime.*` | BGRL: depth 4, unlimited neighbours, 64 roots per step from the union of all units (`merge_units=true`) | - -## 3) Task Config Group (`configs/task/*.yaml`) - -## 3.1 `configs/task/finetune_mean.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `task.name` | `finetune_mean` | str | label | always | `grass-mil-train task.name=custom_mean` | -| `task.aggregation` | `mean` | str | `mean` | supervised module | `grass-mil-train task.aggregation=mean` | -| `task.bag_key` | `region_id` | str | batch attribute name | supervised bagging | `grass-mil-train task.bag_key=sample_id` | -| `task.bag_fallback_key` | `sample_id` | str | batch attribute name | bag key missing fallback | `grass-mil-train task.bag_fallback_key=patch_id` | -| `task.max_instances_per_bag` | `0` | int | `0` means all, otherwise `>=1` | bag aggregation | `grass-mil-train task.max_instances_per_bag=64` | -| `task.instance_sampling` | `all` | str | `all`, `random` | if max instances > 0 | `grass-mil-train task.instance_sampling=random` | -| `task.target_type` | `binary` | str | `binary`, `regression`, `survival` | loss routing | `grass-mil-train task.target_type=regression` | -| `task.target_columns` | `null` | list[str]/null | label names | target selection | `grass-mil-train task.target_columns="[outcome]"` | -| `task.loss` | `categorical_bce` | str | model loss presets | loss build | `grass-mil-train task.loss=regression_mse` | -| `task.lr_warmup.enabled` | `false` | bool | `true/false` | scheduler wrapping | `grass-mil-train task.lr_warmup.enabled=true` | -| `task.lr_warmup.warmup_steps` | `0` | int | `>=0` | warmup enabled | `grass-mil-train task.lr_warmup.warmup_steps=100` | -| `task.lr_warmup.start_factor` | `0.1` | float | `(0,1]` (clamped) | warmup enabled | `grass-mil-train task.lr_warmup.start_factor=0.2` | - -## 3.2 `configs/task/finetune_mil.yaml` - -Includes the keys of 3.1 (with `aggregation=mil_attention`, `target_type=categorical`, -`loss=categorical_ce`, and per-class widths `model.graph_head.output_dim=2`, -`model.attention.n_classes=2`) plus the MIL controls below. The `regime=` -presets (section 12a) set them to the preprint values. - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `task.bag_key` | `region_id` | str | `region_id`, `patch_id`, `sample_id` | what one bag is | `grass-mil-train task.bag_key=patch_id` | -| `task.class_weights` | `auto_train` | str/list/null | `auto_train` = inverse frequency on the training units (mean 1); a list; `null` | categorical loss | `grass-mil-train task.class_weights=[1.0,2.0]` | -| `task.optimization.encoder_lr_scale` | `1.0` | float | `>=0`; `0` freezes the encoder (no gradients, eval mode) | optimizer groups | `grass-mil-train task.optimization.encoder_lr_scale=0.3` | -| `task.optimization.head_lr_scale` | `1.0` | float | `>=0` | optimizer groups | `grass-mil-train task.optimization.head_lr_scale=0.01` | -| `task.optimization.attention_lr_scale` | `2.0` | float | `>=0` | optimizer groups | `grass-mil-train task.optimization.attention_lr_scale=1.0` | -| `task.optimization.head_weight_decay` | `null` | float/null | overrides `optim.weight_decay` for the head group | optimizer groups | `grass-mil-train task.optimization.head_weight_decay=1e-3` | -| `task.optimization.backbone_lr` / `attention_lr` | `null` | float/null | legacy absolute pair (both or neither); ignored when scales are used | | | -| `task.finetune_augmentation.drop_edge_p` | `0.0` | float | `[0,1]` | training batches only; edges dropped with both directions | `grass-mil-train task.finetune_augmentation.drop_edge_p=0.05` | -| `task.finetune_augmentation.drop_feat_p` | `0.0` | float | `[0,1]` | elementwise feature drop, writes the unassigned cell-type token | `grass-mil-train task.finetune_augmentation.drop_feat_p=0.05` | -| `task.finetune_augmentation.seed` | `0` | int | | augmentation generator | | -| `task.head_init_path` | `null` | str/null | probe head file (`graph_head.net.0.*`), or `auto` (resolved per split by `grass-mil-nested-cv`) | linear head (`model.graph_head.num_layers=1`) | `grass-mil-train task.head_init_path=probe_head.pt` | -| `task.freeze_head` | `false` | bool | | head excluded from the optimizer | `grass-mil-train task.freeze_head=true` | -| `task.grad_clip_norm` | `1.0` | float/null | `>0` or null | gradient-norm clipping in the region flush | `grass-mil-train task.grad_clip_norm=null` | -| `task.scheduler_interval` | `step` | str | `step`, `epoch` | with `scheduler=cosine_mil_step` (`T_max: auto`) the cosine spans every optimizer step of the run; skipped AMP steps do not advance it | `grass-mil-train task.scheduler_interval=epoch` | -| `task.region_accumulation.enabled` | `true` | bool | `true/false` | MIL training | `grass-mil-train task.region_accumulation.enabled=false` | -| `task.region_accumulation.hyperbatch_size` | `12` | int | `>=1` | regions per optimizer step | `grass-mil-train task.region_accumulation.hyperbatch_size=8` | -| `task.region_accumulation.flush_on_epoch_end` | `true` | bool | | | | -| `task.max_instances_per_bag` / `instance_sampling` | `0` / `all` | | as in 3.1 | | | - -Region-level metrics are logged per epoch (`val/loss_region`, `val/auroc`, -`val/bal_acc`, `val/acc`, `val/pred_pos_frac`; the same under `test/`) and written -to `metrics_per_epoch.jsonl`; `test_predictions.csv` holds one row per bag. - -## 3.3 `configs/task/pretrain_bgrl.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `defaults` | `override /scheduler: cosine_step` | Hydra defaults entry | scheduler group override | pretrain task composition | `grass-mil-train task=pretrain_bgrl model=bgrl_module` | -| `task.name` | `pretrain_bgrl` | str | label | always | `grass-mil-train task.name=ssl_run` | -| `task.aggregation` | `mean` | str | `mean` | BGRL module ignores bag loss semantics but keeps task schema | `grass-mil-train task.aggregation=mean` | -| `task.bag_key` | `region_id` | str | batch attr | task metadata | `grass-mil-train task.bag_key=sample_id` | -| `task.bag_fallback_key` | `sample_id` | str | batch attr | task metadata | `grass-mil-train task.bag_fallback_key=patch_id` | -| `task.max_instances_per_bag` | `0` | int | `0`/`>=1` | task metadata | `grass-mil-train task.max_instances_per_bag=0` | -| `task.instance_sampling` | `all` | str | `all`/`random` | task metadata | `grass-mil-train task.instance_sampling=all` | -| `task.target_type` | `binary` | str | task labels | task metadata | `grass-mil-train task.target_type=binary` | -| `task.target_columns` | `null` | list/null | optional | task metadata | `grass-mil-train task.target_columns='[label]'` | -| `task.loss` | `categorical_bce` | str | loss label | task metadata | `grass-mil-train task.loss=categorical_bce` | -| `task.lr_warmup.enabled` | `true` | bool | `true/false` | optimizer scheduler warmup | `grass-mil-train task.lr_warmup.enabled=true` | -| `task.lr_warmup.warmup_steps` | `1000` | int | `>=0` | warmup enabled | `grass-mil-train task.lr_warmup.warmup_steps=500` | -| `task.lr_warmup.start_factor` | `0.1` | float | `(0,1]` | warmup enabled | `grass-mil-train task.lr_warmup.start_factor=0.2` | -| `task.warmup_steps` | `1000` | int | `>=0` | BGRL momentum scheduler | `grass-mil-train task.warmup_steps=500` | -| `task.total_steps` | `10000` | int | `>0` preferred | BGRL momentum scheduler | `grass-mil-train task.total_steps=20000` | -| `task.momentum` | `0.99` | float | `[0,1]` practical range | target network EMA | `grass-mil-train task.momentum=0.996` | -| `task.momentum_min` | `0.99` | float | `[0,1]` practical range | cosine momentum floor | `grass-mil-train task.momentum_min=0.9` | -| `task.drop_edge_p1` | `0.2` | float | `[0,1]` | view-1 edge dropout | `grass-mil-train task.drop_edge_p1=0.1` | -| `task.drop_edge_p2` | `0.2` | float | `[0,1]` | view-2 edge dropout | `grass-mil-train task.drop_edge_p2=0.3` | -| `task.drop_feat_p1` | `0.2` | float | `[0,1]` | view-1 feature dropout | `grass-mil-train task.drop_feat_p1=0.1` | -| `task.drop_feat_p2` | `0.2` | float | `[0,1]` | view-2 feature dropout | `grass-mil-train task.drop_feat_p2=0.3` | - -`task.momentum_min` is the legacy name of `task.momentum_final`. - -## 3.4 `configs/task/pretrain_bgrl_a3.yaml` (GRASS-MIL pretraining) - -Selected by `regime=pretrain_a3` (with `model=bgrl_a3`, `callbacks=pretrain_checkpoint`, -`data/sampler=ego_graph_pretrain`); `experiment=grass_mil_pretrain_a3` is an alias. - -| Key | Default | Meaning | -|---|---|---| -| `optim.lr` / `optim.weight_decay` | `5e-5` / `1e-5` | AdamW | -| `task.lr_warmup.*` | enabled, `1000` steps, `start_factor=0.0` | linear warmup, then cosine | -| `scheduler.T_max` | `auto` | `= total_steps - warmup` so the LR reaches `eta_min=1e-7` at the end | -| `task.total_steps` | `200000` | optimizer steps (`trainer.max_steps` follows it) | -| `task.momentum` / `task.momentum_final` | `0.99` / `1.0` | target-encoder EMA, cosine, applied after every optimizer step | -| `task.momentum_warmup_steps` | `0` | no momentum warmup (`null` falls back to `task.warmup_steps`) | -| `task.augmentation` | `mode=importance`, `mu=0.25`, `p_lambda=0.45`, `feature_noise_std=0.2`, `feature_noise_columns=[SIZE]` | symmetric views: degree-importance drop with the unassigned token + Gaussian noise on the named continuous feature(s); set `[]` without a size feature | -| `model.encoder.dropout` | `0.1` (in `model=bgrl_a3`) | encoder dropout during pretraining only | -| `callbacks.model_checkpoint.every_n_train_steps` | `5000` | all step checkpoints kept, plus `last.ckpt` | -| `data.split.train_val_test_split` | `[1,0,0]` | every region used label-free | - -## 4) Model Config Group (`configs/model/**/*.yaml`) - -## 4.1 `configs/model/supervised_module.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `model._target_` | `grass_mil.models.supervised_module.SupervisedModule` | str | import path | supervised runs | `grass-mil-train model=supervised_module` | -| `model._recursive_` | `false` | bool | `true/false` | Hydra instantiate semantics | `grass-mil-train model._recursive_=false` | -| `model.flags.use_ssl` | `false` | bool | `true/false` | optional SSL wrapper in supervised module | `grass-mil-train model.flags.use_ssl=true` | -| `model.encoder.input_dim` | `0` | int | `>0` or `0` for infer | encoder build | `grass-mil-train model.encoder.input_dim=128` | -| `model.encoder.hidden_dim` | `128` | int | `>=1` | encoder | `grass-mil-train model.encoder.hidden_dim=256` | -| `model.encoder.out_dim` | `128` | int/null | `>=1` or null | encoder projection | `grass-mil-train model.encoder.out_dim=256` | -| `model.encoder.num_layers` | `3` | int | `>=1` | encoder | `grass-mil-train model.encoder.num_layers=4` | -| `model.encoder.dropout` | `0.1` | float | `[0,1]` | encoder | `grass-mil-train model.encoder.dropout=0.2` | -| `model.encoder.conv_type` | `gin` | str | `gin`, `gcn`, `gat`, `graphsage`, `gine` | encoder | `grass-mil-train model.encoder.conv_type=gine` | -| `model.encoder.norm` | `batchnorm` | str | `batchnorm`, `layernorm`, `none` | encoder | `grass-mil-train model.encoder.norm=layernorm` | -| `model.encoder.jk` | `last` | str | `last`, `concat`, `max`, `sum` | encoder | `grass-mil-train model.encoder.jk=concat` | -| `model.encoder.act` | `relu` | str | `relu`, `gelu`, `leaky_relu` | encoder | `grass-mil-train model.encoder.act=gelu` | -| `model.encoder.pooling` | `mean` | str | `sum`,`mean`,`max` | graph embedding | `grass-mil-train model.encoder.pooling=max` | -| `model.encoder.gat_heads` | `4` | int | `>=1` and divides hidden dim | `conv_type=gat` | `grass-mil-train model.encoder.gat_heads=8` | -| `model.encoder.use_edge_attr` | `false` | bool | `true/false` | `gcn`/`gine` meaningful | `grass-mil-train model.encoder.use_edge_attr=true` | -| `model.encoder.edge_weight_index` | `0` | int | column index | `gcn` with edge attrs | `grass-mil-train model.encoder.edge_weight_index=1` | -| `model.encoder.edge_attr_dim` | `null` | int/null | required for `gine` | `conv_type=gine` | `grass-mil-train model.encoder.edge_attr_dim=2` | -| `model.encoder.edge_feature_index` | `null` | int/null | column of `edge_attr` fed to the convolution (`gine_length`: `0`) | `conv_type=gine` | `grass-mil-train model.encoder.edge_feature_index=0` | -| `model.encoder.train_eps` | `true` | bool | learnable GIN/GINE epsilon | `gin`/`gine` | `grass-mil-train model.encoder.train_eps=false` | -| `model.encoder.input_init` | `xavier` | str | `xavier`, `default` | input embedding + linear init | `grass-mil-train model.encoder.input_init=default` | -| `model.encoder.categorical_embedding.label` | `cell_type` | str | categorical label name | summed cell-type embedding | | -| `model.encoder.categorical_embedding.num_embeddings` / `column_index` | `null` | int/null | resolved from the processed dataset | | | -| `model.encoder.categorical_embedding.reserve_unassigned` | `true` | bool | reserve a trailing row as the unassigned token | augmentations | | -| `model.encoder.categorical_embedding.unassigned_label` / `unassigned_index` | `null` | str/int/null | reuse an existing category (e.g. `other`) as the unassigned token | | `grass-mil-train model.encoder.categorical_embedding.unassigned_label=other` | -| `model.attention.attention_type` | `gated_projected` | str | `gated`, `gated_projected` | `task.aggregation=mil_attention` | `grass-mil-train model.attention.attention_type=gated` | -| `model.attention.input_dim` | `128` | int | `>=1` | MIL attention | `grass-mil-train model.attention.input_dim=256` | -| `model.attention.projection_dim` | `32` | int | `>=1` | projected attention | `grass-mil-train model.attention.projection_dim=64` | -| `model.attention.hidden_dim` | `8` | int | `>=1` | MIL attention (the "attention space" width) | `grass-mil-train model.attention.hidden_dim=16` | -| `model.attention.dropout` | `true` | bool | `true/false` | dropout 0.25 in the attention net while training | `grass-mil-train model.attention.dropout=false` | -| `model.attention.init` | `kaiming` | str | `kaiming`, `default` | attention linears init | `grass-mil-train model.attention.init=default` | -| `model.attention.n_classes` | `1` (`2` with `task=finetune_mil`) | int | `>=1` | one attention channel per class | `grass-mil-train model.attention.n_classes=3` | -| `model.graph_head.input_dim` | `128` | int | `>=1` | graph head | `grass-mil-train model.graph_head.input_dim=256` | -| `model.graph_head.output_dim` | `1` | int | `>=1` | graph head | `grass-mil-train model.graph_head.output_dim=2` | -| `model.graph_head.hidden_dim` | `128` | int | `>=1` | graph head | `grass-mil-train model.graph_head.hidden_dim=256` | -| `model.graph_head.num_layers` | `2` | int | `>=1`; `1` = linear instance head (required by the probe regimes) | graph head | `grass-mil-train model.graph_head.num_layers=1` | -| `model.graph_head.dropout` | `0.0` | float | `[0,1]` | graph head | `grass-mil-train model.graph_head.dropout=0.2` | -| `model.ssl.method` | `bgrl` | str | `bgrl` | SSL factory | `grass-mil-train model.ssl.method=bgrl` | -| `model.ssl.predictor.hidden_size` | `512` | int | `>=1` | SSL enabled | `grass-mil-train model.ssl.predictor.hidden_size=256` | -| `model.loss.loss_type` | `${task.loss}` | str | loss name | loss factory | `grass-mil-train task.loss=regression_mse` | -| `model.optim` | `${optim}` | DictConfig | optimizer config object | always | `grass-mil-train optim.lr=5e-4` | -| `model.scheduler` | `${scheduler}` | DictConfig | scheduler config object | always | `grass-mil-train scheduler.T_max=200` | -| `model.task` | `${task}` | DictConfig | task config object | always | `grass-mil-train task=finetune_mil` | -| `model.init_from_ckpt` | `null` | str/null | checkpoint path | optional init | `grass-mil-train model.init_from_ckpt=/abs/pretrain.ckpt` | -| `model.init_strict` | `false` | bool | `true/false` | checkpoint init | `grass-mil-train model.init_strict=true` | -| `model.encoder_init_map` | `auto_bgrl_or_identity` | str | `identity`, `auto_bgrl_or_identity`, `spacegm_bgrl` (legacy SPACE-GM BGRL checkpoints) | checkpoint init | `grass-mil-train model.encoder_init_map=spacegm_bgrl` | -| `model.freeze_encoder` | `false` | bool | `true/false` | checkpoint transfer/finetuning | `grass-mil-train model.freeze_encoder=true` | - -## 4.2 `configs/model/bgrl_module.yaml` - -Same encoder block as supervised module plus SSL/pretrain init keys: - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `model._target_` | `grass_mil.models.bgrl_module.BGRLModule` | str | import path | pretrain task | `grass-mil-train model=bgrl_module` | -| `model._recursive_` | `false` | bool | `true/false` | instantiate behavior | `grass-mil-train model._recursive_=false` | -| `model.encoder.*` | same as supervised defaults | mixed | see section 4.1 | BGRL encoder | `grass-mil-train model.encoder.conv_type=gin` | -| `model.ssl.method` | `bgrl` | str | `bgrl` | SSL module | `grass-mil-train model.ssl.method=bgrl` | -| `model.ssl.predictor.hidden_size` | `512` | int | `>=1` | predictor width | `grass-mil-train model.ssl.predictor.hidden_size=1024` | -| `model.optim` | `${optim}` | DictConfig | optimizer object | always | `grass-mil-train optim.lr=1e-3` | -| `model.scheduler` | `${scheduler}` | DictConfig | scheduler object | always | `grass-mil-train scheduler.T_max=20000` | -| `model.task` | `${task}` | DictConfig | pretrain task object | always | `grass-mil-train task=pretrain_bgrl model=bgrl_module` | -| `model.init_from_ckpt` | `null` | str/null | checkpoint path | optional init | `grass-mil-train model.init_from_ckpt=/abs/ssl.ckpt` | -| `model.init_strict` | `false` | bool | `true/false` | checkpoint init | `grass-mil-train model.init_strict=true` | -| `model.encoder_init_map` | `auto_bgrl_or_identity` | str | `identity`, `auto_bgrl_or_identity` | checkpoint init | `grass-mil-train model.encoder_init_map=identity` | - -## 4.3 Model Presets And Component Groups - -### `configs/model/grass_mil.yaml` (GRASS-MIL supervised model) - -`defaults: [encoder: gine_length, attention: gated_projected]`, MLP instance head -(`graph_head.num_layers=2`, hidden = `encoder.out_dim`), bag logit pooled in logit -space (`L_c = sum_i A[i,c] f_c(h_i)`). Use `model=grass_mil task=finetune_mil` or a -`regime=` preset. `model.graph_head.num_layers=1` gives the linear head of the -probe regimes. - -### `configs/model/bgrl_a3.yaml` (GRASS-MIL pretraining model) - -`defaults: [encoder: gine_length]`, `encoder.dropout=0.1`, predictor hidden 512. - -### `configs/model/supervised_module.yaml`, `configs/model/bgrl_module.yaml` - -Generic modules: `defaults: [encoder: gin, attention: gated_projected]`. Switch -the encoder with `model/encoder=gine_length` and the attention with -`model/attention=gated`. - -Attention is **not** toggled by a model flag. `SupervisedModule` derives it from -the task: `use_attention = (task.aggregation == "mil_attention")`. Use -`task=finetune_mil` to enable it and `task=finetune_mean` to disable it. The -component sub-configs (`model.encoder`, `model.attention`, `model.graph_head`, -`model.ssl.predictor`) are plain option dicts consumed by -`build_supervised_components`, not `_target_` instantiations: select another -attention implementation with `model.attention.attention_type=gated`. - -`task=finetune_mil` sets the per-class widths itself -(`model.graph_head.output_dim=2`, `model.attention.n_classes=2`); `task` is merged -after `model` in the root defaults list (see `configs/train.yaml`). - -### Encoder group (`configs/model/encoder/*.yaml`) - -Files: `gin`, `gcn`, `gat`, `graphsage`, `gine`, `gine_length`. All expose the -keys of section 4.1; per-file defaults differ in `conv_type`, `norm`, -`use_edge_attr`, `edge_attr_dim`: - -- `gin.yaml`: `conv_type=gin`, `norm=batchnorm`, `use_edge_attr=false` -- `gcn.yaml`: `conv_type=gcn`, `norm=batchnorm`, `use_edge_attr=true` -- `gat.yaml`: `conv_type=gat`, `norm=batchnorm` -- `graphsage.yaml`: `conv_type=graphsage`, `norm=batchnorm` -- `gine.yaml`: `conv_type=gine`, `norm=batchnorm`, `use_edge_attr=true`, `edge_attr_dim=2` -- `gine_length.yaml` (GRASS-MIL): `conv_type=gine`, 4 x 256, `norm=layernorm`, - `pooling=max`, `edge_attr_dim=1`, `edge_feature_index=0` (the bounded log - edge length), `train_eps=true`, `input_init=xavier`, `categorical_embedding` - with a reserved unassigned row. Supports cell-type-only cohorts (`input_dim=0`). - -### Attention group (`configs/model/attention/*.yaml`) - -| File | Keys | Defaults | -|---|---|---| -| `gated.yaml` | `attention_type`, `input_dim`, `hidden_dim`, `dropout`, `n_classes`, `init` | `gated`, `${model.encoder.out_dim}`, `8`, `true`, `1`, `kaiming` | -| `gated_projected.yaml` | `attention_type`, `input_dim`, `projection_dim`, `hidden_dim`, `dropout`, `n_classes`, `init` | `gated_projected`, `${model.encoder.out_dim}`, `32`, `8`, `true`, `1`, `kaiming` | - -### Losses - -`task.loss` selects the loss: `categorical_bce` (`pos_weight`), `categorical_ce` -(`class_weight`, filled by `task.class_weights=auto_train`), `regression_mse`, -`regression_huber` (`delta`), `survival_coxsgd` (`top_n`, `regularizer_weight`). -Extra loss kwargs go under `model.loss.*`. - -### Augmentation preset files (`configs/augmentation/*.yaml`) - -Selects how BGRL builds its two views. Applies to `task=pretrain_bgrl`. - -Covered files: - -- `configs/augmentation/bgrl_paper.yaml` (production view generator) -- `configs/augmentation/uniform.yaml` (structure-independent baseline) - -**This group is not in any defaults list, so it must be attached explicitly.** -`bgrl_module` reads `task.augmentation`; attach a preset there with: - -```bash -grass-mil-train task=pretrain_bgrl model=bgrl_module \ - +augmentation@task.augmentation=bgrl_paper -``` - -The leading `+` and the `@task.augmentation` package are both required. A bare -`augmentation.mode=uniform` fails with `Could not override 'augmentation.mode'`, -because no `augmentation` key exists in the composed config until the preset is -attached. - -If `task.augmentation` is left unset, the module falls back to per-view uniform -drop rates read from `task.drop_edge_p1/p2` and `task.drop_feat_p1/p2` (see the -pretrain task table above). That fallback, not `bgrl_paper`, is what you get by -default. - -Once attached, individual keys are overridden under `task.augmentation.*`: - -| Key | Default (`bgrl_paper`) | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `task.augmentation.mode` | `importance` | str | `importance`, `uniform` | always | `grass-mil-train task.augmentation.mode=uniform` | -| `task.augmentation.mu` | `0.25` | float | `>0` | `mode=importance` | `grass-mil-train task.augmentation.mu=0.3` | -| `task.augmentation.p_lambda` | `0.45` | float | `(0, 1]` | `mode=importance` | `grass-mil-train task.augmentation.p_lambda=0.5` | -| `task.augmentation.feature_noise_std` | `0.2` | float | `>=0` | `mode=importance` | `grass-mil-train task.augmentation.feature_noise_std=0.1` | -| `task.augmentation.feature_noise_columns` | `["SIZE"]` | list[str] | column names | `mode=importance` | `grass-mil-train 'task.augmentation.feature_noise_columns=[AREA]'` | -| `task.augmentation.drop_edge_p` | `0.1` (`uniform`) | float | `[0, 1)` | `mode=uniform` | `grass-mil-train task.augmentation.drop_edge_p=0.2` | -| `task.augmentation.drop_feat_p` | `0.1` (`uniform`) | float | `[0, 1)` | `mode=uniform` | `grass-mil-train task.augmentation.drop_feat_p=0.2` | - -Two further constraints are easy to get wrong: - -- **Noise columns are resolved by name, never by index.** A hardcoded index that - means cell size in one cohort means something else in another, which silently - feeds the encoder pure noise. Names are matched against the dataset's - molecular feature names. -- **`mode=importance` needs the reserved embedding slot.** Dropped nodes have - their features zeroed and their cell type set to an "unassigned" code, so it - requires `model.encoder.categorical_embedding.reserve_unassigned=true`. - Cohorts with no continuous node features cannot use the Gaussian-noise step at - all; use `uniform` there. - -Node drop probability under `importance` is -`min((1 − importance) · mu, p_lambda)`, where importance is the mean-max -normalised log-degree — so hubs are preserved and low-degree nodes are masked. - -## 5) Optimizer And Scheduler Configs - -## 5.1 `configs/optim/adamw.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `optim._target_` | `torch.optim.AdamW` | str | optimizer class path | always | `grass-mil-train optim._target_=torch.optim.SGD` | -| `optim.lr` | `0.001` | float | `>0` | optimizer step | `grass-mil-train optim.lr=5e-4` | -| `optim.weight_decay` | `0.01` | float | `>=0` | optimizer step | `grass-mil-train optim.weight_decay=1e-4` | - -## 5.2 Scheduler files - -| File | Key | Default | Type | Example Override | -|---|---|---|---|---| -| `configs/scheduler/cosine_epoch.yaml` | `_target_` | `torch.optim.lr_scheduler.CosineAnnealingLR` | str | `grass-mil-train scheduler._target_=...` | -| `configs/scheduler/cosine_epoch.yaml` | `T_max` | `100` | int | `grass-mil-train scheduler.T_max=200` | -| `configs/scheduler/cosine_epoch.yaml` | `eta_min` | `1.0e-6` | float | `grass-mil-train scheduler.eta_min=1.0e-7` | -| `configs/scheduler/cosine_step.yaml` | `_target_` | `torch.optim.lr_scheduler.CosineAnnealingLR` | str | `grass-mil-train scheduler._target_=...` | -| `configs/scheduler/cosine_step.yaml` | `T_max` | `10000` | int | `grass-mil-train scheduler.T_max=20000` | -| `configs/scheduler/cosine_step.yaml` | `eta_min` | `1.0e-7` | float | `grass-mil-train scheduler.eta_min=1.0e-8` | -| `configs/scheduler/cosine_mil_step.yaml` | `T_max` | `auto` (`= max_epochs * ceil(n_train_units / hyperbatch)`, per optimizer step) | str/int | `grass-mil-train scheduler.T_max=5000` | -| `configs/scheduler/cosine_mil_step.yaml` | `eta_min` | `1.0e-7` | float | `grass-mil-train scheduler.eta_min=0` | - -## 6) Trainer Config Group (`configs/trainer/*.yaml`) - -## 6.1 `configs/trainer/default.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `trainer._target_` | `lightning.pytorch.trainer.Trainer` | str | trainer import path | always | `grass-mil-train trainer._target_=...` | -| `trainer.default_root_dir` | `${paths.output_dir}` | str | path | trainer init | `grass-mil-train trainer.default_root_dir=/tmp/out` | -| `trainer.min_epochs` | `1` | int | `>=0` | fit loop | `grass-mil-train trainer.min_epochs=0` | -| `trainer.max_epochs` | `10` | int | `>=1` | fit loop | `grass-mil-train trainer.max_epochs=100` | -| `trainer.accelerator` | `cpu` | str | lightning-supported accelerators | runtime | `grass-mil-train trainer.accelerator=gpu` | -| `trainer.devices` | `1` | int/list/str | lightning-supported device spec | runtime | `grass-mil-train trainer.devices=2` | -| `trainer.check_val_every_n_epoch` | `1` | int | `>=1` | fit loop | `grass-mil-train trainer.check_val_every_n_epoch=2` | -| `trainer.deterministic` | `False` | bool | `true/false` | reproducibility | `grass-mil-train trainer.deterministic=true` | - -## 6.2 Preset overrides - -| File | Keys In File | Meaning | -|---|---|---| -| `configs/trainer/cpu.yaml` | `defaults: [default]`, `accelerator=cpu`, `devices=1` | explicit CPU runtime | -| `configs/trainer/gpu.yaml` | `defaults: [default]`, `accelerator=gpu`, `devices=1` | single-GPU runtime | -| `configs/trainer/mps.yaml` | `defaults: [default]`, `accelerator=mps`, `devices=1` | Apple MPS runtime | -| `configs/trainer/ddp.yaml` | `defaults: [default]`, `strategy=ddp`, `accelerator=gpu`, `devices=4`, `num_nodes=1`, `sync_batchnorm=True` | multi-GPU distributed preset | -| `configs/trainer/ddp_sim.yaml` | `defaults: [default]`, `accelerator=cpu`, `devices=2`, `strategy=ddp_spawn` | CPU DDP simulation for debugging | - -## 7) Callback Config Group (`configs/callbacks/*.yaml`) - -## 7.1 `configs/callbacks/default.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `defaults` | `[model_checkpoint, early_stopping, model_summary, rich_progress_bar, _self_]` | list | callback group names | always | `grass-mil-train callbacks=none` | -| `callbacks.model_checkpoint.dirpath` | `${paths.output_dir}/checkpoints` | str | path | checkpoint callback enabled | `grass-mil-train callbacks.model_checkpoint.dirpath=/tmp/ckpt` | -| `callbacks.model_checkpoint.filename` | `epoch_{epoch:03d}` | str | format string | checkpoint callback enabled | `grass-mil-train callbacks.model_checkpoint.filename=best` | -| `callbacks.model_checkpoint.monitor` | `val/loss` | str | logged metric name | checkpoint callback enabled | `grass-mil-train callbacks.model_checkpoint.monitor=val/acc` | -| `callbacks.model_checkpoint.mode` | `min` | str | `min`, `max` | checkpoint callback enabled | `grass-mil-train callbacks.model_checkpoint.mode=max` | -| `callbacks.model_checkpoint.save_last` | `True` | bool | `true/false` | checkpoint callback enabled | `grass-mil-train callbacks.model_checkpoint.save_last=false` | -| `callbacks.model_checkpoint.auto_insert_metric_name` | `False` | bool | `true/false` | checkpoint callback enabled | `grass-mil-train callbacks.model_checkpoint.auto_insert_metric_name=true` | -| `callbacks.early_stopping.monitor` | `val/loss` | str | logged metric name | early stopping enabled | `grass-mil-train callbacks.early_stopping.monitor=val/acc` | -| `callbacks.early_stopping.patience` | `100` | int | `>=0` | early stopping enabled | `grass-mil-train callbacks.early_stopping.patience=20` | -| `callbacks.early_stopping.mode` | `min` | str | `min`, `max` | early stopping enabled | `grass-mil-train callbacks.early_stopping.mode=max` | -| `callbacks.model_summary.max_depth` | `-1` | int | integer depth | model summary enabled | `grass-mil-train callbacks.model_summary.max_depth=2` | - -## 7.2 `configs/callbacks/model_checkpoint.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `model_checkpoint._target_` | `lightning.pytorch.callbacks.ModelCheckpoint` | str | import path | always | `grass-mil-train callbacks.model_checkpoint._target_=...` | -| `model_checkpoint.dirpath` | `null` | str/null | path or null | callback active | `grass-mil-train callbacks.model_checkpoint.dirpath=/tmp/ckpt` | -| `model_checkpoint.filename` | `null` | str/null | filename template or null | callback active | `grass-mil-train callbacks.model_checkpoint.filename=epoch_{epoch}` | -| `model_checkpoint.monitor` | `null` | str/null | metric name or null | callback active | `grass-mil-train callbacks.model_checkpoint.monitor=val/loss` | -| `model_checkpoint.verbose` | `False` | bool | `true/false` | callback active | `grass-mil-train callbacks.model_checkpoint.verbose=true` | -| `model_checkpoint.save_last` | `null` | bool/null | bool or null | callback active | `grass-mil-train callbacks.model_checkpoint.save_last=true` | -| `model_checkpoint.save_top_k` | `1` | int | integer | callback active | `grass-mil-train callbacks.model_checkpoint.save_top_k=3` | -| `model_checkpoint.mode` | `min` | str | `min`, `max` | callback active | `grass-mil-train callbacks.model_checkpoint.mode=max` | -| `model_checkpoint.auto_insert_metric_name` | `True` | bool | `true/false` | callback active | `grass-mil-train callbacks.model_checkpoint.auto_insert_metric_name=false` | -| `model_checkpoint.save_weights_only` | `False` | bool | `true/false` | callback active | `grass-mil-train callbacks.model_checkpoint.save_weights_only=true` | -| `model_checkpoint.every_n_train_steps` | `null` | int/null | integer or null | callback active | `grass-mil-train callbacks.model_checkpoint.every_n_train_steps=100` | -| `model_checkpoint.train_time_interval` | `null` | duration/null | Lightning duration format or null | callback active | `grass-mil-train callbacks.model_checkpoint.train_time_interval='00:10:00'` | -| `model_checkpoint.every_n_epochs` | `null` | int/null | integer or null | callback active | `grass-mil-train callbacks.model_checkpoint.every_n_epochs=1` | -| `model_checkpoint.save_on_train_epoch_end` | `null` | bool/null | bool or null | callback active | `grass-mil-train callbacks.model_checkpoint.save_on_train_epoch_end=true` | - -## 7.3 `configs/callbacks/early_stopping.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `early_stopping._target_` | `lightning.pytorch.callbacks.EarlyStopping` | str | import path | always | `grass-mil-train callbacks.early_stopping._target_=...` | -| `early_stopping.monitor` | `???` | str | metric name | early stopping enabled | `grass-mil-train callbacks.early_stopping.monitor=val/loss` | -| `early_stopping.min_delta` | `0.` | float | any float | early stopping enabled | `grass-mil-train callbacks.early_stopping.min_delta=0.001` | -| `early_stopping.patience` | `3` | int | `>=0` | early stopping enabled | `grass-mil-train callbacks.early_stopping.patience=10` | -| `early_stopping.verbose` | `False` | bool | `true/false` | early stopping enabled | `grass-mil-train callbacks.early_stopping.verbose=true` | -| `early_stopping.mode` | `min` | str | `min`, `max` | early stopping enabled | `grass-mil-train callbacks.early_stopping.mode=max` | -| `early_stopping.strict` | `True` | bool | `true/false` | early stopping enabled | `grass-mil-train callbacks.early_stopping.strict=false` | -| `early_stopping.check_finite` | `True` | bool | `true/false` | early stopping enabled | `grass-mil-train callbacks.early_stopping.check_finite=true` | -| `early_stopping.stopping_threshold` | `null` | float/null | float or null | early stopping enabled | `grass-mil-train callbacks.early_stopping.stopping_threshold=0.1` | -| `early_stopping.divergence_threshold` | `null` | float/null | float or null | early stopping enabled | `grass-mil-train callbacks.early_stopping.divergence_threshold=10.0` | -| `early_stopping.check_on_train_epoch_end` | `null` | bool/null | bool or null | early stopping enabled | `grass-mil-train callbacks.early_stopping.check_on_train_epoch_end=true` | - -## 7.4 Other callback files - -Covered files: - -- `configs/callbacks/model_summary.yaml` -- `configs/callbacks/rich_progress_bar.yaml` -- `configs/callbacks/none.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `model_summary._target_` | `lightning.pytorch.callbacks.RichModelSummary` | str | import path | `callbacks=model_summary` | `grass-mil-train callbacks.model_summary._target_=...` | -| `model_summary.max_depth` | `1` | int | integer depth | `callbacks=model_summary` | `grass-mil-train callbacks.model_summary.max_depth=2` | -| `rich_progress_bar._target_` | `lightning.pytorch.callbacks.RichProgressBar` | str | import path | `callbacks=rich_progress_bar` | `grass-mil-train callbacks.rich_progress_bar._target_=...` | -| `configs/callbacks/none.yaml` | empty file | n/a | no fields | `callbacks=none` | `grass-mil-train callbacks=none` | - -## 7.5 Selection presets - -| File | Content | -|---|---| -| `configs/callbacks/mil_selection.yaml` | checkpoint on `val/auroc` (max, top 1 + last), early stopping on `val/loss_region` (patience 8, after `trainer.min_epochs`) | -| `configs/callbacks/refit.yaml` | fixed-length refit: no early stopping, no metric selection, `last.ckpt` only | -| `configs/callbacks/pretrain_checkpoint.yaml` | checkpoint every `every_n_train_steps=5000`, all kept | - -## 8) Logger Config Group (`configs/logger/*.yaml`) - -## 8.1 `configs/logger/csv.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `csv._target_` | `lightning.pytorch.loggers.csv_logs.CSVLogger` | str | import path | `logger=csv` | `grass-mil-train logger.csv._target_=...` | -| `csv.save_dir` | `${paths.output_dir}` | str | path | `logger=csv` | `grass-mil-train logger.csv.save_dir=/tmp/logs` | -| `csv.name` | `csv/` | str | name | `logger=csv` | `grass-mil-train logger.csv.name=metrics` | -| `csv.prefix` | `""` | str | prefix string | `logger=csv` | `grass-mil-train logger.csv.prefix=exp1` | - -## 8.2 `configs/logger/tensorboard.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `tensorboard._target_` | `lightning.pytorch.loggers.tensorboard.TensorBoardLogger` | str | import path | `logger=tensorboard` | `grass-mil-train logger.tensorboard._target_=...` | -| `tensorboard.save_dir` | `${paths.output_dir}/tensorboard/` | str | path | `logger=tensorboard` | `grass-mil-train logger.tensorboard.save_dir=/tmp/tb` | -| `tensorboard.name` | `null` | str/null | name or null | `logger=tensorboard` | `grass-mil-train logger.tensorboard.name=runA` | -| `tensorboard.log_graph` | `False` | bool | `true/false` | `logger=tensorboard` | `grass-mil-train logger.tensorboard.log_graph=true` | -| `tensorboard.default_hp_metric` | `True` | bool | `true/false` | `logger=tensorboard` | `grass-mil-train logger.tensorboard.default_hp_metric=false` | -| `tensorboard.prefix` | `""` | str | prefix string | `logger=tensorboard` | `grass-mil-train logger.tensorboard.prefix=train` | - -## 8.3 `configs/logger/wandb.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `wandb._target_` | `lightning.pytorch.loggers.wandb.WandbLogger` | str | import path | `logger=wandb` | `grass-mil-train logger.wandb._target_=...` | -| `wandb.save_dir` | `${paths.output_dir}` | str | path | `logger=wandb` | `grass-mil-train logger.wandb.save_dir=/tmp/wb` | -| `wandb.offline` | `False` | bool | `true/false` | `logger=wandb` | `grass-mil-train logger.wandb.offline=true` | -| `wandb.id` | `null` | str/null | run ID or null | `logger=wandb` | `grass-mil-train logger.wandb.id=my_id` | -| `wandb.anonymous` | `null` | str/null | WandB anonymous mode or null | `logger=wandb` | `grass-mil-train logger.wandb.anonymous=allow` | -| `wandb.project` | `grass-mil` | str | project name | `logger=wandb` | `grass-mil-train logger.wandb.project=grass-mil` | -| `wandb.log_model` | `False` | bool | `true/false` | `logger=wandb` | `grass-mil-train logger.wandb.log_model=true` | -| `wandb.prefix` | `""` | str | prefix string | `logger=wandb` | `grass-mil-train logger.wandb.prefix=exp` | -| `wandb.group` | `""` | str | group name | `logger=wandb` | `grass-mil-train logger.wandb.group=ablation` | -| `wandb.tags` | `[]` | list[str] | any list | `logger=wandb` | `grass-mil-train logger.wandb.tags='[ssl,mil]'` | -| `wandb.job_type` | `""` | str | job type | `logger=wandb` | `grass-mil-train logger.wandb.job_type=train` | - -## 8.4 `configs/logger/mlflow.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `mlflow._target_` | `lightning.pytorch.loggers.mlflow.MLFlowLogger` | str | import path | `logger=mlflow` | `grass-mil-train logger.mlflow._target_=...` | -| `mlflow.tracking_uri` | `${paths.log_dir}/mlflow/mlruns` | str | URI/path | `logger=mlflow` | `grass-mil-train logger.mlflow.tracking_uri=/tmp/mlruns` | -| `mlflow.tags` | `null` | dict/null | tags dict or null | `logger=mlflow` | `grass-mil-train logger.mlflow.tags='{stage:dev}'` | -| `mlflow.prefix` | `""` | str | prefix | `logger=mlflow` | `grass-mil-train logger.mlflow.prefix=exp` | -| `mlflow.artifact_location` | `null` | str/null | URI/path or null | `logger=mlflow` | `grass-mil-train logger.mlflow.artifact_location=/tmp/artifacts` | - -## 8.5 `configs/logger/comet.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `comet._target_` | `lightning.pytorch.loggers.comet.CometLogger` | str | import path | `logger=comet` | `grass-mil-train logger.comet._target_=...` | -| `comet.api_key` | `${oc.env:COMET_API_TOKEN}` | str | API key/env ref | `logger=comet` | `grass-mil-train logger.comet.api_key=...` | -| `comet.save_dir` | `${paths.output_dir}` | str | path | `logger=comet` | `grass-mil-train logger.comet.save_dir=/tmp/comet` | -| `comet.project_name` | `grass-mil` | str | project name | `logger=comet` | `grass-mil-train logger.comet.project_name=grass-mil` | -| `comet.rest_api_key` | `null` | str/null | key or null | `logger=comet` | `grass-mil-train logger.comet.rest_api_key=...` | -| `comet.experiment_key` | `null` | str/null | key or null | `logger=comet` | `grass-mil-train logger.comet.experiment_key=...` | -| `comet.offline` | `False` | bool | `true/false` | `logger=comet` | `grass-mil-train logger.comet.offline=true` | -| `comet.prefix` | `""` | str | prefix | `logger=comet` | `grass-mil-train logger.comet.prefix=train` | - -## 8.6 `configs/logger/neptune.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `neptune._target_` | `lightning.pytorch.loggers.neptune.NeptuneLogger` | str | import path | `logger=neptune` | `grass-mil-train logger.neptune._target_=...` | -| `neptune.api_key` | `${oc.env:NEPTUNE_API_TOKEN}` | str | API key/env ref | `logger=neptune` | `grass-mil-train logger.neptune.api_key=...` | -| `neptune.project` | `org/grass-mil` | str | project slug | `logger=neptune` | `grass-mil-train logger.neptune.project=org/project` | -| `neptune.log_model_checkpoints` | `True` | bool | `true/false` | `logger=neptune` | `grass-mil-train logger.neptune.log_model_checkpoints=false` | -| `neptune.prefix` | `""` | str | prefix | `logger=neptune` | `grass-mil-train logger.neptune.prefix=exp` | - -## 8.7 `configs/logger/aim.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `aim._target_` | `aim.pytorch_lightning.AimLogger` | str | import path | `logger=aim` | `grass-mil-train logger.aim._target_=...` | -| `aim.repo` | `${paths.root_dir}` | str | path/remote URI | `logger=aim` | `grass-mil-train logger.aim.repo=/tmp/.aim` | -| `aim.experiment` | `null` | str/null | experiment name or null | `logger=aim` | `grass-mil-train logger.aim.experiment=baseline` | -| `aim.train_metric_prefix` | `train/` | str | prefix | `logger=aim` | `grass-mil-train logger.aim.train_metric_prefix=train` | -| `aim.val_metric_prefix` | `val/` | str | prefix | `logger=aim` | `grass-mil-train logger.aim.val_metric_prefix=val` | -| `aim.test_metric_prefix` | `test/` | str | prefix | `logger=aim` | `grass-mil-train logger.aim.test_metric_prefix=test` | -| `aim.system_tracking_interval` | `10` | int/null | seconds or null | `logger=aim` | `grass-mil-train logger.aim.system_tracking_interval=5` | -| `aim.log_system_params` | `true` | bool | `true/false` | `logger=aim` | `grass-mil-train logger.aim.log_system_params=false` | -| `aim.capture_terminal_logs` | `false` | bool | `true/false` | `logger=aim` | `grass-mil-train logger.aim.capture_terminal_logs=true` | - -## 8.8 `configs/logger/many_loggers.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `defaults` | `[csv, tensorboard, wandb]` | list | logger group names | `logger=many_loggers` | `grass-mil-train logger=many_loggers logger.defaults='[csv,tensorboard]'` | - -## 9) Debug Config Group (`configs/debug/*.yaml`) - -## 9.1 `configs/debug/default.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `task_name` | `debug` | str | any string | `debug=default` | `grass-mil-train debug=default task_name=dbg` | -| `callbacks` | `null` | null/config | null or callback config | `debug=default` | `grass-mil-train debug=default callbacks=default` | -| `logger` | `null` | null/config | null or logger config | `debug=default` | `grass-mil-train debug=default logger=csv` | -| `extras.ignore_warnings` | `False` | bool | `true/false` | `debug=default` | `grass-mil-train debug=default extras.ignore_warnings=true` | -| `extras.enforce_tags` | `False` | bool | `true/false` | `debug=default` | `grass-mil-train debug=default extras.enforce_tags=false` | -| `hydra.job_logging.root.level` | `DEBUG` | str | logging level | `debug=default` | `grass-mil-train debug=default hydra.job_logging.root.level=INFO` | -| `trainer.max_epochs` | `1` | int | `>=1` | `debug=default` | `grass-mil-train debug=default trainer.max_epochs=2` | -| `trainer.accelerator` | `cpu` | str | accelerator value | `debug=default` | `grass-mil-train debug=default trainer.accelerator=cpu` | -| `trainer.devices` | `1` | int | `>=1` | `debug=default` | `grass-mil-train debug=default trainer.devices=1` | -| `trainer.detect_anomaly` | `true` | bool | `true/false` | `debug=default` | `grass-mil-train debug=default trainer.detect_anomaly=true` | -| `data.num_workers` | `0` | int | `>=0` | `debug=default` | `grass-mil-train debug=default data.num_workers=0` | -| `data.pin_memory` | `False` | bool | `true/false` | `debug=default` | `grass-mil-train debug=default data.pin_memory=false` | - -## 9.2 Other debug presets - -Covered files: - -- `configs/debug/fdr.yaml` -- `configs/debug/limit.yaml` -- `configs/debug/overfit.yaml` -- `configs/debug/profiler.yaml` - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `debug/fdr.defaults` | `[default]` | list | debug group list | `debug=fdr` | `grass-mil-train debug=fdr` | -| `debug/fdr.trainer.fast_dev_run` | `true` | bool | `true/false` | `debug=fdr` | `grass-mil-train debug=fdr trainer.fast_dev_run=true` | -| `debug/limit.defaults` | `[default]` | list | debug group list | `debug=limit` | `grass-mil-train debug=limit` | -| `debug/limit.trainer.max_epochs` | `3` | int | `>=1` | `debug=limit` | `grass-mil-train debug=limit trainer.max_epochs=5` | -| `debug/limit.trainer.limit_train_batches` | `0.01` | float | `(0,1]` or int | `debug=limit` | `grass-mil-train debug=limit trainer.limit_train_batches=0.1` | -| `debug/limit.trainer.limit_val_batches` | `0.05` | float | `(0,1]` or int | `debug=limit` | `grass-mil-train debug=limit trainer.limit_val_batches=0.2` | -| `debug/limit.trainer.limit_test_batches` | `0.05` | float | `(0,1]` or int | `debug=limit` | `grass-mil-train debug=limit trainer.limit_test_batches=0.2` | -| `debug/overfit.defaults` | `[default]` | list | debug group list | `debug=overfit` | `grass-mil-train debug=overfit` | -| `debug/overfit.trainer.max_epochs` | `20` | int | `>=1` | `debug=overfit` | `grass-mil-train debug=overfit trainer.max_epochs=10` | -| `debug/overfit.trainer.overfit_batches` | `3` | int | `>=1` | `debug=overfit` | `grass-mil-train debug=overfit trainer.overfit_batches=1` | -| `debug/overfit.callbacks` | `null` | null/config | null or callback config | `debug=overfit` | `grass-mil-train debug=overfit callbacks=none` | -| `debug/profiler.defaults` | `[default]` | list | debug group list | `debug=profiler` | `grass-mil-train debug=profiler` | -| `debug/profiler.trainer.max_epochs` | `1` | int | `>=1` | `debug=profiler` | `grass-mil-train debug=profiler trainer.max_epochs=2` | -| `debug/profiler.trainer.profiler` | `simple` | str | `simple`, `advanced`, `pytorch`, etc. | `debug=profiler` | `grass-mil-train debug=profiler trainer.profiler=advanced` | - -## 10) Hydra Runtime Config (`configs/hydra/default.yaml`) - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `defaults` | overrides for `hydra_logging=colorlog`, `job_logging=colorlog` | list | hydra logging groups | always | `grass-mil-train hydra/job_logging=stdout` | -| `hydra.run.dir` | `${paths.log_dir}/${task_name}/runs/${now:%Y-%m-%d}_${now:%H-%M-%S}` | str | directory template | single run | `grass-mil-train hydra.run.dir=/tmp/single_run` | -| `hydra.sweep.dir` | `${paths.log_dir}/${task_name}/multiruns/${now:%Y-%m-%d}_${now:%H-%M-%S}` | str | directory template | multirun | `grass-mil-train -m hydra.sweep.dir=/tmp/multi` | -| `hydra.sweep.subdir` | `${hydra.job.num}` | str | template string | multirun | `grass-mil-train -m hydra.sweep.subdir=job_${hydra.job.num}` | -| `hydra.job_logging.handlers.file.filename` | `${hydra.runtime.output_dir}/${task_name}.log` | str | file path template | Hydra file logging | `grass-mil-train hydra.job_logging.handlers.file.filename=/tmp/train.log` | - -## 11) Paths Config (`configs/paths/default.yaml`) - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `paths.root_dir` | `${oc.env:PROJECT_ROOT}` | str | existing path | always | `grass-mil-train paths.root_dir=/abs/project` | -| `paths.data_dir` | `${paths.root_dir}/data/` | str | path | always | `grass-mil-train paths.data_dir=/abs/data` | -| `paths.log_dir` | `${paths.root_dir}/logs/` | str | path | always | `grass-mil-train paths.log_dir=/abs/logs` | -| `paths.output_dir` | `${hydra:runtime.output_dir}` | str | hydra runtime path | always | `grass-mil-train paths.output_dir=/tmp/out` | -| `paths.work_dir` | `${hydra:runtime.cwd}` | str | working directory path | always | `grass-mil-train paths.work_dir=/tmp` | - -## 12) Extras Config (`configs/extras/default.yaml`) - -| Key | Default | Type | Valid Values | Active When | Example Override | -|---|---|---|---|---|---| -| `extras.ignore_warnings` | `False` | bool | `true/false` | `extras` enabled | `grass-mil-train extras.ignore_warnings=true` | -| `extras.enforce_tags` | `True` | bool | `true/false` | `extras` enabled | `grass-mil-train extras.enforce_tags=false` | -| `extras.print_config` | `True` | bool | `true/false` | `extras` enabled | `grass-mil-train extras.print_config=false` | - -## 12a) Regime Presets (`configs/regime/*.yaml`) - -`regime=` composes `model=grass_mil`, `task=finetune_mil`, -`scheduler=cosine_mil_step`, `callbacks=mil_selection`, `data/sampler=ego_graph_mil` -and sets the keys below; each is an ordinary override afterwards. - -| Regime | `optim.lr` / `weight_decay` | `encoder_lr_scale` | head | `head_lr_scale` / `head_weight_decay` | `attention_lr_scale` | roots per region | `encoder.dropout` | drop edge / feat | epochs (min) | -|---|---|---|---|---|---|---|---|---|---| -| `mil` | 1e-4 / 1e-3 | 1.0 | MLP (2 layers) | 1.0 / null | 2.0 | 256 | 0.0 | 0 / 0 | 25 (10) | -| `mil_frozen` | 1e-4 / 1e-3 | 0.0 | MLP | 1.0 / null | 2.0 | 512 | 0.0 | 0 / 0 | 25 (10) | -| `linear_probehead` | 1e-4 / 1e-3 | 0.0 | linear, probe pre-fit, frozen | 0.0 / null | 2.0 | 512 | 0.0 | 0 / 0 | 25 (10) | -| `linear_ftall` | 1e-4 / 1e-3 | 1.0 | linear, probe pre-fit | 0.01 / 1e-3 | 2.0 | 256 | 0.0 | 0 / 0 | 25 (10) | -| `linear_ftall_reg` | 1e-4 / 1e-2 | 0.3 | linear, probe pre-fit | 0.01 / 1e-3 | 2.0 | 256 | 0.2 | 0.05 / 0.05 | 25 (10) | - -`regime=pretrain_a3` is the pretraining recipe (section 3.4): it swaps the model, task, callbacks and sampler groups and needs no labels. - -Shared by the fine-tuning regimes: `task.region_accumulation.hyperbatch_size=12`, `task.grad_clip_norm=1.0`, -`trainer.precision=16-mixed`, `data.val_sampler` = fixed 1024-root subset per -region (`fixed_per_unit_seed=0`), `data.test_sampler` = all roots. - -## 12b) Analysis CLIs - -| CLI | Config | Main keys | -|---|---|---| -| `grass-mil-probe` | `configs/probe.yaml` | `probe.roots_per_unit=96`, `probe.seed=0`, `probe.batch_size=256`, `probe.C=1.0`, `probe.class_weight=balanced`, `probe.max_iter=3000`, `probe.eval_splits=[val]`, `probe.output_path`; splits from `data.split.explicit` | -| `grass-mil-nested-cv` | `configs/nested_cv.yaml` | `nested_cv.folds_file`, `nested_cv.id_kind=region|sample|patch`, `nested_cv.{label_file,id_column,label_column,patient_column}` (default from `data.graph_labels`), `nested_cv.inner_k=3`, `inner_seed=0`, `selection_metric=val/auroc`, `epoch_reduce=median`, `min_refit_epochs=1`, `outer_folds`, `resume`, `nested_cv.probe.*` (as the probe keys, `enabled=auto`, `cache_dir`), `patient_aggregation=mean` | -| `grass-mil-extract` | `configs/extract.yaml` | `extract.nested_cv_summary` or `+extract.models.=` (+ `+extract.held_out.=[...]`), `extract.unit_id=auto|region|sample|patch`, `extract.regions`, `extract.batch_size=512`, `extract.autocast`, `extract.device`, `extract.resume`, `extract.output_dir` | -| `grass-mil-scores` | `configs/scores.yaml` | `input`, `label_column`, `scores`, `fold_column`, `patient_column`, `patient_level`, `standardise=z|rank|none`, `strata`, `n_boot=2000`, `seed=0`, `operating_rules`, `prevalences` | -| `grass-mil-niches` | `configs/niches.yaml` | `cells_dir`, `regions=held_out|all|list|file`, `labels.*`, `space=gnn|attention`, `attention_ckpt`, `reduction.{method,n_components}`, `clustering.{method,n_clusters,params.*}`, `sample_size`, `l2_normalise`, `assign_k`, `seed`, `present_threshold`, `n_boot`, `figures`, `umap.*`, `positive_class`, `positive_name` | -| `grass-mil-metaclusters` | `configs/metaclusters.yaml` | `cells_dir`, `models`, `stage=all|cluster|match|transfer|basis`, `labels.*`, the clustering keys of `niches.yaml`, `match.{distance,threshold,scan.{min,max,n},min_folds}`, `transfer.{min_cells,present_threshold}`, `basis.{enabled,space_model,k,sample_size,n_init,seed}` | -| `grass-mil-spatial-stats` | `configs/spatial_stats.yaml` | `cells_dir`, `regions`, `labels.*`, `analyses`, `group_map`, `reference`, `partners`, `infiltrating`, `nhood.{n_perms,backend}`, `cn.{window_k,n_cn,sample_size}`, `crossk.radii`, `seed`, `min_per_group` | -| `grass-mil-predict` | `configs/inference/predict.yaml` | `aggregation.mode=mean|max|none|attention_weighted|mean_probability`, `aggregation.bag_scope=patch|region|sample|patient`, `aggregation.patient_map.{label_file,id_column,patient_column}` | - -## 13) Repo-Enforced Behavioral Constraints (From Code) - -These constraints are not just documentation conventions; they are enforced in runtime code. - -| Constraint | Enforced In | Behavior | -|---|---|---| -| `task.target_type` must be one of `binary`, `regression`, `survival` | `src/grass_mil/models/training/builders.py` | raises `ValueError` for unsupported target types | -| `task.instance_sampling` must be `all` or `random` | `src/grass_mil/models/training/builders.py` | raises `ValueError` for invalid sampling mode | -| `task.region_accumulation.hyperbatch_size >= 1` when enabled | `src/grass_mil/models/supervised_module.py` | raises `ValueError` | -| `conv_type=gine` requires `model.encoder.edge_attr_dim` and runtime `edge_attr` | `src/grass_mil/models/components/backbones.py` | raises `ValueError` | -| `conv_type=gat` requires `hidden_dim % gat_heads == 0` | `src/grass_mil/models/components/backbones.py` | raises `ValueError` | -| `model.encoder.use_edge_attr=true` is meaningful only for `gcn`/`gine` | `src/grass_mil/models/components/backbones.py` | warning/ignored for unsupported convs | -| `data.split.train_val_test_split` must sum to `1.0` | `src/grass_mil/data/components/precompute.py` split helpers | raises `ValueError` | -| `data.feature_reducer.fit_mode` must be `global` or `train_only` | `src/grass_mil/data/components/precompute.py` | raises `ValueError` | -| `data.sampler.runtime.weight_mode` must be `inverse`, `sqrt_inverse`, or `proportional` | `src/grass_mil/data/components/samplers.py` | raises `ValueError` | -| `model.encoder_init_map` must be `identity` or `auto_bgrl_or_identity` | `src/grass_mil/models/training/checkpoint_init.py` | raises `ValueError` | -| `model.init_from_ckpt` path must exist | checkpoint init helper | raises `FileNotFoundError` | - -## 14) Trainer/Callback/Logger Exhaustiveness Note - -This file is exhaustive for keys defined in repository config presets. - -Hydra still allows passing additional Lightning kwargs that are not pre-declared in these YAML files. Those extra keys are valid if accepted by the target Lightning class. - -Example: - -```bash -grass-mil-train trainer.log_every_n_steps=1 trainer.enable_progress_bar=true -``` - -Use this flexibility with caution and keep custom overrides version-controlled in experiment config files when possible. diff --git a/docs/inference_utility_suite.md b/docs/inference_utility_suite.md deleted file mode 100644 index f11fc62..0000000 --- a/docs/inference_utility_suite.md +++ /dev/null @@ -1,161 +0,0 @@ -# Inference Utility Suite - -This document describes the artifact-oriented inference workflow implemented in -`src/grass_mil/inference/predict.py`. - -## Purpose - -The suite extends checkpoint evaluation with: - -- optionally deterministic prediction export -- task-aware metrics for categorical/regression/survival tasks -- graph embedding export for downstream interpretability workflows -- optional node embedding export with bag-id mapping -- single-pass collection for predictions + embeddings (no second dataloader pass) - -## CLI Entry Point - -```bash -grass-mil-predict \ - ckpt_path=/absolute/path/to/checkpoint.ckpt \ - task=finetune_mil \ - model=supervised_module \ - data=spatial_omics -``` - -Base config: - -- `configs/inference/predict.yaml` - -Sub-config groups: - -- `configs/inference/metrics/default.yaml` -- `configs/inference/aggregation/default.yaml` -- `configs/inference/embeddings/default.yaml` -- `configs/inference/interpretability/default.yaml` - -## Aggregation Subsampling Semantics - -`aggregation.subsample_fraction` supports efficient inference-time subsampling: - -- when `< 1.0`, `src/grass_mil/inference/predict.py` injects predict-time sampler runtime - subsampling so fewer roots/subgraphs are forwarded through the model -- sampling is uniform at random over candidate roots, with optional reproducibility - via `aggregation.subsample_seed` -- this path applies to aggregation modes `mean`, `max`, and `attention_weighted` - while avoiding post-aggregation double subsampling -- if subsampling is requested but cannot be effectively applied preforward - (for example, `identity` sampler or `runtime.enabled=false`), inference fails - fast with a remediation error -- `aggregation.subsample_fraction=1.0` disables subsampling and keeps full usage -- `aggregation.subsample_fraction=0.0` is invalid - -Inference summary metadata includes: - -- `aggregation.preforward_subsampling_applied` (backward-compatible bool) -- `aggregation.preforward_subsampling` (structured decision payload) - -## Output Schema - -Outputs are written under: - -- `${paths.output_dir}/${predict.output_subdir}` - -### `predictions.csv` - -Columns: - -- `bag_id` -- `logit_*` -- optional `target_*` -- optional `attention` (JSON-encoded list) - -Rows are sorted by `bag_id` for deterministic export. - -### `metrics.json` - -Structure: - -- `global`: overall metrics -- `per_group`: per-bag metrics keyed by `bag_id` (when `metrics.per_group=true`) - -Task metric sets: - -- binary/categorical: `accuracy`, `precision`, `recall`, `f1`, `roc_auc` -- regression: `r2`, `mae`, `rmse` -- survival: `c_index` - -### `graph_embeddings.csv` - -Columns: - -- `bag_id` -- `graph_emb_*` - -### `node_embeddings.csv` (optional) - -Enabled via `embeddings.extract_node=true`. - -Columns: - -- `bag_id` -- `node_emb_*` - -### `inference_summary.json` - -Quick run summary with artifact paths and row count. - -### `instance_table.csv` (optional) - -Enabled via `interpretability.enabled=true`. - -Columns include: - -- `instance_id`, `bag_id`, `patch_id`, `region_id`, `sample_id` -- `score`, `logit_*`, `attention` -- `inst_emb_*` -- `comp_*` (when composition metadata is available) -- `center_x`, `center_y` (when centroid metadata is available) - -Notes: - -- `attention` is normalized within each bag. -- when attention logits are unavailable, `attention` defaults to uniform within bag. -- `score` export is configurable via `interpretability.score_mode` and - `interpretability.score_logit_index`: - - `score_mode=sigmoid` (default): `sigmoid(logit[index])` - - `score_mode=identity`: raw `logit[index]` (for hazard/log-risk style outputs) - - `score_mode=softmax`: class probability `softmax(logits)[index]` -- by default, missing composition metadata raises (`interpretability.require_composition=true`). - -### `spatial_table.csv` (optional) - -Enabled via `interpretability.spatial.enabled=true`. - -Columns: - -- `source_id` -- `target_id` -- `distance` - -Edges are taken from original subgraph/root connectivity by reading `instance_graphs` -stored alongside prediction payloads. -No kNN graph is reconstructed during interpretability export. - -Unit contract: - -- Preprocessing converts coordinates/distances to micrometers using `data.coord_scale_um`. -- `data.coord_scale_um=1.0` means input coordinates are already in micrometers. -- Exported `distance` values are consumed as micrometers by interpretability tier-2 analyses. -- If `data.coord_scale_um` changes, regenerate processed artifacts before running prediction/export. - -## Legacy Parity Mapping - -High-value legacy utilities from `working_version` are mapped as: - -- prediction collectors (`collect_predict_for_all_nodes*`) -> unified - `collect_inference_payload` in `src/grass_mil/inference/collectors.py` -- full-graph aggregation (`full_graph_*`) -> `src/grass_mil/inference/aggregation.py` -- evaluation metrics (`graph_*_evaluate_fn`) -> `src/grass_mil/inference/metrics.py` -- embedding collectors -> `src/grass_mil/inference/collectors.py` + `src/grass_mil/inference/io.py` -- interpretability table bridge -> `src/grass_mil/inference/interpretability_export.py` diff --git a/docs/interpretability_extension_guide.md b/docs/interpretability_extension_guide.md deleted file mode 100644 index 051ad65..0000000 --- a/docs/interpretability_extension_guide.md +++ /dev/null @@ -1,200 +0,0 @@ -# Interpretability Extension Guide - -This guide describes the recommended workflow for adding new interpretability analyses. - -Core rule: - -- add new analysis logic in `core/` or `tier2/` -- expose it through a plugin -- enable it via config -- validate with unit + pipeline tests - -The pipeline should not need structural edits for each new method. - -## 1) Decide Where Logic Belongs - -Use `core/` when: - -- analysis is generally applicable and does not require spatial graph topology. - -Use `tier2/` when: - -- analysis uses spatial neighborhood/connectivity structures or higher-order post-hoc behavior. - -Examples: - -- spatial entropy: `tier2/` -- boundary rugosity: `tier2/` -- non-spatial embedding reliability score: `core/` - -## 2) Implement Pure Analysis Functions - -Create a module: - -- `src/grass_mil/interpretability/core/.py` -- or `src/grass_mil/interpretability/tier2/.py` - -Guidelines: - -1. Keep function signatures explicit. -2. Validate required columns early and fail with clear errors. -3. Return typed outputs/dataclasses (or pandas frames) without side effects. -4. Avoid direct plotting in compute functions. - -Example skeleton: - -```python -from __future__ import annotations - -import pandas as pd - -def compute_spatial_entropy( - node_table: pd.DataFrame, - spatial_table: pd.DataFrame, - *, - label_column: str = "niche_label", -) -> pd.DataFrame: - if label_column not in node_table.columns: - raise ValueError(f"Missing label column {label_column!r}.") - # compute ... - return pd.DataFrame(...) -``` - -## 3) Create A Plugin Wrapper - -Add plugin class in: - -- `src/grass_mil/interpretability/plugins/builtin.py` -- or a new plugin module imported at registration time. - -Plugin responsibilities: - -1. Read context state (e.g., niche labels). -2. Call your pure analysis function. -3. Return `PluginResult` with: -- `payload` for machine consumption -- `sections` for report composition - -Example skeleton: - -```python -@dataclass -class SpatialEntropyPlugin(InterpretabilityPlugin): - name: str = "spatial_entropy" - - def required_inputs(self) -> list[str]: - return ["spatial_table", "niche_labels"] - - def run(self, dataset, context: PluginContext, **params): - table = dataset.instance_table.copy() - table["niche_label"] = context.state["niche_labels"] - entropy_df = compute_spatial_entropy( - table, - dataset.spatial_table, - label_column=str(params.get("label_column", "niche_label")), - ) - return PluginResult( - name=self.name, - payload={"entropy": entropy_df}, - sections=[ - ReportSection( - title="Spatial Entropy", - description="Entropy by niche.", - tables={"entropy": entropy_df}, - ) - ], - ) -``` - -`required_inputs()` entries are checked by `run_interpretability_pipeline(...)` before -`run(...)` is called. If any required input is missing from dataset fields/context state, -the pipeline raises `ValueError` immediately. - -## 4) Register The Plugin - -Register in `register_builtin_plugins(registry)`: - -```python -def register_builtin_plugins(registry: PluginRegistry) -> None: - registry.register(SpatialEntropyPlugin()) -``` - -The pipeline builds a fresh per-run registry (`create_builtin_registry()`), so plugin -registration is instance-scoped and not process-global. - -Duplicate plugin names are rejected by default. If you intentionally want to -replace a previously registered plugin implementation, use -`registry.register(plugin, replace=True)`. - -## 5) Add Config Wiring - -Update: - -- `configs/interpretability/plugins/default.yaml` - -Add: - -1. plugin name to `enabled` list (if default-on behavior desired) -2. plugin parameter block under `params.` - -Example: - -```yaml -enabled: - - niche_profiles - - spatial_entropy -params: - spatial_entropy: - label_column: niche_label - radius: 50.0 -``` - -## 6) Optional: Add Custom Plotly Visuals - -If table-only output is insufficient: - -1. Add figure builder in `src/grass_mil/interpretability/reporting/plotly_builders.py`. -2. Detect plugin payload and create a figure. -3. Ensure figure naming is stable for snapshot paths. - -Guideline: - -- keep plotting dependent on payload shape contracts, not ad hoc column discovery. - -## 7) Add Tests (Required) - -Minimum: - -1. Unit test for compute function. -2. Plugin test validating payload/section structure. -3. Pipeline integration test with plugin enabled. - -Suggested files: - -- `tests/interpretability/test_.py` -- extend `tests/interpretability/test_plugins_and_tier2.py` -- extend `tests/interpretability/test_pipeline_and_report.py` - -## 8) Extension Checklist - -Before merging, ensure: - -1. Compute function is side-effect free. -2. Plugin handles missing required inputs with explicit errors. -3. Config defaults exist and are documented. -4. Report section renders without runtime-only assumptions. -5. Unit and integration tests pass. -6. Doc updates added to `interpretability_suite.md` if behavior is user-facing. - -## 9) Recommended Pattern For Future Methods - -For methods like spatial entropy or boundary rugosity: - -1. Add feature calculator in `tier2/`. -2. Add plugin + default params. -3. Add one table section and one plot section. -4. Keep pipeline unchanged. - -This preserves the architecture promise: - -- new analyses are plug-ins, not pipeline rewrites. diff --git a/docs/interpretability_suite.md b/docs/interpretability_suite.md deleted file mode 100644 index 098b051..0000000 --- a/docs/interpretability_suite.md +++ /dev/null @@ -1,760 +0,0 @@ -# Interpretability Suite - -This document describes the interpretability subsystem in `grass-mil`. - -Scope: - -- architecture and design principles -- normalized analysis inputs -- reduction/clustering APIs and config controls -- plugin execution model (core + tier-2) -- report generation (interactive HTML + static PDF) -- operational runbook and troubleshooting - -## 1) Purpose And Design Principles - -The interpretability suite is designed as a first-class system, not a notebook helper. - -Key principles: - -1. Extensibility first. -2. Parameterized algorithms (no hardcoded DR/clustering assumptions). -3. Reproducible artifacts and reports. -4. Separation of concerns: -- analysis numerics -- plugin orchestration -- report rendering -5. Stable Python API for both core and tier-2 analyses. - -Package roots: - -- `src/grass_mil/interpretability/` -- `configs/interpretability/` -- `tests/interpretability/` - -## 2) High-Level Architecture - -Core layers: - -1. Data contracts and loaders -- `src/grass_mil/contracts.py` (shared package-level contracts, not - interpretability-specific) -- `src/grass_mil/interpretability/core/data.py` - -2. Core numerical analysis -- `src/grass_mil/interpretability/core/reduction.py` -- `src/grass_mil/interpretability/core/clustering.py` -- `src/grass_mil/interpretability/core/niches.py` (naming and dendrogram order) -- `src/grass_mil/interpretability/core/biomarkers.py` -- `src/grass_mil/interpretability/core/attribution.py` (margin attribution) -- `src/grass_mil/interpretability/core/agreement.py` (cross-space agreement) -- `src/grass_mil/interpretability/core/transfer.py` (niche label transfer) - -3. Tier-2 analyses -- `src/grass_mil/interpretability/tier2/neighborhood.py` -- `src/grass_mil/interpretability/tier2/filtration.py` -- `src/grass_mil/interpretability/tier2/autocorrelation.py` (Moran's I) -- `src/grass_mil/interpretability/tier2/ripley.py` (Ripley cross-L) -- `src/grass_mil/interpretability/tier2/tissue_graph.py` - -4. Plugin orchestration -- `src/grass_mil/interpretability/plugins/base.py` -- `src/grass_mil/interpretability/plugins/registry.py` -- `src/grass_mil/interpretability/plugins/builtin.py` -- `src/grass_mil/interpretability/pipeline.py` - -5. Reporting -- `src/grass_mil/interpretability/reporting/plotly_builders.py` -- `src/grass_mil/interpretability/reporting/sections.py` -- `src/grass_mil/interpretability/reporting/html.py` -- `src/grass_mil/interpretability/reporting/snapshot.py` -- `src/grass_mil/interpretability/reporting/pdf.py` -- `src/grass_mil/interpretability/reporting/render.py` - -6. CLI -- `src/grass_mil/interpretability/report_cli.py` - -## 3) Normalized Input Contracts - -The suite consumes normalized tables, independent of model internals. - -### Required table: `instance_table` (`.csv` or `.parquet`) - -Required minimal columns: - -- `instance_id` -- `bag_id` -- embedding columns with one of supported prefixes: - - `inst_emb_*` - - `emb_*` - - `graph_emb_*` - -Common optional columns: - -- `region_id` -- `sample_id` -- `cell_type` -- `score` -- `attention` -- `hazard` -- composition columns `comp_*` - -`score` convention: - -- exported inference tables use configurable score projection from `instance_logits`: - - `score_mode=sigmoid` (default): `score = sigmoid(logit[score_logit_index])` - - `score_mode=identity`: `score = logit[score_logit_index]` -- `score_mode=softmax`: `score = softmax(logits)[score_logit_index]` -- defaults (`sigmoid`, index `0`) keep notebook-era binary classification parity. -- for survival/hazard-style interpretation, prefer `score_mode=identity`. -- `instance_attention_logits` with one column per class are exported as - `attention_c` (within-bag softmax) and `attention_logit_c` (raw) columns; - `attention` holds the canonical class (`attention_class_index`, default the last). - -Note: `niche_profiles` requires `comp_*` columns. -There is no fallback to one-hot `cell_type` composition. - -`variance_estimator` convention (biomarker z-scoring): - -- `unbiased` (default): sample variance/std (`ddof=1`) -- `biased`: population variance/std (`ddof=0`, notebook-style scaling) - -### Optional table: `bag_table` - -Used for downstream joins and metadata context. Not required by default pipeline math. - -### Optional table: `spatial_table` - -Required for tier-2 spatial plugins. - -Required columns for current tier-2 methods: - -- `source_id` -- `target_id` -- `distance` (required by filtration curves) -- `weight` (optional; currently ignored by neighborhood enrichment, which is unweighted/count-based) - -Unit contract: - -- During preprocessing, spatial coordinates are converted to micrometers (`um`) via `data.coord_scale_um`. -- `data.coord_scale_um=1.0` means input coordinates are already in `um`. -- `spatial_table.distance` is treated as micrometers by tier-2 analyses. -- If `data.coord_scale_um` changes, regenerate processed artifacts (`data.force_precompute=true` or a new `data.processed_dir`) before export/report. - -## 4) Python API - -Public exports: - -- `run_reduction(...)` -- `run_clustering(...)` -- `niche_composition_summary(...)` -- `run_neighborhood_enrichment(...)` -- `run_diff_neighborhood_enrichment(...)` -- `compute_filtration_curves(...)` -- `prepare_tissue_graph_view(...)` -- `build_tissue_graph_figure(...)` -- `fit_niche_transfer_from_report_bundle(...)` -- `apply_niche_transfer(...)` -- `save_niche_transfer_bundle(...)` -- `load_niche_transfer_bundle(...)` -- `run_interpretability_pipeline(...)` -- `render_interpretability_report(...)` - -Import path: - -```python -from grass_mil.interpretability import ( - run_reduction, - run_clustering, - prepare_tissue_graph_view, - build_tissue_graph_figure, - fit_niche_transfer_from_report_bundle, - apply_niche_transfer, - run_interpretability_pipeline, - render_interpretability_report, -) -``` - -Composition summary API: - -- `niche_composition_summary` operates on explicit composition columns - (`comp_` by default, via `composition_prefix`). -- `cell_type_column` is not part of the signature. - -The attention-based summaries that used to sit here -(`cluster_attention_summary`, `cluster_survival_attention_summary`) have been -removed. Attention alone is not an attribution: it says where the model looked, -not which way that evidence pushed the prediction. Use -`niche_attribution_summary` from `grass_mil.interpretability.core.attribution` -instead, which reports attention-weighted logit margins and satisfies the -attribution identity in [`method_fidelity.md`](method_fidelity.md). Survival -routes through the same function with hazard as the single class. - -### Neighborhood Enrichment Behavior - -`run_neighborhood_enrichment(...)`: - -- Uses unweighted edge counts (edge `weight` is ignored). -- `undirected=true` mirrors each edge (`u->v` and `v->u`). -- In `undirected=true`, unordered node pairs are deduplicated before mirroring so input tables - that already contain both directions are not double-counted. -- `enrichment_mode` controls the returned enrichment matrix: - - `zscore`: `(observed - expected) / std` - - `obs-exp`: `observed - expected` - - `log2fc`: `log2(observed / expected)` with epsilon clipping for stability -- `n_perms > 0`: - - expected and std are estimated from label permutations - - enrichment uses the selected `enrichment_mode` - - p-values are two-sided empirical permutation p-values -- `n_perms = 0`: - - uses analytical expected/std approximation - - emits a warning that analytical mode is active - -`run_diff_neighborhood_enrichment(...)`: - -- Computes pairwise condition differences: - `metric(cond_a) - metric(cond_b)` where `metric` is selected by `enrichment_mode`. -- Per-condition baselines are analytical. -- If `n_perms > 0`, only the differential p-values are permutation-based. -- Differential permutations are performed at the library/sample level via `permutation_group_column` (defaults to `sample_id`), matching notebook behavior. -- Each permutation group id must map to exactly one condition. -- Emits a warning that analytical per-condition baselines are used. - -## 5) Reduction And Clustering APIs - -### `run_reduction(method, X, **params)` - -Supported methods: - -- `pca` -- `umap` -- `tsne` - -Common parameters: - -- `n_components` -- `random_state` -- method-specific options (`n_neighbors`, `min_dist`, `metric`, `perplexity`, etc.) - -Returns a `ReductionResult` dataclass with: - -- reduced embedding matrix -- resolved parameters -- fitted object - -### `run_clustering(method, X, **params)` - -Supported methods: - -- `kmeans` -- `agglomerative` -- `hdbscan` -- `spectral` -- `optics` -- `dbscan` - -Common parameters: - -- `n_clusters` (where required) -- `random_state` -- method-specific options (`min_cluster_size`, `metric`, `linkage`, `eps`, etc.) - -Returns a `ClusteringResult` dataclass with: - -- niche labels -- resolved parameters -- fitted object - -## 6) Plugin Execution Model - -The pipeline runs plugin list from config after reduction/clustering. - -Plugin protocol (`InterpretabilityPlugin`): - -- `name` -- `required_inputs()` -- `run(dataset, context, **params) -> PluginResult` - -`required_inputs()` is enforced by the pipeline before plugin execution; missing inputs -fail fast with a `ValueError` naming the plugin and missing keys. - -Plugin registration is per pipeline run (fresh `PluginRegistry`). - -`PluginResult` includes: - -- `payload`: machine-readable outputs -- `sections`: report sections (tables/metadata, optional visuals) - -Built-in plugins: - -| Plugin | Unit | -|---|---| -| `niche_profiles` | instances, composition over all their cells | -| `margin_attribution` | instances | -| `cell_type_enrichment_per_niche` | cells, within each niche | -| `cell_type_moran_per_niche` | cells, within each niche | -| `cell_filtration_curves` | cell-cell edges, within each niche | -| `cell_type_ripley_per_niche` | cells, within each niche | -| `cell_type_diff_enrichment_by_condition` | cells, differenced between conditions | -| `niche_label_moran` | instances, globally | -| `niche_agreement` | instances | -| `tissue_graph` | instances (visualisation) | - -The tier-2 statistics are cell-level: the instance-level versions were removed, -because treating a whole ego-graph as one node measures how neighbourhood -*labels* relate rather than how cells do. Two niches with identical composition -but opposite spatial arrangement are indistinguishable to an instance-level -analysis. - -The cell-level plugins need `cell_table` and `cell_edge_table`, exported by -`grass-mil-predict interpretability.cells.enabled=true` and passed to the report -as `data.cell_table` / `data.cell_edge_table`. - -Three plugins are not in any shipped preset, because each needs an input the -pipeline cannot infer: - -| Plugin | Additional input required | -|---|---| -| `cell_type_diff_enrichment_by_condition` | a condition column, plus a permutation group column in permutation mode | -| `tissue_graph` | an explicit `sample_value` (no auto-pick) | -| `niche_agreement` | a second niche labelling column to compare against | - -`cell_filtration_curves` default threshold grid: - -- `np.linspace(0.0, 55.0, 500)` -- thresholds are interpreted in micrometers - -`tissue_graph` behavior: - -- disabled by default in `plugins.enabled` -- requires explicit `sample_value` selection (no auto-pick) -- requires `center_x` and `center_y` coordinates in `instance_table` -- colors nodes by pipeline `niche_labels` context -- uses `spatial_table` edges filtered to selected sample nodes - -## 7) Pipeline Orchestration - -`run_interpretability_pipeline(...)` is the single coordinator. - -Workflow: - -1. Load embedding matrix from `instance_table`. -2. Run configured reduction. -3. Run configured clustering: -- `cluster_on_pca=true`: niche on PCA projection used for clustering. -- `cluster_on_pca=false`: niche on raw embedding columns. -4. Build core niche summaries. -5. Execute enabled plugins in order. -6. Persist canonical artifacts: -- `artifacts/instance_with_niches.csv` -- `artifacts/reduction.csv` (if enabled) -- `artifacts/niche_labels.csv` (if enabled) -- `artifacts/pipeline_summary.json` -7. Return `ReportBundle`. - -`ReportBundle` includes `cluster_feature_reduction`: - -- populated with clustering PCA result when `cluster_on_pca=true` -- `None` when clustering runs on raw embeddings - -## 8) Niche Transfer Workflow (Library API) - -Notebook-parity niche transfer is available as a library workflow: - -1. Fit transfer bundle on a source interpretability run (niche labels + clustering feature space). -2. Persist bundle (`.joblib`) for deterministic reuse. -3. Apply bundle to a query `instance_table` to obtain transferred labels. - -Design notes: - -- Uses kNN over the exact clustering feature space used in the source run. -- If source clustering used PCA (`cluster_on_pca=true`), transfer uses that same PCA projector. -- HDBSCAN noise label `-1` is retained as a transferable class. -- This replaces fragile notebook assumptions like `clusterer._embedding_`. - -Example: - -```python -from pathlib import Path - -import pandas as pd - -from grass_mil.interpretability.core.data import load_interpretability_dataset -from grass_mil.interpretability.core.transfer import ( - apply_niche_transfer, - fit_niche_transfer_from_report_bundle, - load_niche_transfer_bundle, - save_niche_transfer_bundle, -) -from grass_mil.interpretability.pipeline import run_interpretability_pipeline - -dataset = load_interpretability_dataset(instance_table_path=Path("/abs/source_instance_table.csv")) -bundle = run_interpretability_pipeline( - dataset, - { - "reduction": {"enabled": True, "method": "pca", "params": {"n_components": 2}}, - "clustering": { - "enabled": True, - "method": "hdbscan", - "cluster_on_pca": True, - "pca_components": 10, - "params": {"min_cluster_size": 100, "min_samples": 1}, - }, - "plugins": {"enabled": [], "params": {}}, - }, - artifacts_dir=Path("/abs/source_artifacts"), -) -transfer = fit_niche_transfer_from_report_bundle(dataset, bundle, n_neighbors=15) -save_niche_transfer_bundle(transfer, "/abs/cluster_transfer_bundle.joblib") - -loaded = load_niche_transfer_bundle("/abs/cluster_transfer_bundle.joblib") -query_table = pd.read_csv("/abs/query_instance_table.csv") -transferred = apply_niche_transfer(loaded, query_table) -``` - -Notebook helper example for single tissue graph: - -```python -from grass_mil.interpretability.tier2.tissue_graph import ( - build_tissue_graph_figure, - prepare_tissue_graph_view, -) - -view = prepare_tissue_graph_view( - node_table=instance_with_clusters_df, - spatial_table=spatial_table_df, - sample_column="sample_id", - sample_value="sx", - id_column="instance_id", - x_column="center_x", - y_column="center_y", - label_column="niche_label", -) -fig = build_tissue_graph_figure(view, show_edges=True) -fig.show() -``` - -## 9) Reporting System - -`render_interpretability_report(...)` generates: - -1. Interactive report: `report.html` -2. Static snapshots: `artifacts/snapshots/*.png` (Plotly + kaleido) -3. Static PDF: `report.pdf` (Playwright print pipeline) - -### HTML report - -- Plotly interactive sections. -- Tabular summaries from core and plugin outputs. -- Structured section breakdown with metadata. - -### PDF report - -Design intent: - -- publication-stable, deterministic layout -- static figure rendering -- explicit print CSS and margins - -Engine: - -- Playwright Chromium print-to-PDF. - -## 10) CLI Usage - -Generate interpretability-ready tables directly from inference: - -```bash -grass-mil-predict \ - ckpt_path=/abs/path/model.ckpt \ - data=spatial_omics \ - interpretability.enabled=true \ - interpretability.spatial.enabled=true -``` - -This writes `instance_table.csv` (and optionally `spatial_table.csv`) into -`${paths.output_dir}/${predict.output_subdir}` and records paths in -`inference_summary.json` under `interpretability.*`. - -Entrypoint: - -```bash -grass-mil-report \ - data.instance_table=/abs/path/instance_table.csv -``` - -With spatial plugins: - -```bash -grass-mil-report \ - data.instance_table=/abs/path/instance_table.csv \ - data.spatial_table=/abs/path/spatial_table.csv \ - plugins.enabled="[niche_profiles,margin_attribution,neighborhood_enrichment,diff_neighborhood_enrichment,filtration_curves]" -``` - -Render a single selected tissue graph colored by niche labels: - -```bash -grass-mil-report \ - data.instance_table=/abs/path/instance_table.csv \ - data.spatial_table=/abs/path/spatial_table.csv \ - plugins.enabled="[tissue_graph]" \ - plugins.params.tissue_graph.sample_column=sample_id \ - plugins.params.tissue_graph.sample_value=sx -``` - -Override algorithms: - -```bash -grass-mil-report \ - data.instance_table=/abs/path/instance_table.csv \ - reduction=umap \ - clustering=dbscan \ - clustering.params.eps=0.25 \ - clustering.params.min_samples=10 -``` - -Tune HDBSCAN selection behavior explicitly: - -```bash -grass-mil-report \ - data.instance_table=/abs/path/instance_table.csv \ - clustering=hdbscan \ - clustering.params.min_cluster_size=100 \ - clustering.params.min_samples=5 \ - clustering.params.cluster_selection_method=leaf \ - clustering.params.cluster_selection_epsilon=0.0 -``` - -Enable/disable report outputs: - -```bash -grass-mil-report \ - data.instance_table=/abs/path/instance_table.csv \ - report.pdf.enabled=true \ - report.html.enabled=true -``` - -## 11) Config Groups - -Main config: - -- `configs/interpretability/report.yaml` - -Subgroups: - -- `configs/interpretability/reduction/*.yaml` -- `configs/interpretability/clustering/*.yaml` -- `configs/interpretability/plugins/default.yaml` — `niche_profiles` only -- `configs/interpretability/plugins/full.yaml` — composes `default` and adds - `margin_attribution`, the four cell-level analyses, and `niche_label_moran` - -Select a preset with: - -```bash -grass-mil-report interpretability/plugins@plugins=full -``` - -The `@plugins` package part is required. `report.yaml` is `# @package _global_`, -so the shorter `plugins=full` does **not** load the preset — it silently -overwrites the `plugins` node with the literal string `"full"`, and the run then -fails downstream with a confusing error rather than at the override. - -`full.yaml` is materially slower than the default: Moran's I and Ripley both -scale with niche size and permutation count. It does not enable the four -plugins listed above that need inputs the pipeline cannot infer. - -Important knobs: - -- `reduction.method` -- `reduction.params.*` -- `clustering.method` -- `clustering.params.*` -- `clustering.cluster_on_pca` -- `clustering.pca_components` -- `plugins.enabled` -- `plugins.params..*` -- `variance_estimator` -- `embedding_prefixes` -- `report.html.enabled` -- `report.pdf.enabled` -- `report.pdf.snapshot_dpi_scale` - -Notebook-aligned defaults: - -- `reduction=umap` with `n_neighbors=50`, `min_dist=0.1`, `random_state=null` -- `clustering=hdbscan` with `min_cluster_size=100`, `min_samples=1`, `cluster_selection_method=eom` -- `clustering.cluster_on_pca=true`, `clustering.pca_components=10` -- effect: clustering runs on PCA(10) while report scatter uses UMAP output -- note: PCA components are capped to `min(n_samples, n_features)` on small datasets - -Parameter contract: - -- `clustering.params.*` is forwarded directly to the underlying sklearn clustering estimator. -- No method-specific parameter filtering is applied beyond common `n_clusters`/`random_state` handling. - -## 12) Environment Prerequisites - -From project root: - -```bash -pip install -e ".[interpretability,report,spatial-stats,io]" -``` - -For PDF export, install Playwright browser binaries: - -```bash -playwright install chromium -``` - -If PDF export fails, check: - -1. `playwright` package installed. -2. Chromium binaries installed. -3. sandbox constraints for headless browser runtime. - -If snapshot export fails, check: - -1. `kaleido` package installed. -2. Plotly-kaleido compatibility in your environment. - -## 13) Typical Operational Flow - -1. Generate/prepare normalized instance/spatial tables. -2. Run report CLI with selected reduction/clustering presets. -3. Inspect `report.html` interactively. -4. Share `report.pdf` for static dissemination. -5. Consume `artifacts/*.csv` and JSON for downstream analysis scripts. - -## 14) Troubleshooting - -### "No embedding columns found" - -Cause: - -- `instance_table` has no columns matching configured prefixes. - -Fix: - -- rename columns to `inst_emb_*` (or set `embedding_prefixes` override). - -### "Spatial table is required for ... plugin" - -Cause: - -- tier-2 plugin enabled without `data.spatial_table`. - -Fix: - -- provide `spatial_table`, or remove plugin from `plugins.enabled`. - -### "Playwright is unavailable" - -Cause: - -- Python package missing or browser binaries not installed. - -Fix: - -- `pip install playwright` -- `playwright install chromium` - -### "Failed to export plot snapshot" - -Cause: - -- `kaleido` missing or incompatible runtime. - -Fix: - -- `pip install kaleido` - -## 15) Relationship To Other Docs - -- prediction/inference payload generation: `docs/inference_utility_suite.md` -- training/eval workflows: `docs/workflows_train_infer.md` -- extension instructions: `docs/interpretability_extension_guide.md` - -## Niche vocabulary - -Outputs use the language a biologist reads, not the algorithm's: - -- The groups are **niches** — recurring local tissue neighbourhoods — not - "clusters". Tables carry a `niche_label` column and figures name niches - `Niche 3`, keeping the numeric id so a figure joins back to the table. -- Label `-1` is **Background**, not "noise". Density-based clustering assigns it - to neighbourhoods that fit no dense group; those are usually transitional or - sparsely populated tissue, which is a meaningful category rather than a - failure. Set `skip_background: false` on the per-niche plugins to include it. -- Niches are ordered by a **dendrogram over their cellular composition** - (`core/niches.py::order_niches_by_composition`), so compositionally similar - niches sit next to each other in every table and heatmap instead of appearing - in arbitrary label order. Background is pinned last: it is a catch-all, and - letting it join the tree distorts the ordering. Correlation distance is the - default, so two niches with the same makeup at different densities read as - similar. Pass `order_by_composition=False` for plain numeric order. - -The *algorithm* is still clustering, so `run_clustering`, `ClusteringResult` and -the sklearn parameters (`n_clusters`, `min_cluster_size`) keep their names. The -distinction is deliberate: those are the method, the niches are the result. - -The report's clustering step also accepts the GRASS-MIL recipe used by -`grass-mil-niches`: `clustering.sample_size` (fit on a random sample), -`clustering.l2_normalise` (cluster on L2-normalised coordinates, i.e. cosine -geometry) and `clustering.assign_k` (k-NN assignment of the remaining instances). -The margin attribution plugin checks the additive identity against the bag -logits of `data.bag_table` (`grass-mil-predict` predictions) and reports the -residual. - -## What the tier-2 statistics are computed over - -The instance-level tier-2 analyses run on the **instance graph**; the cell-level -analyses (next section) run on the cell graph when `data.cell_table` and -`data.cell_edge_table` are given. Confusing the two makes the results easy to -over-read, so be explicit about it: - -- A row of `instance_table` is one **k-hop ego-graph**, rooted at one cell. -- An edge in `spatial_table` connects two **instances**, i.e. two roots. -- `comp_*` are the cell-type fractions among **all cells inside** that - ego-graph. This is the only place the full cell content survives. -- `cell_type` is the **root cell's own type**, so that the instance graph can be - treated as a (sampled) cell graph. -- `center_x`/`center_y` are the **mean position of all cells** in the ego-graph, - not the root cell's coordinates. - -What each analysis therefore measures: - -| Analysis | Unit | Measures | -|---|---|---| -| `cell_type_enrichment_per_niche` | cell pairs inside a niche | how often two cell types are in contact, versus chance | -| `cell_type_moran_per_niche` | cells inside a niche | whether each cell type forms contiguous patches | -| `cell_filtration_curves` | cell-cell edges inside a niche | cells of each type reached as the distance threshold grows | -| `cell_type_ripley_per_niche` | cells inside a niche | cross-L between cell types, from their own coordinates | -| `niche_label_moran` | instances, globally | whether niches form contiguous territories | -| `niche_profiles` | instances | mean `comp_*` per niche, which reflects every cell in the ego-graphs | - -**Roots are a weighted sample, not every cell.** With -`data.sampler.runtime.proportional_root_sampling=true` (the default) roots are -drawn with inverse-frequency weighting by `property_name`, so rare cell types -are deliberately over-represented and one cell can root several instances. Any -cell-type distribution over instances therefore reflects that sampling design -rather than the tissue. Set `proportional_root_sampling=false` for an unweighted -draw. - -### Cell-level analyses - -Export `cell_table` and `cell_edges` with -`grass-mil-predict interpretability.cells.enabled=true`, pass them to the report -as `data.cell_table` / `data.cell_edge_table`, and these become available: - -| Analysis | Unit | Measures | -|---|---|---| -| `cell_type_enrichment_per_niche` | cell pairs within a niche | cell-type x cell-type contact enrichment, using only edges whose both endpoints lie in the niche; optionally differenced against Background | -| `cell_type_moran_per_niche` | cells within a niche | Moran's I of each cell-type indicator, as a niche x cell-type table | -| `cell_filtration_curves` | cell pairs within a niche | cells of each type reached as the edge-distance threshold grows, counting both endpoints | -| `cell_type_ripley_per_niche` | cells within a niche | centred cross-L between cell types, from the cells' own coordinates | -| `niche_label_moran` | instances, globally | whether niches form contiguous territories rather than interleaving | - -Each cell inherits the niche of the ego-graph it was sampled into, so a cell -appearing in several overlapping ego-graphs contributes to each. This mirrors -the pooling the legacy NSCLC reports do. diff --git a/docs/method_correspondence.md b/docs/method_correspondence.md deleted file mode 100644 index ce8bc82..0000000 --- a/docs/method_correspondence.md +++ /dev/null @@ -1,52 +0,0 @@ -# Method Correspondence With The Preprint - -Each item of the preprint's method maps to a preset, a module and a test. The -settings are the defaults of the presets; all remain overridable. - -| Preprint element | Default | Config | Code | Test | -|---|---|---|---|---| -| Delaunay cell graph, contact edges, bounded log edge length `clip((ln d - 2)/3, 0, 1)` | `edge_features=[log_distance, neighbor]`, `neighbor_cutoff_um` in the data's units | `data/spatial_omics.yaml` (`graph_builder`) | `data/components/graph_builders.py::log_distance_feature` | `tests/models/test_fidelity_phase1.py` | -| Node input: cell-type embedding + linear projection of continuous features, xavier init | `input_init=xavier`, `categorical_embedding` | `model/encoder/gine_length.yaml` | `models/components/embeddings.py` | phase 1 | -| GINE encoder 4 x 256, learnable epsilon, layer norm, max pooling over the ego-graph | `num_layers=4`, `hidden_dim=out_dim=256`, `train_eps=true`, `norm=layernorm`, `pooling=max` | `model/encoder/gine_length.yaml` | `models/components/backbones.py` | phase 1 | -| Ego-graphs: k = 4 hops, all neighbours, no radius cap | `depth=4`, `num_neighbors=-1`, `ego_radius.enabled=false` | `data/sampler/ego_graph_mil.yaml` | `data/components/samplers.py` | `tests/models/test_fidelity_phase2.py` | -| BGRL pretraining "A3": roots from all regions, 64 per step, AdamW 5e-5 / 1e-5, warmup 1000, 200k steps, momentum 0.99 -> 1, encoder dropout 0.1, DropImportance (mu 0.25, p_lambda 0.45) + feature noise 0.2 on size, checkpoint every 5000 steps | `regime/pretrain_a3.yaml` (alias `experiment/grass_mil_pretrain_a3.yaml`) | `model/bgrl_a3.yaml`, `task/pretrain_bgrl_a3.yaml`, `data/sampler/ego_graph_pretrain.yaml`, `callbacks/pretrain_checkpoint.yaml` | `models/bgrl_module.py`, `data/components/merge.py`, `models/training/augmentations.py` | `tests/models/test_fidelity_phase4.py` | -| Projected gated attention (32 / 8, dropout 0.25, one channel per class, kaiming) | `model/attention/gated_projected.yaml` | | `models/components/attention.py` | phase 1 | -| Logit-space pooling `L_c = sum_i A[i,c] l[i,c]`, exact per-cell contributions | `task.aggregation=mil_attention` | `task/finetune_mil.yaml` | `models/training/bagging.py`, `interpretability/core/attribution.py` | `tests/models/test_bagging.py`, phase 6 | -| One region = one bag, one softmax, one loss; class-weighted CE (train split, mean 1); fixed 1024-root validation; all-cells test | `unit_as_bag=true`, `subsample_count`, `class_weights=auto_train`, `val_sampler`, `test_sampler=null` | `regime/mil.yaml` | `models/supervised_module.py`, `models/training/class_weights.py` | phase 2 | -| Probe pre-fit: 96 roots per region, StandardScaler + LogisticRegression(C 1, balanced), folded into a linear head | `probe.yaml` | `regime/linear_*.yaml` (`task.head_init_path=auto`) | `models/training/probe.py` | `tests/models/test_fidelity_phase3.py` | -| Regimes probehead / ftALL / ftALL-reg: learning-rate scales (encoder 0 / 1 / 0.3, head 0 / 0.01, attention 2), weight decay 1e-3 / 1e-2, dropout 0.2 and edge/feature drop 0.05 in the regularised regime, bag 512 / 256, per-step cosine over the run, hyperbatch 12, clip 1, AMP, early stopping on region loss (patience 8, min 10 epochs), checkpoint on region AUROC | `regime/*.yaml`, `scheduler/cosine_mil_step.yaml`, `callbacks/mil_selection.yaml` | | `models/training/optimization.py`, `models/supervised_module.py` | phase 2 | -| Nested CV: patient-grouped outer folds, inner `StratifiedGroupKFold(3, seed 0)`, best epoch = argmax val AUROC, median refit, test once | `nested_cv.yaml` | | `nested_cv.py`, `data/components/splits.py` | phase 3 | -| Patient score: mean of core probabilities (merged bag as alternative) | `aggregation.bag_scope=patient` | `inference/aggregation/default.yaml` | `inference/aggregation.py` | `tests/models/test_fidelity_phase5.py` | -| Cell cache of every fold model (E fp16, instance logits, raw attention, bias) | `extract.yaml` | | `inference/cell_cache.py` | phase 5 | -| relattn margin `N_r [A_1 (l_1 - b_1) - A_0 (l_0 - b_0)]` | | | `interpretability/niches/margins.py` | `tests/interpretability/test_fidelity_phase6.py` | -| Cosine HDBSCAN: PCA-10 on 100k cells, L2, mcs 300, min_samples 1, eom, 15-NN assignment; niche table, co-localisation, figures | `niches.yaml` | | `interpretability/core/clustering.py::cosine_hdbscan`, `interpretability/niches/*` | phase 6 | -| Metaclusters: per-fold clustering, Hellinger distance, constrained average linkage, threshold scan or 0.40, patient shares (each core labelled once), transfer, pairwise ARI, shared k-means basis | `metaclusters.yaml` | | `interpretability/metaclusters/*` | `tests/interpretability/test_fidelity_phase7.py` | -| Bootstrap AUROC CI (2000, seed 0), DeLong, within-fold z / rank, operating points at within-fold quantile rules | `scores.yaml` | | `evaluation/*` | phase 5 | -| Neighbourhood enrichment (200 permutations), cellular neighbourhoods (10-NN windows, k-means 10 on 200k windows, CLR), contacts, cross-K (r 20/50/100, pi r^2) | `spatial_stats.yaml` | | `spatial_stats/*` | `tests/spatial_stats/test_fidelity_phase8.py` | - -## Known deviations - -- The LR warmup starts at `start_factor=0.0` through `LinearLR`, which clamps to - 1e-4 of the base rate at step 0 (the reference starts at exactly 0); the - difference is one step. -- `grass-mil-nested-cv` rejects inner-fold seeds that leave a fold without one - class and tries the next seed; on cohorts with more than a handful of patients - per class the seed is never changed. -- The niche pipeline has no clustering sweep; it reports the settings it was - given (the preprint chose its settings by a sweep once and then fixed them). -- Validation is computed on a fixed 1024-root subset per region and test on all - cells, as in the reference; both are configurable. - -## Validation against the original runs - -Recorded after the fidelity check of the public package against the preprint's -NSCLC runs (private data; the converter lives outside the repository): - -| Check | Reference | Public package | Status | -|---|---|---|---| -| Frozen embeddings of three cores (per-cell cosine with the reference fp16 cache; 875-5150 cells each) | 1.0 | median 0.993-0.995, minimum 0.97, 91-95 % of cells >= 0.99 | passed | -| Probe region AUROC, outer fold 2: inner folds 0 / 1 / 2 and refit (scored on the held-out fold) | 0.482 / 0.304 / 0.564 / 0.661 | 0.471 / 0.354 / 0.555 / 0.685 | passed (max deviation 0.05 on inner fold 1, 0.024 on the refit) | -| Held-out AUROC, `linear_ftall_reg`, outer fold 2 (inner best epochs, refit length E) | 0.664 (1 / 1 / 1, E = 1) | 0.660 (3 / 1 / 3, E = 3) | passed (deviation 0.005) | -| Held-out AUROC, `linear_probehead`, outer fold 2 (inner best epochs, refit length E) | 0.627 (10 / 14 / 1, E = 10) | 0.665 (16 / 15 / 1, E = 15) | deviation 0.038, slightly above the 0.03 target; single-seed runs on 86 held-out regions (bootstrap CI about +/- 0.12) | -| Niches of the fold-2 refit (`linear_ftall_reg`, held-out cells, default cosine HDBSCAN) | 11 clusters, noise 0.78, Spearman(emp, relattn) +0.65; hypoxic-tumour clusters most positive-associated with positive margins, lymphoid aggregates strongly negative (-1.49) | 14 clusters, noise 0.75, Spearman +0.44; hypoxic-tumour clusters most positive-associated (relattn +0.45 and +0.29), B/T-cell aggregate -1.81, noise negative | qualitatively reproduced; the public refit is a different trained model (E = 3 vs 1), so cluster boundaries differ | -| Neighbourhood-enrichment z-scores on all 426 cores (numpy backend, 200 permutations) vs the reference squidpy table: Spearman of the patient-mean z per group pair, 225 pairs | > 0.9 | 0.984 (STAS+) and 0.977 (STAS-); the same six pairs lead the ranked test table | passed | diff --git a/docs/model_component_contracts.md b/docs/model_component_contracts.md deleted file mode 100644 index 40b0cb1..0000000 --- a/docs/model_component_contracts.md +++ /dev/null @@ -1,150 +0,0 @@ -# Model Component Contracts - -This document defines the canonical contracts for model components, losses, and samplers. - -## Goals - -- Keep interfaces minimal and composable. -- Keep behavior deterministic under fixed seeds. -- Avoid hidden mutable global state. -- Keep modules data-layer agnostic. - -## Graph Unit Contract - -A graph unit is any `torch_geometric.data.Data` object emitted by the data layer. -It may represent a patch, a region-level graph, or a full-sample graph. - -Required fields: - -- `x`: `torch.Tensor` with shape `[num_nodes, num_node_features]`, float dtype. -- `edge_index`: `torch.LongTensor` with shape `[2, num_edges]`. - -Optional fields: - -- `edge_attr`: `torch.Tensor` with shape `[num_edges, num_edge_features]`. -- `batch`: `torch.LongTensor` with shape `[num_nodes]` for batched data. -- `pos`: `torch.Tensor` with shape `[num_nodes, 2]` or `[num_nodes, d]`. -- `graph_y`: `torch.Tensor` with shape `[1, num_tasks]` or `[batch_size, num_tasks]`. -- `graph_w`: `torch.Tensor` with same shape as `graph_y` for weighted objectives. -- `categorical_codes`: `torch.LongTensor` for categorical annotations. - -Non-requirements: - -- Samplers must not assume graph units are patches. -- Model components must not depend on `sample_id`, `region_id`, or `patch_id`. - -## Encoder Contract - -Signature: - -- `forward(x, edge_index, edge_attr=None, batch=None, categorical_codes=None, - return_graph_embedding=False)`. - -Encoders return: - -- Node embeddings: shape `[num_nodes, hidden_dim_or_projected_dim]`. -- Optional graph embedding via pooling: shape `[num_graphs, pooled_dim]`. - -Supported jump-knowledge modes: - -- `last`, `concat`, `max`, `sum`. - -Categorical inputs: - -- `categorical_codes` is `[num_nodes, num_labels]` long. The encoder embeds one - configured column and **sums** it with the projection of `x`; it is not - concatenated, so the embedding width is `hidden_dim`. -- `input_dim == 0` is legal and means the cohort has no continuous node - features. A categorical embedding is then required. -- Out-of-range categorical codes raise. They are never clamped: clamping trains - on a wrong cell type silently. - -## Attention Contract - -Attention modules accept node/subgraph embeddings: - -- Input shape `[n_items, emb_dim]`. - -Attention modules return: - -- Unnormalized attention logits: shape `[n_items, n_classes]`. -- Pass-through embeddings (for compatibility): shape `[n_items, emb_dim]`. - -Normalization is handled by the caller to keep modules reusable: softmax over -`dim=0` (instances), applied **per class column**. - -Width rule: - -- `n_classes` must equal the head's `output_dim` **exactly**. A binary head needs - two attention channels, not one shared channel broadcast across both — a - shared channel cannot express a niche that pushes toward one class and away - from the other. -- Enforced by `grass_mil.contracts.validate_attention_width` at build time. - Width 1 is legal only when the head emits a single logit. - -`head_spaces()` exposes intermediate activations (`a_pre_tanh`, `gate_product`) -so analysis code does not reach into module internals. - -## Loss Contracts - -Categorical losses: - -- Predictions: logits with shape `[N, C]` or `[N, 1]` for binary. -- Targets: class ids `[N]` or binary labels `[N, 1]`. -- Optional weights: - - class weights for CE-style loss. - - per-sample weights shape `[N]` or `[N, 1]`. - -Regression losses: - -- Predictions and targets shape-compatible (`[N, D]` or `[N]`). -- Optional per-sample weights shape `[N]` or broadcast-compatible. - -Survival loss (`CoxSGD`): - -- `y_pred`: shape `[N]` or `[N, 1]`. -- `length`: observed times shape `[N]` or `[N, 1]`. -- `event`: event indicator shape `[N]` or `[N, 1]`, expected binary values. - -## Sampler Contracts - -Sampler strategies operate on graph units produced by the data layer. - -Identity strategy: - -- Batches graph units without neighborhood expansion. - -Neighborhood strategies: - -- Operate per graph unit. -- Must preserve graph-unit metadata unless explicitly documented otherwise. -- Must support optional post-hoc transform hooks. - -Native-first policy: - -- Prefer native PyG samplers when contracts can be satisfied. -- Use custom sampler only when required behavior is unavailable in native path. - -## Anti-Patterns - -- No script-specific hardcoding in components. -- No mutable global constants that alter runtime behavior. -- No implicit side effects in constructors. -- No reliance on deprecated APIs when modern equivalent exists. - -## Config-First Constants - -- Keep hyperparameters and architecture constants in configs with explicit names. -- Avoid undocumented inline numeric defaults in implementation modules. -- Structural literals are acceptable only for clear invariants (for example, rank checks). - -## Attribution Contract - -The bag logit is an attention-weighted sum of instance logits with no added -bias, so for every bag and class: - - sum_i A[i,c] * l[i,c] == L_c - -This is an invariant, not an approximation. Any change to bag pooling must -preserve it, and `niche_attribution_summary` reports `identity_residual` so a -violation surfaces rather than propagating into the niche summaries. diff --git a/docs/nested_cv.md b/docs/nested_cv.md deleted file mode 100644 index d89f770..0000000 --- a/docs/nested_cv.md +++ /dev/null @@ -1,88 +0,0 @@ -# Nested Cross-Validation (the preprint protocol) - -`grass-mil-nested-cv` runs the evaluation protocol of the preprint: fixed, -patient-grouped outer folds; stratified, patient-grouped inner folds for epoch -selection; a refit on the whole development set; one held-out prediction per -region and a mean per patient. `grass-mil-loocv` (leave-one-region-out) is -still available for very small cohorts, but nested CV is the protocol all -preprint numbers come from. - -## Inputs - -| Input | Form | -|---|---| -| processed graphs | any `data.processed_dir` built by `grass-mil-train` / `prepare_data` (whole regions, `data.sample_unit=full`, one graph per region) | -| labels | CSV with a region-id column, a label column (0/1 or category names) and, for patient grouping, a patient column: `data.graph_labels.{label_file,id_column,tasks,patient_column}` | -| outer folds | JSON list of lists of region ids: `nested_cv.folds_file` (patients never split across folds; build it once and version it) | -| pretrained encoder | `model.init_from_ckpt=.ckpt` (optional but recommended) | -| regime | `regime=mil` (default MLP head), `mil_frozen`, `linear_probehead`, `linear_ftall`, `linear_ftall_reg` | - -```bash -grass-mil-nested-cv regime=linear_ftall_reg \ - data.raw_manifest_path=data/raw/manifest.csv data.processed_dir=data/processed/cohort \ - data.sample_unit=full data.csv.categorical_label_columns=[cell_type] data.csv.molecular_columns=[SIZE] \ - data.graph_builder.kwargs.neighbor_cutoff_um=55 \ - data.graph_labels.label_file=data/raw/labels.csv data.graph_labels.id_column=REGION_ID \ - data.graph_labels.tasks=[STAS] data.graph_labels.patient_column=PATIENT \ - nested_cv.folds_file=data/raw/folds.json \ - model.init_from_ckpt=runs/pretrain_a3/checkpoints/last.ckpt \ - paths.output_dir=runs/nested_ftall_reg -``` - -## What happens per outer fold `i` - -1. Development set = every other fold, sorted by region id. -2. `StratifiedGroupKFold(nested_cv.inner_k, shuffle=True, random_state=nested_cv.inner_seed)` over the development regions, grouped by patient. A seed that leaves a fold without one class is rejected and the next seed tried (only matters for tiny cohorts). -3. For every inner fold: if the regime uses a probe head (`task.head_init_path: auto`, the `linear_*` regimes), `grass-mil-probe` is run on the inner training regions (cached under `nested_cv.probe.cache_dir`); then `grass-mil-train` with the regime, `callbacks=mil_selection` (checkpoint on region-level `val/auroc`, early stopping on `val/loss_region`, `trainer.min_epochs`), writing `metrics_per_epoch.jsonl`. The best epoch is `argmax val/auroc` (earliest on ties; `val/loss_region` is the fallback when AUROC is undefined). -4. `E = median` (`nested_cv.epoch_reduce`) of the inner best epochs. -5. Refit on all development regions for exactly `E` epochs (no early stopping, last weights; the probe is refitted on the development set), then test on fold `i` with every cell of every region (`test_sampler` = all roots). -6. Outputs: `outer_/result.json`, `outer_/refit/test_predictions.csv`, `outer_/oof_predictions.csv`. - -The run resumes: finished inner runs and outer folds are skipped (`nested_cv.resume=true`). - -## Outputs - -| File | Content | -|---|---| -| `nested_cv_summary.json` | `auc_mean`, `auc_sd`, `auc_folds`, per-fold `E`, inner epochs, run directories, pooled region AUROC, patient AUROC | -| `oof_predictions.csv` | one row per region: `outer`, `region_id`, `patient_id`, `label`, `logit_c`, `prob_class_c` | -| `oof_patient_predictions.csv` | mean of the region probabilities per patient | -| `probes/*.pt` | folded probe heads (`graph_head.net.0.{weight,bias}`) per split | - -Score the table with `grass-mil-scores input=/oof_predictions.csv` (bootstrap -CI, per-fold mean, operating points) and extract cells of the five refit models -with `grass-mil-extract extract.nested_cv_summary=/nested_cv_summary.json`. - -## Regimes - -| `regime=` | encoder | instance head | attention | bag | notes | -|---|---|---|---|---|---| -| `mil` (default) | fine-tuned (lr x1) | MLP (2 layers), trained | lr x2 | 256 roots | GRASS-MIL default | -| `mil_frozen` | frozen | MLP | lr x2 | 512 | encoder in eval mode, no gradients | -| `linear_probehead` | frozen | linear, pre-fitted by the probe, frozen | lr x2 | 512 | attention is the only trainable part | -| `linear_ftall` | fine-tuned (lr x1) | linear, pre-fitted, lr x0.01, wd 1e-3 | lr x2 | 256 | | -| `linear_ftall_reg` | fine-tuned (lr x0.3) | linear, pre-fitted, lr x0.01 | lr x2 | 256 | wd 1e-2, encoder dropout 0.2, edge/feature drop 0.05 | - -All regimes share AdamW (lr 1e-4), a per-step cosine schedule over the whole run -(`scheduler=cosine_mil_step`, `T_max: auto`), inverse-frequency class weights from -the training regions, gradient accumulation over 12 regions, gradient clipping -at 1, mixed precision, a fixed 1024-root validation subset per region and -all-cells testing. Every value is a Hydra key, e.g. `regime=mil -task.optimization.encoder_lr_scale=0.3 data.sampler.runtime.subsample_count=128 -trainer.max_epochs=40 model.encoder.num_layers=3 data.sampler.runtime.depth=3`. - -## Patient-level scores - -`oof_patient_predictions.csv` averages the region probabilities (the preprint -default). The merged-bag alternative, where all cells of a patient's regions -form one bag with a joint attention softmax, is available at prediction time: - -```bash -grass-mil-predict ckpt_path=/checkpoints/last.ckpt regime=linear_ftall_reg data=... \ - aggregation.mode=attention_weighted aggregation.bag_scope=patient \ - aggregation.patient_map.label_file=labels.csv aggregation.patient_map.id_column=REGION_ID \ - aggregation.patient_map.patient_column=PATIENT -``` - -`aggregation.mode=mean_probability aggregation.bag_scope=patient` reproduces the -mean-of-regions score from the same run. diff --git a/docs/niche_analysis.md b/docs/niche_analysis.md new file mode 100644 index 0000000..16b412d --- /dev/null +++ b/docs/niche_analysis.md @@ -0,0 +1,133 @@ +# Niche analysis + +Niches are recurrent cellular neighbourhoods that a trained model attends to. +The analysis runs on the per-cell outputs of a model and describes every niche +by what it contains, how it relates to the phenotype, how much the model relies +on it, and how its cells are arranged. + +## 1. Per-cell outputs + +```bash +grass-mil-extract experiment=my_cohort extract.nested_cv_summary=/nested_cv_summary.json +``` + +Every model of the cross-validation (`fold0`, `fold1`, ...) scores every cell of +every unit. For each cell the cache stores the encoder embedding, the attention +head's low-dimensional pre-tanh activation (the *attention space*), the +instance logits and the attention scores, with the cell's type, position and +id. The units each model held out are recorded. Explicit checkpoints can be +added with `+extract.models.=`. + +## 2. Niches of one model + +```bash +grass-mil-niches cells_dir=/cells/fold0 +``` + +**Cells.** By default the units the model did not see during training +(`units=held_out`; or `all`, or a list). + +**Margin.** Within a unit, a cell's contribution to class `c` is +`A[i, c] (l[i, c] - b_c)`, where `b_c` is the head bias; summed over the unit it +gives the unit's logit minus the bias. A cell's margin is its contribution to +the positive class minus its contribution to the other class(es), multiplied +by the unit's number of cells so that units of different size are comparable. + +**Clustering.** Cells are clustered in the attention space +(`clustering.space=attention`; `embedding` uses the encoder embedding, for which +`clustering.pca_components=10` is recommended): L2 normalisation (cosine +geometry), HDBSCAN (`min_cluster_size` 300, `min_samples` 1, excess of mass) on +a sample of at most `clustering.sample_size` cells, and assignment of the rest +by their 15 nearest neighbours. Cells HDBSCAN leaves unassigned form the +**Background**. Niches are numbered by their margin: `Niche 1` is the most +positive. + +### Outputs + +| File | Content | +|---|---| +| `niches.csv` | one row per niche (see below) | +| `unit_niche_shares.csv` | the share of every niche in every unit | +| `cells.csv` | every cell: unit, id, position, cell type, niche, margin | +| `descriptors/*.csv` | spatial descriptors of every niche (see below) | +| `niche_colocalisation.csv` | `log2` observed/expected contacts between niches | +| `report.md` | the niche table and the main descriptors | +| `figures/` | see below | +| `niche_model.joblib` | the fitted clustering (labels new cells of the same model) | + +`niches.csv` columns: + +* `share`: fraction of the analysed cells in the niche; `n_units`: units where + it exceeds `present_threshold` (2%); +* `association`: fraction of the niche's cells that come from positive units + (cell-weighted; compare with the base rate in the report). For survival + models a positive unit is one with an event; +* `margin` [`margin_lo`, `margin_hi`]: the mean over units of the niche's mean + cell margin in the unit, with a 95% bootstrap interval over units; + `sign_consistency`: the fraction of units where it has that sign; +* `composition`, `enriched` (cell types over-represented more than + `enrichment_threshold` times) and `frac_`. + +### Spatial descriptors + +Computed on every niche's subgraph within each unit (its cells and the edges +between them) and averaged over the units where the niche has at least +`descriptors.min_cells` cells: + +| File | Statistic | +|---|---| +| `moran.csv` | Moran's I of every cell type (are cells of the type clustered within the niche?), mean and s.e.m. over units, permutation p-value and BH q-value | +| `enrichment.csv` | neighbourhood-enrichment z-score of every cell-type pair (label permutations over the fixed graph) | +| `cross_l.csv` | Ripley's centred cross-L, `L(r*) - r*`, of every cell-type pair, at the radius `r*` where the pairs deviate most from randomness | +| `contacts.csv` | `log2` of a pair's share of the niche's contacts over its share of all contacts; masked below `descriptors.min_contacts` | + +### Figures + +* `niche_panel.png`: composition, association, margin with its interval and share per niche; +* `embedding.png`: the clustering space (first two principal components) coloured by niche, margin and empirical association; +* `unit_composition.png`: the niche shares of every unit, by class; +* `units/*.png`: the most positive and most negative units in the tissue (cell types, margins, niches); +* `moran.png`, `enrichment.png`, `cross_l.png`, `contacts.png`, `niche_colocalisation.png`. + +### Survival + +With `survival.file` (one row per patient with `time` and `event` columns), the +niche shares of each patient (the mean over the patient's units) are related +to survival: a Cox hazard ratio per standard deviation of the log share of +every niche (`survival_cox.csv`), and with `survival.score=["Niche 1", "-Niche 4"]` +Kaplan-Meier curves of the score's tertiles with a log-rank test +(`figures/survival.png`). + +## 3. Replication across models (optional) + +```bash +grass-mil-metaclusters cells_dir=/cells +``` + +Each model of the cross-validation has its own embedding, so its niches are +discovered separately. This step asks which niches recur: + +1. `cluster`: the niches of every model on the units it held out; +2. `match`: niches of different models are grouped into metaclusters by the + Hellinger distance of their compositions, with average linkage that never + merges two niches of the same model; the threshold maximises the number of + metaclusters found by at least `match.min_models` models. Outputs + `metaclusters.csv`, the metacluster shares of every unit and patient-level + tests; +3. `transfer`: every model labels every unit, giving training versus held-out + behaviour and the agreement between models (`transfer/`); +4. `basis`: one k-means partition of a single model's space, attributed by + every model (`basis/`). + +## 4. Model-free spatial statistics (optional) + +```bash +grass-mil-spatial-stats cells_dir=/cells/fold0 reference=[tumor] +``` + +Classical statistics of every unit compared between positive and other +patients (Mann-Whitney on patient means, BH-corrected): neighbourhood +enrichment of every cell-type pair, the shares of cellular neighbourhoods +(k-means of every cell's 10-nearest-neighbour composition, CLR-transformed), +contact fractions and cross-L at fixed radii between the `reference` group(s) +and the others. `group_map` merges fine cell types into groups first. diff --git a/docs/niche_discovery.md b/docs/niche_discovery.md deleted file mode 100644 index 250d978..0000000 --- a/docs/niche_discovery.md +++ /dev/null @@ -1,89 +0,0 @@ -# Niche Discovery And Cross-Fold Replication - -Everything downstream of a trained model works on its **cell cache**: for every -cell of every region, the encoder embedding of the ego-graph rooted at that -cell (`E`, fp16), the per-class instance logits (`LG`) and the raw per-class -attention logits (`AR`), plus the head bias, cell types, positions and the cell -graph once per region. - -## 1. Extract the cells - -```bash -# every refit model of a nested CV (held-out regions recorded per model) -grass-mil-extract regime=linear_ftall_reg data=... \ - extract.nested_cv_summary=runs/nested_ftall_reg/nested_cv_summary.json \ - paths.output_dir=runs/cells - -# one explicit checkpoint -grass-mil-extract regime=mil data=... +extract.models.final=runs/train/checkpoints/last.ckpt -``` - -Output: `/cells//index.json` + one `.npz` per region. -The model is composed from the same `regime`/`model` keys as the training run so -the checkpoint fits; ego-graph settings come from `data.sampler.runtime`. - -## 2. Niches of one model (`grass-mil-niches`) - -```bash -grass-mil-niches cells_dir=runs/cells/fold2 output_dir=runs/niches_fold2 \ - positive_name=STAS held_out_auroc=0.664 -``` - -Default settings (the preprint recipe; every key is overridable, no sweep): - -| key | default | meaning | -|---|---|---| -| `regions` | `held_out` | cells of the regions the model never trained on (`all`, a list or a file are accepted) | -| `space` | `gnn` | encoder embedding; `attention` = the 8-d pre-tanh attention activation (`attention_ckpt`) | -| `reduction` | PCA, 10 components | fitted on a random sample of `sample_size` cells | -| `l2_normalise` | true | cluster on L2-normalised PCs, i.e. the cosine geometry | -| `clustering` | HDBSCAN, `min_cluster_size` 300, `min_samples` 1, `eom` | any backend of the clustering registry (`kmeans` with `n_clusters`, `agglomerative`, `dbscan`, `optics`, `spectral`) | -| `assign_k` | 15 | k-NN assignment of unsampled cells on the PCs (backends with `predict`, e.g. k-means, assign directly) | -| `sample_size` | 100000 | cells used to fit reducer and clustering (`null` = all) | - -Per cell, the **attention-weighted margin** is -`relattn_i = N_r [A_{i,1}(l_{i,1} - b_1) - A_{i,0}(l_{i,0} - b_0)]` -with the attention normalised over the region's cells per class, so that the -margins of a region sum to `N_r` times its bias-free bag margin. - -Outputs: `labels.npy`, `cluster_model.pkl` (joblib; reducer + assigner, reusable -on other cells), `cells.csv` (region, cell index, cluster, cell type, relattn, -x, y), `clusters.csv`, `coloc_log2.npy`, `report.md`, `pca_triptych.png`, -`umap_triptych.png` (if `umap-learn` is installed), `clusters_overview.png`, -`colocalisation.png`. - -`clusters.csv` columns: share of cells, regions present (> `present_threshold`), -positive / negative regions present, composition (`comp_`), top types, -enriched types (> 1.5x the cohort frequency), empirical fraction of cells from -positive regions, AUROC of the per-region share, `relattn` with a bootstrap CI -over regions (`n_boot`), signed share of the region margin, sign consistency, -attention lift. The report adds the Spearman concordance between empirical -association and margin and the co-localisation `log2(obs/exp)` of cell-cell edges -between clusters. - -## 3. Cross-fold replication (`grass-mil-metaclusters`) - -A clustering of one model lives in that model's space. To ask which niches -replicate, cluster every fold model's held-out cells with the same fixed -settings and match the clusters across folds by composition: - -```bash -grass-mil-metaclusters cells_dir=runs/cells output_dir=runs/meta \ - labels.label_file=data/raw/labels.csv labels.label_column=STAS labels.patient_column=PATIENT \ - positive_name=STAS -``` - -Stages (`stage=all|cluster|match|transfer|basis`): - -1. **cluster**: per fold model, `grass-mil-niches` settings on its held-out cells -> `cluster/fold_clusters.csv`, `labels_.npy`, `cluster_model_.pkl`. -2. **match**: Hellinger distance between cluster compositions (`match.distance=aitchison` for CLR euclidean), constrained average linkage that never merges two clusters of the same fold, threshold from a scan maximising the number of metaclusters replicated in `>= match.min_folds` folds (ties -> fewer metaclusters) or fixed (`match.threshold=0.4`). Outputs `match/metaclusters.csv` (members, within distance, per-fold empirical association and margin, sign agreement, CI exclusion count), `metacluster_patient_shares.csv` (each region labelled once, by the model that held it out) and `metacluster_patient_tests.csv` (Mann-Whitney + BH, AUROC of the share), `report.md`, `metaclusters_overview.png`. -3. **transfer**: every model labels every region in its own space with its saved cluster model; `transfer/transfer_meta_by_model.csv` (share, association, margin on train vs held-out regions per metacluster and model), `transfer_model_ari.csv` (pairwise ARI on cells both models assign), consensus patient shares and tests. -4. **basis**: a shared label-free k-means basis (`basis.k`, standardised embeddings of `basis.space_model`, `basis.sample_size` cells) with the attribution of every fold model split into held-out / train / all (model, region) pairs: `basis/basis_attribution.csv`, `report_basis.md`, `basis_attribution.png`. - -## 4. The interpretability report (`grass-mil-report`) - -The HTML/PDF report of `docs/interpretability_suite.md` works on instance -tables exported by `grass-mil-predict`. Its clustering step accepts the same -sample / L2 / assignment recipe (`clustering.sample_size`, -`clustering.l2_normalise`, `clustering.assign_k`) and its margin attribution -checks the additive identity against the bag table (`data.bag_table`). diff --git a/docs/repository_map.md b/docs/repository_map.md deleted file mode 100644 index 2c864c4..0000000 --- a/docs/repository_map.md +++ /dev/null @@ -1,379 +0,0 @@ -# Repository Map - -Every tracked file (`git ls-files`) with a one-line purpose. Regenerate with -`python tools/update_repository_map.py` after adding files. - -## Root metadata and tooling - -- `.gitignore`: Ignore rules for generated artefacts, environments, logs, data and vignette outputs. -- `.pre-commit-config.yaml`: Code-quality hooks (ruff lint + format, file hygiene checks). -- `.project-root`: Root marker used by `rootutils.setup_root(...)`. -- `CHANGELOG.md`: Release notes. -- `CITATION.cff`: Machine-readable citation (preprint DOI 10.64898/2026.10.08.757722). -- `LICENSE`: MIT license. -- `Makefile`: Convenience targets: `test`, `format`, `clean`, `clean-logs`. -- `README.md`: Landing page: method summary, install, typical usage, workflow table, citation. -- `environment.yaml`: Conda environment alternative for local development. -- `pyproject.toml`: Package metadata, dependencies, extras, console entry points (`grass-mil-*`), pytest/ruff configuration. - -## CI and contribution files - -- `.github/PULL_REQUEST_TEMPLATE.md`: Pull-request template. -- `.github/codecov.yml`: Codecov configuration. -- `.github/dependabot.yml`: Dependabot update policy. -- `.github/release-drafter.yml`: Release drafter configuration. -- `.github/workflows/code-quality-main.yaml`: CI: pre-commit on every push to main. -- `.github/workflows/code-quality-pr.yaml`: CI: pre-commit on changed files of a pull request. -- `.github/workflows/release-drafter.yml`: CI: draft release notes from merged pull requests. -- `.github/workflows/test.yml`: CI: pytest on Linux and macOS, coverage, install-from-source smoke test. -- `.vscode/settings.json`: Editor settings. - -## Maintenance tools - -- `tools/update_repository_map.py`: Regenerates this map. - -## Package: entry points and shared modules - -- `src/grass_mil/__init__.py`: GRASS-MIL: graph-based representation and discovery of phenotype-associated. -- `src/grass_mil/_compat.py`: Cross-version compatibility shims applied at package import. -- `src/grass_mil/_paths.py`: Project-root resolution shared by the CLI entrypoints. -- `src/grass_mil/batching.py`: Splitting batched graphs back into their per-subgraph parts. -- `src/grass_mil/callbacks.py`: Lightweight run-artefact callbacks. -- `src/grass_mil/contracts.py`: Cross-cutting contracts. -- `src/grass_mil/eval.py`: Module. -- `src/grass_mil/extract.py`: ``grass-mil-extract``: per-cell cache of trained models (see inference/cell_cache.py). -- `src/grass_mil/loocv.py`: Module. -- `src/grass_mil/nested_cv.py`: ``grass-mil-nested-cv``: nested, patient-grouped cross-validation (preprint protocol). -- `src/grass_mil/probe.py`: ``grass-mil-probe``: fit the linear probe head on a frozen encoder. -- `src/grass_mil/train.py`: Module. - -## Package: data pipeline - -- `src/grass_mil/data/__init__.py`: Package marker / re-exports. -- `src/grass_mil/data/components/__init__.py`: Package marker / re-exports. -- `src/grass_mil/data/components/datasets.py`: Module. -- `src/grass_mil/data/components/ego_radius.py`: Euclidean radius cutoff for k-hop ego-graphs. -- `src/grass_mil/data/components/feature_reducers.py`: Module. -- `src/grass_mil/data/components/graph_builders.py`: Module. -- `src/grass_mil/data/components/loaders.py`: Module. -- `src/grass_mil/data/components/merge.py`: Merge several graph units into one disconnected graph. -- `src/grass_mil/data/components/patching.py`: Module. -- `src/grass_mil/data/components/precompute.py`: Module. -- `src/grass_mil/data/components/samplers.py`: Module. -- `src/grass_mil/data/components/seed_sampling.py`: Interior-only seed sampling for ego-graph extraction. -- `src/grass_mil/data/components/spatial_types.py`: Module. -- `src/grass_mil/data/components/splits.py`: Explicit splits, fold files and patient-grouped inner folds. -- `src/grass_mil/data/components/transforms.py`: Module. -- `src/grass_mil/data/spatial_omics_datamodule.py`: Module. - -## Package: models - -- `src/grass_mil/models/__init__.py`: Package marker / re-exports. -- `src/grass_mil/models/bgrl_module.py`: Module. -- `src/grass_mil/models/components/__init__.py`: Package marker / re-exports. -- `src/grass_mil/models/components/attention.py`: Module. -- `src/grass_mil/models/components/backbones.py`: Module. -- `src/grass_mil/models/components/embeddings.py`: Input embeddings for node features. -- `src/grass_mil/models/components/factory.py`: Module. -- `src/grass_mil/models/components/heads.py`: Module. -- `src/grass_mil/models/components/losses.py`: Module. -- `src/grass_mil/models/components/pooling.py`: Module. -- `src/grass_mil/models/components/ssl.py`: Module. -- `src/grass_mil/models/supervised_module.py`: Module. -- `src/grass_mil/models/training/__init__.py`: Package marker / re-exports. -- `src/grass_mil/models/training/augmentations.py`: Graph augmentations for self-supervised pretraining. -- `src/grass_mil/models/training/bagging.py`: Module. -- `src/grass_mil/models/training/builders.py`: Module. -- `src/grass_mil/models/training/checkpoint_init.py`: Module. -- `src/grass_mil/models/training/class_weights.py`: Class weights for the region cross-entropy, computed from the TRAINING split only. -- `src/grass_mil/models/training/loss_utils.py`: Module. -- `src/grass_mil/models/training/optimization.py`: Module. -- `src/grass_mil/models/training/probe.py`: Linear probe pre-fit on a frozen encoder (preprint "probe head"). -- `src/grass_mil/models/training/ssl_runtime.py`: Module. - -## Package: inference - -- `src/grass_mil/inference/__init__.py`: Package marker / re-exports. -- `src/grass_mil/inference/aggregation.py`: Module. -- `src/grass_mil/inference/cell_cache.py`: Per-cell cache of a trained GRASS-MIL model: the input of niche discovery. -- `src/grass_mil/inference/collectors.py`: Module. -- `src/grass_mil/inference/interpretability_export.py`: Module. -- `src/grass_mil/inference/io.py`: Module. -- `src/grass_mil/inference/metrics.py`: Module. -- `src/grass_mil/inference/model_loading.py`: Build a supervised module from config and load trained weights outside ``Trainer``. -- `src/grass_mil/inference/predict.py`: Module. -- `src/grass_mil/inference/schemas.py`: Module. - -## Package: evaluation statistics - -- `src/grass_mil/evaluation/__init__.py`: Evaluation statistics for held-out scores (AUROC CIs, DeLong, operating points). -- `src/grass_mil/evaluation/cli.py`: ``grass-mil-scores``: AUROC CIs, DeLong comparisons and operating points. -- `src/grass_mil/evaluation/operating_points.py`: Sensitivity / specificity / PPV / NPV at pre-specified operating rules. -- `src/grass_mil/evaluation/stats.py`: AUROC statistics for held-out (out-of-fold) scores. - -## Package: interpretability (niches, metaclusters, report) - -- `src/grass_mil/interpretability/__init__.py`: Package marker / re-exports. -- `src/grass_mil/interpretability/core/__init__.py`: Package marker / re-exports. -- `src/grass_mil/interpretability/core/agreement.py`: Agreement between partitions of the same instances. -- `src/grass_mil/interpretability/core/attribution.py`: Exact attribution of a bag prediction to its instances. -- `src/grass_mil/interpretability/core/biomarkers.py`: Module. -- `src/grass_mil/interpretability/core/clustering.py`: Module. -- `src/grass_mil/interpretability/core/data.py`: Module. -- `src/grass_mil/interpretability/core/niches.py`: Niche identity: naming, the background label, and composition ordering. -- `src/grass_mil/interpretability/core/reduction.py`: Module. -- `src/grass_mil/interpretability/core/transfer.py`: Module. -- `src/grass_mil/interpretability/metaclusters/__init__.py`: Cross-fold niche replication: per-fold clustering, matching into metaclusters,. -- `src/grass_mil/interpretability/metaclusters/basis.py`: Stage 4: a shared, label-free k-means basis with per-model attribution. -- `src/grass_mil/interpretability/metaclusters/match.py`: Stage 2: match fold-clusters across folds into metaclusters. -- `src/grass_mil/interpretability/metaclusters/records.py`: Stage 1: cluster every fold model's held-out cells in its own space. -- `src/grass_mil/interpretability/metaclusters/transfer.py`: Stage 3: transfer every fold model's clustering to ALL regions in its own space. -- `src/grass_mil/interpretability/metaclusters_cli.py`: ``grass-mil-metaclusters``: cross-fold niche replication (cluster / match / transfer / basis). -- `src/grass_mil/interpretability/niches/__init__.py`: Niche discovery on a trained GRASS-MIL model's per-cell cache. -- `src/grass_mil/interpretability/niches/colocalisation.py`: Niche-niche co-localisation on the cell graph: log2(observed / expected) edge counts. -- `src/grass_mil/interpretability/niches/figures.py`: Figures of the niche report (matplotlib; UMAP optional). -- `src/grass_mil/interpretability/niches/margins.py`: Per-cell attention-weighted margins and the attention feature space. -- `src/grass_mil/interpretability/niches/pipeline.py`: Niche discovery on one model's cell cache. -- `src/grass_mil/interpretability/niches/tables.py`: Per-niche summary table (the preprint's cluster table). -- `src/grass_mil/interpretability/niches_cli.py`: ``grass-mil-niches``: niche discovery on one model's cell cache. -- `src/grass_mil/interpretability/pipeline.py`: Module. -- `src/grass_mil/interpretability/plugins/__init__.py`: Package marker / re-exports. -- `src/grass_mil/interpretability/plugins/base.py`: Module. -- `src/grass_mil/interpretability/plugins/builtin.py`: Module. -- `src/grass_mil/interpretability/plugins/registry.py`: Module. -- `src/grass_mil/interpretability/report_cli.py`: Module. -- `src/grass_mil/interpretability/reporting/__init__.py`: Package marker / re-exports. -- `src/grass_mil/interpretability/reporting/html.py`: Module. -- `src/grass_mil/interpretability/reporting/pdf.py`: Module. -- `src/grass_mil/interpretability/reporting/plotly_builders.py`: Module. -- `src/grass_mil/interpretability/reporting/render.py`: Module. -- `src/grass_mil/interpretability/reporting/sections.py`: Module. -- `src/grass_mil/interpretability/reporting/snapshot.py`: Module. -- `src/grass_mil/interpretability/tier2/__init__.py`: Package marker / re-exports. -- `src/grass_mil/interpretability/tier2/autocorrelation.py`: Global spatial autocorrelation (Moran's I). -- `src/grass_mil/interpretability/tier2/cell_level.py`: Cell-level tier-2 analyses, run within each niche. -- `src/grass_mil/interpretability/tier2/filtration.py`: Module. -- `src/grass_mil/interpretability/tier2/neighborhood.py`: Module. -- `src/grass_mil/interpretability/tier2/ripley.py`: Ripley's K/L and the bivariate cross-K. -- `src/grass_mil/interpretability/tier2/tissue_graph.py`: Module. - -## Package: classical spatial statistics - -- `src/grass_mil/spatial_stats/__init__.py`: Classical spatial statistics on the cell graph, per region, compared between groups. -- `src/grass_mil/spatial_stats/cellular_neighbourhoods.py`: Cellular neighbourhoods: k-NN window composition clustered with k-means. -- `src/grass_mil/spatial_stats/cli.py`: ``grass-mil-spatial-stats``: classical spatial statistics on a cell cache. -- `src/grass_mil/spatial_stats/contacts.py`: Contact fractions on the cell graph. -- `src/grass_mil/spatial_stats/cross_k.py`: Ripley cross-K between a reference group and partner groups, normalised by pi r^2. -- `src/grass_mil/spatial_stats/nhood_enrichment.py`: Neighbourhood enrichment z-scores on a cell graph (label permutation null). -- `src/grass_mil/spatial_stats/run.py`: Run the classical spatial statistics over a cell cache and compare groups. -- `src/grass_mil/spatial_stats/tests.py`: Per-region -> per-patient summaries and two-group tests with BH correction. - -## Package: utilities - -- `src/grass_mil/utils/__init__.py`: Package marker / re-exports. -- `src/grass_mil/utils/instantiators.py`: Module. -- `src/grass_mil/utils/logging_utils.py`: Module. -- `src/grass_mil/utils/pylogger.py`: Module. -- `src/grass_mil/utils/rich_utils.py`: Module. -- `src/grass_mil/utils/utils.py`: Module. - -## Hydra configs - -- `src/grass_mil/configs/__init__.py`: Package marker / re-exports. -- `src/grass_mil/configs/augmentation/bgrl_paper.yaml`: Degree-importance-weighted node/edge dropping plus Gaussian noise on the. -- `src/grass_mil/configs/augmentation/uniform.yaml`: Uniform edge dropping and feature masking, independent of graph structure. -- `src/grass_mil/configs/callbacks/default.yaml`: Hydra config: `default`. -- `src/grass_mil/configs/callbacks/early_stopping.yaml`: https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.EarlyStopping.html. -- `src/grass_mil/configs/callbacks/mil_selection.yaml`: Model selection for MIL fine-tuning (preprint): keep the epoch with the best. -- `src/grass_mil/configs/callbacks/model_checkpoint.yaml`: https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html. -- `src/grass_mil/configs/callbacks/model_summary.yaml`: https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.RichModelSummary.html. -- `src/grass_mil/configs/callbacks/none.yaml`: Hydra config: `none`. -- `src/grass_mil/configs/callbacks/pretrain_checkpoint.yaml`: Pretraining: a checkpoint every N optimizer steps (preprint: 5000), all kept,. -- `src/grass_mil/configs/callbacks/refit.yaml`: Refit on all development folds for a fixed number of epochs: no early. -- `src/grass_mil/configs/callbacks/rich_progress_bar.yaml`: https://lightning.ai/docs/pytorch/latest/api/lightning.pytorch.callbacks.RichProgressBar.html. -- `src/grass_mil/configs/data/sampler/ego_graph_mil.yaml`: Ego-graph bags for MIL fine-tuning (preprint defaults). -- `src/grass_mil/configs/data/sampler/ego_graph_pretrain.yaml`: Ego-graph batches for BGRL pretraining (preprint "A3" defaults): roots drawn. -- `src/grass_mil/configs/data/sampler/identity.yaml`: Hydra config: `identity`. -- `src/grass_mil/configs/data/sampler/shadow_custom.yaml`: Hydra config: `shadow_custom`. -- `src/grass_mil/configs/data/sampler/shadow_native.yaml`: Hydra config: `shadow_native`. -- `src/grass_mil/configs/data/spatial_omics.yaml`: Hydra config: `spatial_omics`. -- `src/grass_mil/configs/debug/default.yaml`: default debugging setup, runs 1 full epoch. -- `src/grass_mil/configs/debug/fdr.yaml`: runs 1 train, 1 validation and 1 test step. -- `src/grass_mil/configs/debug/limit.yaml`: uses only 1% of the training data and 5% of validation/test data. -- `src/grass_mil/configs/debug/overfit.yaml`: overfits to 3 batches. -- `src/grass_mil/configs/debug/profiler.yaml`: runs with execution time profiling. -- `src/grass_mil/configs/eval.yaml`: Hydra config: `eval`. -- `src/grass_mil/configs/experiment/example_cohort.yaml`: Template for a cohort of whole regions (cores, ROIs, slides) given as CSV cell. -- `src/grass_mil/configs/experiment/grass_mil_pretrain_a3.yaml`: Same recipe as `regime=pretrain_a3`, for runs without a cohort experiment file: -- `src/grass_mil/configs/extract.yaml`: grass-mil-extract: per-cell cache (embeddings, instance logits, attention) of. -- `src/grass_mil/configs/extras/default.yaml`: disable python warnings if they annoy you. -- `src/grass_mil/configs/hydra/default.yaml`: https://hydra.cc/docs/configure_hydra/intro/. -- `src/grass_mil/configs/inference/aggregation/default.yaml`: Hydra config: `default`. -- `src/grass_mil/configs/inference/embeddings/default.yaml`: Hydra config: `default`. -- `src/grass_mil/configs/inference/interpretability/default.yaml`: Hydra config: `default`. -- `src/grass_mil/configs/inference/metrics/default.yaml`: Hydra config: `default`. -- `src/grass_mil/configs/inference/predict.yaml`: Hydra config: `predict`. -- `src/grass_mil/configs/interpretability/clustering/agglomerative.yaml`: Hydra config: `agglomerative`. -- `src/grass_mil/configs/interpretability/clustering/dbscan.yaml`: Hydra config: `dbscan`. -- `src/grass_mil/configs/interpretability/clustering/hdbscan.yaml`: Hydra config: `hdbscan`. -- `src/grass_mil/configs/interpretability/clustering/kmeans.yaml`: Hydra config: `kmeans`. -- `src/grass_mil/configs/interpretability/clustering/optics.yaml`: Hydra config: `optics`. -- `src/grass_mil/configs/interpretability/clustering/spectral.yaml`: Hydra config: `spectral`. -- `src/grass_mil/configs/interpretability/plugins/default.yaml`: Hydra config: `default`. -- `src/grass_mil/configs/interpretability/plugins/full.yaml`: Every builtin analysis enabled. -- `src/grass_mil/configs/interpretability/reduction/pca.yaml`: Hydra config: `pca`. -- `src/grass_mil/configs/interpretability/reduction/tsne.yaml`: Hydra config: `tsne`. -- `src/grass_mil/configs/interpretability/reduction/umap.yaml`: Hydra config: `umap`. -- `src/grass_mil/configs/interpretability/report.yaml`: Hydra config: `report`. -- `src/grass_mil/configs/local/.gitkeep`: Placeholder for untracked machine-local configs. -- `src/grass_mil/configs/logger/aim.yaml`: https://aimstack.io/. -- `src/grass_mil/configs/logger/comet.yaml`: https://www.comet.ml. -- `src/grass_mil/configs/logger/csv.yaml`: csv logger built in lightning. -- `src/grass_mil/configs/logger/many_loggers.yaml`: train with many loggers at once. -- `src/grass_mil/configs/logger/mlflow.yaml`: https://mlflow.org. -- `src/grass_mil/configs/logger/neptune.yaml`: https://neptune.ai. -- `src/grass_mil/configs/logger/tensorboard.yaml`: https://www.tensorflow.org/tensorboard/. -- `src/grass_mil/configs/logger/wandb.yaml`: https://wandb.ai. -- `src/grass_mil/configs/loocv.yaml`: Hydra config: `loocv`. -- `src/grass_mil/configs/metaclusters.yaml`: grass-mil-metaclusters: cross-fold niche replication. -- `src/grass_mil/configs/model/attention/gated.yaml`: Plain gated attention (Ilse et al. 2018), one attention channel per class. -- `src/grass_mil/configs/model/attention/gated_projected.yaml`: Projected gated attention (GRASS-MIL default): Linear + LeakyReLU projection,. -- `src/grass_mil/configs/model/bgrl_a3.yaml`: Hydra config: `bgrl_a3`. -- `src/grass_mil/configs/model/bgrl_module.yaml`: Hydra config: `bgrl_module`. -- `src/grass_mil/configs/model/encoder/gat.yaml`: Hydra config: `gat`. -- `src/grass_mil/configs/model/encoder/gcn.yaml`: Hydra config: `gcn`. -- `src/grass_mil/configs/model/encoder/gin.yaml`: Small GIN encoder without edge features: a lightweight default for smoke runs. -- `src/grass_mil/configs/model/encoder/gine.yaml`: Hydra config: `gine`. -- `src/grass_mil/configs/model/encoder/gine_length.yaml`: GRASS-MIL encoder (preprint defaults): 4 GINE layers of width 256 conditioned. -- `src/grass_mil/configs/model/encoder/graphsage.yaml`: Hydra config: `graphsage`. -- `src/grass_mil/configs/model/grass_mil.yaml`: Hydra config: `grass_mil`. -- `src/grass_mil/configs/model/supervised_module.yaml`: Hydra config: `supervised_module`. -- `src/grass_mil/configs/nested_cv.yaml`: grass-mil-nested-cv: the preprint's nested, patient-grouped cross-validation. -- `src/grass_mil/configs/niches.yaml`: grass-mil-niches: niche discovery on a trained model's cell cache. -- `src/grass_mil/configs/optim/adamw.yaml`: Hydra config: `adamw`. -- `src/grass_mil/configs/paths/default.yaml`: path to root directory. -- `src/grass_mil/configs/probe.yaml`: grass-mil-probe: fit the linear probe head on a frozen encoder. -- `src/grass_mil/configs/regime/linear_ftall.yaml`: Preprint "ftALL": encoder fine-tuned end to end (lr x1.0), linear head. -- `src/grass_mil/configs/regime/linear_ftall_reg.yaml`: Preprint "ftALL, lr x0.3 + regularisation" (the STAS model of record): -- `src/grass_mil/configs/regime/linear_probehead.yaml`: Preprint "probehead": frozen encoder, LINEAR head pre-fitted by grass-mil-probe. -- `src/grass_mil/configs/regime/mil.yaml`: GRASS-MIL default regime: BGRL-initialised encoder fine-tuned end to end,. -- `src/grass_mil/configs/regime/mil_frozen.yaml`: Frozen encoder (learning-rate scale 0, kept in eval mode); MLP head and. -- `src/grass_mil/configs/regime/pretrain_a3.yaml`: GRASS-MIL encoder pretraining (preprint "A3"): BGRL on the ego-graphs of. -- `src/grass_mil/configs/scheduler/cosine_epoch.yaml`: Hydra config: `cosine_epoch`. -- `src/grass_mil/configs/scheduler/cosine_mil_step.yaml`: Cosine annealing per optimizer step over the whole run (preprint): with. -- `src/grass_mil/configs/scheduler/cosine_step.yaml`: Hydra config: `cosine_step`. -- `src/grass_mil/configs/scores.yaml`: grass-mil-scores: evaluation statistics for held-out scores. -- `src/grass_mil/configs/spatial_stats.yaml`: grass-mil-spatial-stats: classical spatial statistics per region, compared. -- `src/grass_mil/configs/task/finetune_mean.yaml`: Hydra config: `finetune_mean`. -- `src/grass_mil/configs/task/finetune_mil.yaml`: Hydra config: `finetune_mil`. -- `src/grass_mil/configs/task/finetune_survival.yaml`: Cox survival with MIL attention pooling. The head emits a single scalar. -- `src/grass_mil/configs/task/pretrain_bgrl.yaml`: Hydra config: `pretrain_bgrl`. -- `src/grass_mil/configs/task/pretrain_bgrl_a3.yaml`: BGRL pretraining, preprint "A3" recipe: AdamW 5e-5 / weight decay 1e-5,. -- `src/grass_mil/configs/train.yaml`: specify here default configuration. -- `src/grass_mil/configs/trainer/cpu.yaml`: Hydra config: `cpu`. -- `src/grass_mil/configs/trainer/ddp.yaml`: Hydra config: `ddp`. -- `src/grass_mil/configs/trainer/ddp_sim.yaml`: Hydra config: `ddp_sim`. -- `src/grass_mil/configs/trainer/default.yaml`: Hydra config: `default`. -- `src/grass_mil/configs/trainer/gpu.yaml`: Hydra config: `gpu`. -- `src/grass_mil/configs/trainer/gpu_amp.yaml`: Hydra config: `gpu_amp`. -- `src/grass_mil/configs/trainer/mps.yaml`: Hydra config: `mps`. - -## Documentation - -- `docs/capabilities_and_modes.md`: Capabilities And Modes. -- `docs/configuration_guide.md`: Configuration Guide. -- `docs/dataset_preparation.md`: Dataset Preparation Guide. -- `docs/evaluation_stats.md`: Evaluation Statistics (`grass-mil-scores`). -- `docs/getting_started.md`: Getting Started: Typical Usage On Your Own Cohort. -- `docs/hyperparameter_reference.md`: Hyperparameter Reference. -- `docs/inference_utility_suite.md`: Inference Utility Suite. -- `docs/interpretability_extension_guide.md`: Interpretability Extension Guide. -- `docs/interpretability_suite.md`: Interpretability Suite. -- `docs/method_correspondence.md`: Method Correspondence With The Preprint. -- `docs/model_component_contracts.md`: Model Component Contracts. -- `docs/nested_cv.md`: Nested Cross-Validation (the preprint protocol). -- `docs/niche_discovery.md`: Niche Discovery And Cross-Fold Replication. -- `docs/repository_map.md`: Repository Map. -- `docs/spatial_statistics.md`: Classical Spatial Statistics (`grass-mil-spatial-stats`). -- `docs/workflows_train_infer.md`: Training And Inference Workflows. - -## Vignettes - -- `vignettes/01_anndata_to_spatial_graphs.ipynb`: Vignette 1: AnnData -> manifest -> cellular graphs. -- `vignettes/02_pretraining_and_niches.ipynb`: Vignette 2: BGRL pretraining, fine-tuning, instance export, interpretability report. -- `vignettes/03_nested_cv_niches_and_metaclusters.ipynb`: Vignette 3: nested CV, cell extraction, niches, metaclusters, scores, spatial statistics. -- `vignettes/README.md`: Index of the worked examples on the public MIBI-TOF dataset. - -## Tests - -- `tests/__init__.py`: Package marker / re-exports. -- `tests/configs/test_component_instantiation.py`: Compose the model/encoder and model/attention groups through Hydra and build them. -- `tests/configs/test_task_preset_composition.py`: Compose task presets through Hydra the way the CLI does. -- `tests/conftest.py`: Shared pytest fixtures (Hydra configs on a synthetic datamodule). -- `tests/contracts/test_model_component_contracts.py`: Tests for model component contracts. -- `tests/data/test_datamodule_sampler_exposure.py`: Tests for datamodule sampler exposure. -- `tests/data/test_loader_flags.py`: Tests for loader flags. -- `tests/data/test_precompute_manifest_regions.py`: Tests for precompute manifest regions. -- `tests/data/test_precompute_reducer_modes.py`: Tests for precompute reducer modes. -- `tests/data/test_samplers.py`: Tests for samplers. -- `tests/data/test_seed_sampling.py`: Contracts for interior seed restriction and the ego-radius cutoff. -- `tests/eval/test_aggregation.py`: Tests for aggregation. -- `tests/eval/test_collectors.py`: Tests for collectors. -- `tests/eval/test_embeddings.py`: Tests for embeddings. -- `tests/eval/test_interpretability_export.py`: Tests for interpretability export. -- `tests/eval/test_io.py`: Tests for io. -- `tests/eval/test_metrics.py`: Tests for metrics. -- `tests/eval/test_predict_cli.py`: Tests for predict cli. -- `tests/helpers/__init__.py`: Package marker / re-exports. -- `tests/helpers/config_paths.py`: Locate the packaged Hydra config tree. -- `tests/helpers/package_available.py`: Tests for package available. -- `tests/helpers/run_if.py`: Adapted from: -- `tests/helpers/run_sh_command.py`: Tests for run sh command. -- `tests/helpers/synthetic_cohort.py`: Synthetic cohort on disk (regions, labels, folds) with a planted niche, used by the phase tests. -- `tests/helpers/synthetic_datamodule.py`: In-memory bag datamodule for fast training tests. -- `tests/interpretability/test_attribution.py`: Contracts for the exact additive attribution. -- `tests/interpretability/test_cell_level.py`: Cell-level tier-2 analyses. -- `tests/interpretability/test_core.py`: Tests for core. -- `tests/interpretability/test_fidelity_phase6.py`: Phase 6: niche discovery (cosine HDBSCAN, relattn margins, tables, co-localisation, pipeline). -- `tests/interpretability/test_fidelity_phase7.py`: Phase 7: cross-fold metaclusters, transfer and shared basis. -- `tests/interpretability/test_niches.py`: Niche naming, the Background label, and composition ordering. -- `tests/interpretability/test_pipeline_and_report.py`: Tests for pipeline and report. -- `tests/interpretability/test_plugins_and_tier2.py`: Tests for plugins and tier2. -- `tests/interpretability/test_spatial_statistics.py`: Contracts for Moran's I, Ripley's cross-L, and cross-space agreement. -- `tests/interpretability/test_tissue_graph.py`: Tests for tissue graph. -- `tests/interpretability/test_transfer.py`: Tests for transfer. -- `tests/models/test_attention.py`: Tests for attention. -- `tests/models/test_augmentations.py`: Contracts for the BGRL view generators. -- `tests/models/test_backbones.py`: Tests for backbones. -- `tests/models/test_bagging.py`: Bag-pooling contracts, including the attribution identity. -- `tests/models/test_fidelity_phase1.py`: Phase 1 fidelity: graph edge features, encoder epsilon/init, attention init and space. -- `tests/models/test_fidelity_phase2.py`: Phase 2 fidelity: region bags, root sampling, parameter groups, schedules, class weights,. -- `tests/models/test_fidelity_phase3.py`: Phase 3: explicit splits, probe pre-fit, nested cross-validation. -- `tests/models/test_fidelity_phase4.py`: Phase 4: BGRL "A3" pretraining fidelity (merged-unit sampling, momentum schedule,. -- `tests/models/test_fidelity_phase5.py`: Phase 5: cell cache extraction, patient aggregation, evaluation statistics. -- `tests/models/test_losses.py`: Tests for losses. -- `tests/models/test_node_embeddings.py`: Contracts for the encoder's input embeddings. -- `tests/models/test_optimization.py`: Tests for optimization. -- `tests/models/test_runtime.py`: Tests for runtime. -- `tests/models/test_ssl.py`: Tests for ssl. -- `tests/spatial_stats/__init__.py`: Package marker / re-exports. -- `tests/spatial_stats/test_fidelity_phase8.py`: Phase 8: classical spatial statistics. -- `tests/test_configs.py`: Tests for configs. -- `tests/test_e2e_synthetic.py`: End-to-end smoke test of the public analysis chain on the synthetic cohort. -- `tests/test_eval.py`: Tests for eval. -- `tests/test_loocv.py`: Tests for loocv. -- `tests/test_packaging.py`: Guards on the installable-package contract. -- `tests/test_spatial_omics_datamodule.py`: Tests for spatial omics datamodule. -- `tests/test_train.py`: Tests for train. -- `tests/test_train_regimes.py`: Tests for train regimes. - -## Other - -- `data/.gitkeep`: Tracked file. -- `logs/.gitkeep`: Tracked file. - -## Generated at runtime (untracked) - -- `logs/`: Hydra run directories (checkpoints, `metrics_per_epoch.jsonl`, predictions, caches, reports). -- `data/`: raw manifests and processed graphs (`processed_index.json`, `metadata.json`, `graphs/`). -- `vignettes/vignette_work/`: data, graphs and runs produced by the vignettes. diff --git a/docs/spatial_statistics.md b/docs/spatial_statistics.md deleted file mode 100644 index 3b9a706..0000000 --- a/docs/spatial_statistics.md +++ /dev/null @@ -1,25 +0,0 @@ -# Classical Spatial Statistics (`grass-mil-spatial-stats`) - -Model-free statistics on the cell graph of every region, averaged per patient -and compared between classes with Mann-Whitney U and Benjamini-Hochberg -correction within each analysis. They run on any cell cache written by -`grass-mil-extract` (only cell types, positions and the graph are used). - -```bash -grass-mil-spatial-stats cells_dir=runs/cells/fold0 output_dir=runs/spatial \ - labels.label_file=data/raw/labels.csv labels.label_column=STAS labels.patient_column=PATIENT \ - reference=[tumor] partners=[myeloid,CD8,Treg] infiltrating=[myeloid,CD8,Treg] \ - group_map='{normal: tumor, hypoxic: tumor, cd8: CD8, ki67_cd8: CD8}' -``` - -| analysis | per region | patient level | -|---|---|---| -| `nhood` | neighbourhood enrichment z-scores of every group pair on the cell graph, `nhood.n_perms` label permutations (`nhood.backend=squidpy` delegates to `squidpy.gr.nhood_enrichment` on the same graph) | mean z per pair, test | -| `cn` | `cn.window_k`-nearest-neighbour window composition, `MiniBatchKMeans(cn.n_cn)` fitted on `cn.sample_size` windows of all regions, share of each cellular neighbourhood | mean share, CLR-transformed, test; `cn_centres.csv` gives the CN compositions | -| `contacts` | fraction of `reference` cells touching >= 1 cell of each partner group, homotypic contact, fraction of each `infiltrating` group whose neighbours are mostly reference cells | mean, test | -| `crossk` | Ripley cross-K reference -> partner at `crossk.radii`, normalised by pi r^2 (1 = complete spatial randomness; self pairs removed for reference -> reference) | mean, test | - -Outputs: `*_per_region.csv`, `nhood_tests.csv`, `cn_tests.csv`, `contact_tests.csv`, -`crossk_tests.csv`, `report_spatial.md`, `spatial_stats_summary.json`. -`group_map` collapses fine cell types into coarse groups before everything -(identity by default). Distances are in the cache's coordinate units. diff --git a/docs/training_and_evaluation.md b/docs/training_and_evaluation.md new file mode 100644 index 0000000..cc98456 --- /dev/null +++ b/docs/training_and_evaluation.md @@ -0,0 +1,116 @@ +# Training and evaluation + +## The model + +Every cell is represented by its *k*-hop ego-graph in the cell graph +(`data.sampler.depth`, default 4). The encoder embeds each ego-graph: a learned +cell-type embedding summed with a projection of the continuous features, then +`model.encoder.num_layers` GINE layers that condition every message on the +edge length, and max pooling. A unit is a bag of the ego-graphs of +`data.sampler.roots_per_unit` cells (256 by default, redrawn every epoch). + +For each instance `i`, the instance head gives one logit `l[i, c]` per class +and the projected gated attention gives one score per class; the attention is +a softmax over the bag's instances `A[i, c]`, and the bag logit is + +``` +L_c = sum_i A[i, c] * l[i, c] +``` + +so every unit prediction is an exact sum of per-cell contributions, which the +niche analysis uses. + +## Pretraining + +```bash +grass-mil-train experiment=my_cohort regime=pretrain trainer=gpu +``` + +BGRL on the ego-graphs of every unit, label-free: two augmented views +(importance-weighted node and edge dropping, Gaussian noise on +`pretrain.augmentation.noise_features`), an online encoder with a predictor and +an exponential-moving-average target encoder. The run writes +`checkpoints/last.ckpt` and a checkpoint every 5000 steps +(`pretrain.total_steps`, default 200 000). + +## Fine-tuning regimes + +A regime is the training procedure; select it with `regime=`. + +| Regime | Encoder | Instance head | Trained | +|---|---|---|---| +| `finetune` (default) | fine-tuned (lr x1) | linear, initialised by a probe | encoder, head (lr x0.01), attention (lr x2) | +| `finetune_reg` | fine-tuned (lr x0.3), dropout 0.2 | linear, probe | as `finetune`, with weight decay 1e-2, edge drop and node masking 0.05 | +| `attention_only` | frozen | linear, probe, frozen | attention only (512 roots per unit) | +| `mlp_head` | fine-tuned | two-layer MLP, random | everything | +| `mlp_head_frozen` | frozen | two-layer MLP, random | head and attention (512 roots per unit) | + +The probe: before training, the frozen encoder embeds +`finetune.probe.roots_per_unit` cells of every training unit, each ego-graph +inherits its unit's label, and a standardised logistic regression (ridge for +regression) is folded into the linear head. With uniform attention, the model +then starts from the mean-pooled probe. Survival models start from a random +head. + +Every regime trains one bag per step of the dataloader and takes an optimizer +step every `finetune.hyperbatch_size` units (12), with gradient-norm clipping +(`finetune.grad_clip_norm`), a cosine learning-rate schedule over the whole run +and AdamW (`optim.lr` 1e-4). Validation units are scored on a fixed set of +`data.val_roots_per_unit` cells; the checkpoint with the best validation metric +is kept and training stops early after `task.selection.patience` epochs without +improvement (after `trainer.min_epochs`). + +```bash +grass-mil-train experiment=my_cohort model.init_from_ckpt=/checkpoints/last.ckpt +``` + +A single run splits the units at random (`data.split.ratios`, grouped by +patient and stratified by class) unless `data.split.train/val/test` list ids, +and writes `test_predictions.csv` for the test units. + +## Tasks + +| `task=` | Targets | Loss | Selection | +|---|---|---|---| +| `classification` (default) | one class column, `task.num_classes` classes | cross-entropy, classes weighted by inverse training frequency | `val/auroc` | +| `regression` | one or more numeric columns | mean squared error (`task.loss=huber` for Huber) | `val/rmse` | +| `survival` | `[time, event]` | Cox partial likelihood over the units of each optimizer step | `val/c_index` | + +## Nested cross-validation + +```bash +grass-mil-nested-cv experiment=my_cohort model.init_from_ckpt=/checkpoints/last.ckpt +``` + +For every outer fold of `nested_cv.folds_file`: + +1. the other folds form the development set, split into `nested_cv.inner_k` + (3) inner folds that keep patients together and every class in every fold; +2. every inner run trains on its training part and records the epoch with the + best `task.selection.metric` on its validation part; +3. the model is refitted on the whole development set for the median of those + epochs, without validation, and scored once on the held-out fold, on every + cell. + +| Output | Content | +|---|---| +| `nested_cv_summary.json` | metric of every outer fold, mean and standard deviation, pooled out-of-fold metrics, patient-level metrics | +| `oof_predictions.csv` | one row per unit, scored by the model that held it out | +| `oof_patient_predictions.csv` | the mean of each patient's unit scores | +| `outer_/refit/checkpoints/last.ckpt` | the model of outer fold `i` | + +Runs that finished are reused when the command is repeated +(`nested_cv.resume`); `nested_cv.outer_folds=[2]` runs a single fold. + +## Prediction + +```bash +grass-mil-predict experiment=new_cohort ckpt_path=/checkpoints/last.ckpt +``` + +The model and its ego-graph settings are rebuilt from the checkpoint; `data` +only needs to describe the cohort (its cell-type vocabulary must match). Every +unit is scored on all of its cells (`predictions.csv`). With a patient column, +patients are scored as the mean of their units (`predict.patient_aggregation=mean`) +or as one bag holding every cell of their units (`merged_bag`). With labels, +`metrics.json` reports unit- and patient-level metrics. diff --git a/docs/workflows_train_infer.md b/docs/workflows_train_infer.md deleted file mode 100644 index 7f90c65..0000000 --- a/docs/workflows_train_infer.md +++ /dev/null @@ -1,692 +0,0 @@ -# Training And Inference Workflows - -This guide provides flow-through instructions for all major `grass-mil` workflows. - -> New here? The GRASS-MIL path (pretraining, regimes, nested CV, niches) is -> walked through in [getting_started.md](getting_started.md), and the -> configuration system in [configuration_guide.md](configuration_guide.md). -> This page covers the generic training and inference machinery underneath. - -Interpretability workflows are documented in: - -- `docs/interpretability_suite.md` -- `docs/interpretability_extension_guide.md` - -## 0) Environment And Dependency Prerequisites - -From the repository root: - -```bash -python -m venv .venv -source .venv/bin/activate -pip install -e ".[interpretability,report,spatial-stats,io]" -export PROJECT_ROOT=$(pwd) -``` - -Dependency notes by data type: - -- `csv` / `tsv`: requires `pandas`, `numpy`. -- `h5ad`: requires `anndata`. -- `sce` / `rds`: requires `rpy2` plus an R runtime and Bioconductor-compatible SCE objects. -- Polygon masking/patching: requires `shapely`. -- Graph building (`delaunay`): requires `scipy`. -- Graph runtime/samplers: requires `torch-geometric`; `shadow_custom` transform path also needs `torch-sparse`. - -## 1) Running A Default Pipeline With New Data - -### Step 1: Prepare The Manifest - -Default manifest path is `data/raw/manifest.csv` and default column names come from `configs/data/spatial_omics.yaml`. - -Required columns: - -- `sample_id` -- `input_path` -- `input_type` (`csv`, `tsv`, `h5ad`, `sce`, `rds`) - -Optional columns: - -- `region_id` (used as default bag key for supervised training) -- `polygons_path` (polygon file for filtering/patch segmentation) - -Example: - -```csv -sample_id,input_path,input_type,region_id,polygons_path -s1,data/raw/s1.csv,csv,r1, -s2,data/raw/s2.h5ad,h5ad,r2,data/raw/s2_polygons.geojson -``` - -Path resolution behavior: - -- Absolute paths are used as-is. -- Relative paths are resolved relative to the manifest directory. - -### Step 2: Configure Input-Specific Data Parsing - -All options below live under `data.*` in Hydra. - -Spatial unit expectation: - -- Normalize spatial coordinates to micrometers during preprocessing. -- `data.coord_scale_um` is the preprocessing conversion factor to micrometers. -- `data.coord_scale_um=1.0` means input coordinates are already in micrometers. -- This ensures downstream `spatial_table.distance` values are in micrometers and compatible with interpretability defaults (including filtration thresholds). -- If `data.coord_scale_um` changes, regenerate processed artifacts (`data.force_precompute=true` or use a new `data.processed_dir`). - -#### CSV/TSV inputs - -Key options (`data.csv.*`): - -- `sep` (default `,`) -- `coord_columns` (default `['x','y']`) -- `cell_id_column` -- `categorical_label_columns` -- `molecular_columns` (`null` means infer by excluding coords/cell id/categorical labels) - -Example override: - -```bash -grass-mil-train \ - data.csv.coord_columns="[x_um,y_um]" \ - data.csv.cell_id_column=cell_id \ - data.csv.categorical_label_columns="[cell_type]" -``` - -#### H5AD inputs - -Key options (`data.h5ad.*`): - -- `coord_source`: `obsm` or `obs` -- `coord_key` (for `obsm`, default `spatial`) -- `coord_columns` (for `obs`) -- `cell_id_column` -- `categorical_label_columns` -- `molecular_layer` -- `molecular_features` - -#### SCE/RDS inputs - -Key options (`data.sce.*`): - -- `assay_name` -- `coord_source`: `colData` or `reducedDims` -- `coord_key` (required for `reducedDims`) -- `coord_columns` (required for `colData`) -- `cell_id_column` -- `categorical_label_columns` -- `transpose_assay` -- `molecular_features` - -### Step 3: (Optional) Graph Labels And Polygon Files - -#### Polygon files - -Supported polygon formats: - -- `.pkl` / `.pickle` -- `.wkt` / `.txt` (line-delimited WKT) -- `.json` / `.geojson` - -`polygons_path` behavior: - -- Coordinates are first filtered by polygon union. -- Patch extraction uses polygons when `sample_unit != full`; otherwise tiling/full logic applies. - -#### Graph label file - -Configure `data.graph_labels.*`: - -- `label_file`: CSV path with labels -- `id_column`: column used as lookup key -- `tasks`: task columns (or all columns if `null`) -- `scope`: `region` | `patch` | `sample` - -When enabled, labels are attached as `graph_y` on each graph unit. - -### Step 4: Precompute And Split Control - -Precompute is triggered from `SpatialOmicsDataModule.prepare_data()` during training/eval startup. - -Primary controls: - -- `data.force_precompute`: force rebuild even if index exists -- `data.sample_unit`: `tile` or `full` -- `data.tiling.*`: tile geometry and tile min-cell threshold -- `data.min_cells`: global patch retention threshold -- `data.split.train_val_test_split`: must sum to `1.0` -- `data.split.split_by`: `sample`, `region`, or `patch` -- `data.split.loocv.*`: optional fold-wise holdout overrides -- `data.reducer_scope`: `sample` or `dataset` -- `data.feature_reducer.fit_mode`: `train_only` or `global` (only relevant for dataset scope) - -Important split persistence behavior: - -- Train/val/test labels are assigned during precompute and persisted in - `processed_index.json`. -- The datamodule consumes persisted split labels and does not recompute them at - runtime. -- If you change any `data.split.*` value, you must regenerate processed - artifacts (`data.force_precompute=true` or a new `data.processed_dir`), - otherwise old persisted splits are reused. - -Typical fresh-data command: - -```bash -grass-mil-train \ - task=finetune_mean \ - model=supervised_module \ - data=spatial_omics \ - data.raw_manifest_path=/absolute/path/to/manifest.csv \ - data.processed_dir=/absolute/path/to/processed/spatial_omics \ - data.force_precompute=true -``` - -### Step 4b: Run LOOCV (Single Fold Or Full Loop) - -`src/grass_mil/loocv.py` reuses the train stack and runs fold-specific jobs by overriding -`data.split.loocv.*`. - -Single-fold by explicit fold id: - -```bash -grass-mil-loocv \ - loocv.mode=single \ - data.split.loocv.enabled=true \ - data.split.loocv.fold_unit=region \ - data.split.loocv.holdout_id="sampleA::region3" \ - data.split.loocv.validation_strategy=heldout_fold_items \ - data.split.loocv.val_ratio=0.2 -``` - -Single-fold by fold index (useful for array jobs): - -```bash -grass-mil-loocv \ - loocv.mode=single \ - loocv.fold_index=0 \ - data.split.loocv.enabled=true \ - data.split.loocv.fold_unit=sample -``` - -If both `data.split.loocv.holdout_id` and `loocv.fold_index` are set in single -mode, `loocv.fold_index` takes priority and a warning is emitted. - -Full loop across all discovered folds: - -```bash -grass-mil-loocv \ - loocv.mode=all \ - data.split.loocv.enabled=true \ - data.split.loocv.fold_unit=region \ - data.split.loocv.validation_strategy=patches_from_train_items \ - data.split.loocv.val_ratio=0.15 -``` - -Validation strategy options: - -- `heldout_fold_items`: validation is held out from remaining fold items. -- `patches_from_train_items`: validation patches are sampled from training fold items (weaker independence, useful in low-data settings). - -Safety note: - -- For `data.reducer_scope=dataset` and `data.feature_reducer.fit_mode=train_only`, LOOCV requires fold-specific precompute (`data.force_precompute=true` or per-fold `processed_dir`). -- To prevent stale split reuse across folds, at least one of `loocv.force_precompute_per_fold` or `loocv.per_fold_processed_dir` must be `true`. - -### Step 5: Run Default Training And Evaluation - -Default training (mean aggregation): - -```bash -grass-mil-train task=finetune_mean model=supervised_module data=spatial_omics -``` - -Evaluation of a selected checkpoint: - -```bash -grass-mil-eval \ - ckpt_path=/absolute/path/to/checkpoint.ckpt \ - task=finetune_mean \ - model=supervised_module \ - data=spatial_omics -``` - -### Step 6: Inspect Output Artifacts - -Data artifacts (`data.processed_dir`): - -- `processed_index.json`: graph entry index, label maps, reducer state, split metadata -- `metadata.json`: summary of feature dimensions, labels, and split/reducer metadata -- `graphs//*.pt`: serialized PyG `Data` units - -Training artifacts (`paths.output_dir`): - -- `train.log` (Hydra job logging) -- `checkpoints/last.ckpt` -- `checkpoints/epoch_*.ckpt` (if checkpoint callback enabled) - -## 2) Pretraining With SSL, Then Using It In Fine-Tuning - -### Step 1: Run SSL Pretraining (BGRL) - -The GRASS-MIL recipe ("A3" in the preprint) is one experiment preset: - -```bash -grass-mil-train experiment=my_cohort regime=pretrain_a3 paths.output_dir=runs/pretrain_a3 trainer=gpu -# or, without a cohort file: grass-mil-train experiment=grass_mil_pretrain_a3 data.raw_manifest_path=... data.processed_dir=... -``` - -It composes `model=bgrl_a3`, `task=pretrain_bgrl_a3`, `callbacks=pretrain_checkpoint` -and `data/sampler=ego_graph_pretrain` (see -[`hyperparameter_reference.md`](hyperparameter_reference.md) section 3.4); every -value is overridable, e.g. `task.total_steps=50000 model.encoder.num_layers=3 -data.sampler.runtime.depth=3 task.augmentation.feature_noise_columns=[]`. - -The generic pretraining task remains available: - -```bash -grass-mil-train task=pretrain_bgrl model=bgrl_module data=spatial_omics -``` - -The task preset selects the step-wise scheduler itself via -`defaults: - override /scheduler: cosine_step`. That resolves because `task` is -merged after `scheduler` in the root defaults list. - -The pretraining task also sets: - -- View generation, selected via the `augmentation` config group - (`bgrl_paper` or `uniform`) — see the augmentation section of - [`hyperparameter_reference.md`](hyperparameter_reference.md) -- Momentum schedule (`momentum`, `momentum_final`, `momentum_warmup_steps`, `total_steps`) - -### Step 2: Choose A Checkpoint - -Use either: - -- `checkpoints/last.ckpt` for latest state -- a specific `epoch_*.ckpt` - -If your callback tracks best metric, use best checkpoint path from training logs. - -### Step 3: Fine-Tune Using The Pretrained Encoder - -```bash -grass-mil-train \ - task=finetune_mil \ - model=supervised_module \ - data=spatial_omics \ - model.init_from_ckpt=/absolute/path/to/pretrain.ckpt \ - model.encoder_init_map=auto_bgrl_or_identity \ - model.init_strict=false -``` - -Optional encoder freezing: - -```bash -grass-mil-train \ - task=finetune_mil \ - model=supervised_module \ - model.init_from_ckpt=/absolute/path/to/pretrain.ckpt \ - model.freeze_encoder=true -``` - -### Checkpoint Mapping Behavior And Failure Cases - -Mapping behavior (`encoder_init_map`): - -- `auto_bgrl_or_identity`: remaps known BGRL encoder prefixes such as `online_encoder.*` and `ssl_model.online_encoder.*` to `encoder.*` -- `identity`: no key remapping -- `spacegm_bgrl`: imports encoders pretrained with the original SPACE-GM style BGRL code (`x_embedding`, `feat_embedding`, `gnns.{i}.gine.*`, `batch_norms.{i}`) - -Common failures: - -1. `FileNotFoundError` if `model.init_from_ckpt` path is invalid. -2. Missing/unexpected key mismatch with `model.init_strict=true`. -3. Shape mismatch when encoder dimensions differ between pretrain and fine-tune configs. -4. Invalid `encoder_init_map` value (must be `identity`, `auto_bgrl_or_identity` or `spacegm_bgrl`). - -## 2b) GRASS-MIL Regimes, Probe And Nested Cross-Validation - -The preprint's fine-tuning settings are `regime=` presets on top of -`model=grass_mil task=finetune_mil` (one region = one bag, region-level -selection); see [`nested_cv.md`](nested_cv.md) for the full protocol and -[`hyperparameter_reference.md`](hyperparameter_reference.md) section 12a for the -values. - -```bash -# one run with explicit splits (ids = region ids of the processed cohort) -grass-mil-train regime=mil data=... model.init_from_ckpt=runs/pretrain_a3/checkpoints/last.ckpt \ - data.split.explicit.train=train.json data.split.explicit.val=val.json data.split.explicit.test=test.json - -# linear head pre-fitted by a logistic probe on the frozen encoder, then attention only -grass-mil-probe regime=linear_probehead data=... model.init_from_ckpt=.ckpt \ - data.split.explicit.train=train.json data.split.explicit.val=val.json probe.output_path=probe_head.pt -grass-mil-train regime=linear_probehead data=... model.init_from_ckpt=.ckpt \ - task.head_init_path=probe_head.pt data.split.explicit.train=train.json ... - -# the whole nested CV (fits the probes itself) -grass-mil-nested-cv regime=linear_ftall_reg data=... model.init_from_ckpt=.ckpt \ - nested_cv.folds_file=folds.json data.graph_labels.patient_column=PATIENT -``` - -Downstream of a trained model: `grass-mil-extract` (cells), `grass-mil-niches`, -`grass-mil-metaclusters` ([`niche_discovery.md`](niche_discovery.md)), -`grass-mil-scores` ([`evaluation_stats.md`](evaluation_stats.md)), -`grass-mil-spatial-stats` ([`spatial_statistics.md`](spatial_statistics.md)). - -## 3) Adjusting Model Settings (Aggregation, Backbone, Loss, MIL Options) - -### A) Mean vs MIL Aggregation - -Mean aggregation: - -```bash -grass-mil-train task=finetune_mean model=supervised_module -``` - -MIL attention aggregation: - -```bash -grass-mil-train task=finetune_mil model=supervised_module -``` - -Notes: - -- `task=finetune_mil` activates attention path and enables MIL-only options. -- `task=finetune_mean` ignores MIL-only controls. - -### B) Backbone And Encoder Adjustments - -Switch backbone type: - -```bash -grass-mil-train model.encoder.conv_type=gin -grass-mil-train model.encoder.conv_type=gcn -grass-mil-train model.encoder.conv_type=gat -grass-mil-train model.encoder.conv_type=graphsage -grass-mil-train model.encoder.conv_type=gine -``` - -Important constraints: - -1. `conv_type=gine` requires `model.encoder.edge_attr_dim` and runtime `edge_attr`. -2. `conv_type=gcn` can use scalar edge weights from `edge_attr` when `model.encoder.use_edge_attr=true`; choose feature by `model.encoder.edge_weight_index`. -3. `conv_type=gat` requires `hidden_dim % gat_heads == 0`. -4. `use_edge_attr=true` on unsupported backbones is ignored with warning. - -### C) Pooling, JK, And Core Encoder Hyperparameters - -Common overrides: - -```bash -grass-mil-train \ - model.encoder.hidden_dim=256 \ - model.encoder.out_dim=256 \ - model.encoder.num_layers=4 \ - model.encoder.dropout=0.2 \ - model.encoder.norm=layernorm \ - model.encoder.jk=concat \ - model.encoder.pooling=max -``` - -### D) Attention Variants - -```bash -grass-mil-train task=finetune_mil \ - model.attention.attention_type=gated - -grass-mil-train task=finetune_mil \ - model.attention.attention_type=gated_projected \ - model.attention.projection_dim=128 -``` - -### E) MIL Region Accumulation - -Manual optimizer stepping over region hyperbatches. This is a throughput -mechanism, not an objective: region cross-entropy remains the sole loss. - -```bash -grass-mil-train task=finetune_mil \ - task.region_accumulation.enabled=true \ - task.region_accumulation.hyperbatch_size=8 -``` - -Dual LR groups for MIL backbone/attention: - -```bash -grass-mil-train task=finetune_mil \ - task.optimization.backbone_lr=1e-3 \ - task.optimization.attention_lr=4e-3 -``` - -### F) Loss And Target-Type Switching - -Binary classification (default): - -```bash -grass-mil-train task.target_type=binary task.loss=categorical_bce -``` - -Regression: - -```bash -grass-mil-train task.target_type=regression task.loss=regression_mse -grass-mil-train task.target_type=regression task.loss=regression_huber -``` - -Survival: - -```bash -grass-mil-train task.target_type=survival task.loss=survival_coxsgd -``` - -Target column selection: - -```bash -grass-mil-train task.target_columns="[outcome]" -``` - -Target format requirements: - -- Binary/regression: `graph_y` shape compatible with output logits. -- Survival: requires at least two columns in `graph_y` interpreted as `[time, event]`. - -## 4) Inference Approaches - -### A) Supported CLI Inference Path (Checkpoint Evaluation) - -`src/grass_mil/eval.py` is the official CLI inference/evaluation entrypoint. - -```bash -grass-mil-eval \ - ckpt_path=/absolute/path/to/checkpoint.ckpt \ - task=finetune_mil \ - model=supervised_module \ - data=spatial_omics -``` - -This runs `trainer.test(...)` and logs test metrics. - -### B) Inference Utility Suite (Prediction + Metrics + Embeddings) - -Use `src/grass_mil/inference/predict.py` for artifact-centric inference runs: - -```bash -grass-mil-predict \ - ckpt_path=/absolute/path/to/checkpoint.ckpt \ - task=finetune_mil \ - model=supervised_module \ - data=spatial_omics \ - aggregation.mode=mean \ - embeddings.extract_node=false -``` - -Internally, this utility performs a single `trainer.predict(...)` pass and -collects prediction + embedding payloads together. - -By default this writes: - -- prediction table: `${paths.output_dir}/${predict.output_subdir}/${output.predictions_filename}` -- metrics report: `${paths.output_dir}/${predict.output_subdir}/${output.metrics_filename}` -- graph embeddings: `${paths.output_dir}/${predict.output_subdir}/${embeddings.filename}` -- summary payload: `${paths.output_dir}/${predict.output_subdir}/${output.summary_filename}` - -### C) Advanced Predict Path (Python API) - -`SupervisedModule.predict_step(...)` still supports direct Python-level prediction output. - -```python -import hydra -from pathlib import Path -from omegaconf import open_dict - -# Compose exactly as train/eval does. -project_root = Path.cwd() -with hydra.initialize_config_dir( - config_dir=str(project_root / "configs"), - version_base="1.3", -): - cfg = hydra.compose( - config_name="eval.yaml", - overrides=[ - "task=finetune_mil", - "model=supervised_module", - "data=spatial_omics", - ], - ) - -# Ensure path roots are concrete for interpolation. -with open_dict(cfg): - cfg.paths.root_dir = str(project_root) - -datamodule = hydra.utils.instantiate(cfg.data) -model = hydra.utils.instantiate(cfg.model) -trainer = hydra.utils.instantiate(cfg.trainer) - -preds = trainer.predict( - model=model, - datamodule=datamodule, - ckpt_path="/absolute/path/to/checkpoint.ckpt", -) - -# Each predict output dict can contain: -# - bag_ids -# - bag_logits -# - bag_targets (if present in batch) -# - bag_attention (MIL mode) -# - row_region_ids / row_sample_ids -# - instance_* fields (when instance payload emission is enabled) -# - embedding_* fields (when embedding payload emission is enabled) -print(preds[0].keys()) -``` - -Expected fields from `predict_step`: - -- `bag_ids` -- `bag_logits` -- `row_region_ids` -- `row_sample_ids` -- optional `bag_targets` (if labels present) -- optional `bag_attention` (when MIL aggregation is active) -- optional `instance_logits`, `instance_patch_ids`, `instance_region_ids`, `instance_sample_ids` -- optional `instance_attention_logits` (MIL + instance payload enabled) -- optional `embedding_bag_ids`, `graph_embeddings`, `embedding_bag_counts` -- optional `node_embeddings`, `node_bag_ids` (when node embedding extraction enabled) - -## 5) Anything Else This Repo Can Do - -See the capability matrix in [`capabilities_and_modes.md`](capabilities_and_modes.md) for details. Highlights: - -1. Runtime shadow sampling (`shadow_native`, `shadow_custom`) including weighted root-node sampling. -2. Multiple logger backends (CSV, TensorBoard, WandB, MLflow, Comet, Neptune, Aim). -3. Debug presets (`fdr`, `limit`, `overfit`, `profiler`) for rapid iteration. -4. Device/runtime presets (`cpu`, `gpu`, `mps`, `ddp`, `ddp_sim`). -5. Resume training from checkpoints via `ckpt_path`. -6. Notebook-based smoke tests for data loading, components, and regime behavior. -7. CI and pre-commit checks for quality and compatibility. - -## 5.1) Train-vs-Validation Sampler Split - -Validation can use a dedicated sampler config (`data.val_sampler`) while training -keeps `data.sampler`. This is useful when training should keep proportional root -sampling but validation should approximate tissue-level uniform sampling. - -Example: - -```bash -grass-mil-train \ - data.sampler.runtime.proportional_root_sampling=true \ - data.val_sampler.name=shadow_custom \ - data.val_sampler.kwargs="{}" \ - data.val_sampler.runtime.enabled=true \ - data.val_sampler.runtime.depth=2 \ - data.val_sampler.runtime.num_neighbors=8 \ - data.val_sampler.runtime.subgraph_batch_size=32 \ - data.val_sampler.runtime.proportional_root_sampling=false \ - data.val_sampler.runtime.subsample_fraction=0.5 \ - data.val_sampler.runtime.subsample_seed=7 -``` - -When `data.val_sampler=null`, validation reuses `data.sampler`. - -## 5.2) Validation Sampling And Epoch Metric Semantics - -Validation behavior is split between step-level logging and epoch-end -aggregation: - -- `validation_step(...)` computes region loss on each batch and logs - `val/loss` with `on_step=true` and `on_epoch=true`. -- Validation predictions/metadata are buffered across the epoch. -- `on_validation_epoch_end(...)` rebuilds a full validation payload (across - distributed ranks when DDP is active), then re-aggregates at - `bag_scope='region'`. -- Aggregation mode at validation epoch-end is: - - `attention_weighted` when `task.aggregation=mil_attention` - - `mean` otherwise -- Validation epoch-end aggregation always uses full instance usage - (`subsample_fraction=1.0`). -- Final task metrics are logged once per epoch from aggregated region-level - outputs: - - binary/categorical: `val/acc` - - regression: `val/mae`, `val/rmse`, `val/r2` - - survival: `val/c_index` - -Practical implications: - -- If you want validation sampling behavior different from training, configure - `data.val_sampler.*` explicitly. -- If `data.val_sampler` is not set, validation inherits `data.sampler`. -- Validation metrics are region-aggregated epoch metrics; batch-level - `val/loss` can diverge from the final epoch metric trend. - -Example: keep training weighted root sampling, but validate with deterministic, -non-proportional sampling. - -```bash -grass-mil-train \ - data.sampler.runtime.proportional_root_sampling=true \ - data.val_sampler.name=shadow_custom \ - data.val_sampler.kwargs="{}" \ - data.val_sampler.runtime.enabled=true \ - data.val_sampler.runtime.proportional_root_sampling=false \ - data.val_sampler.runtime.shuffle_subgraphs=false \ - data.val_sampler.runtime.subsample_fraction=1.0 -``` - -## 6) Optional Smoke Validation Commands - -Quick config and runtime checks: - -```bash -pytest tests/test_configs.py -q -pytest tests/test_train_regimes.py -q -pytest tests/test_eval.py -q -``` - -Pre-commit quality checks: - -```bash -pre-commit run -a -``` - -These are optional sanity checks and are not required for every training run. diff --git a/environment.yaml b/environment.yaml index 9ddd717..3b9f8b4 100644 --- a/environment.yaml +++ b/environment.yaml @@ -1,68 +1,18 @@ -# Conda environment for local development. -# -# This is an alternative to `pip install -e ".[interpretability,report,io]"`. -# It exists because conda resolves the compiled stack (PyTorch, PyG, scipy) -# across all packages simultaneously, and installs precompiled binaries without -# requiring system compilers. -# -# Keep the versions here consistent with `pyproject.toml` (`requires-python` and -# `dependencies`). CI does not use this file: it installs with pip and pins the -# stack inline in `.github/workflows/test.yml`. - +# Conda alternative to the pip install in the README; conda resolves the +# compiled stack (PyTorch, PyG, torch-sparse) together. name: grass-mil channels: - pyg - pytorch - conda-forge - - defaults -# Major versions are pinned so all users share the same majors, while minor and -# patch versions stay free. dependencies: - # Must satisfy pyproject's `requires-python = ">=3.11"`. - python=3.12 - - pytorch=2.* - - lightning=2.* - - torchmetrics=1.* - - hydra-core=1.* - - rich=13.* - - numpy - - pandas - - scipy - - scikit-learn - - joblib - - lightning-utilities - - shapely - - anndata - # PyG and its compiled extensions resolve far more reliably from conda than - # from pip; torch-sparse in particular has no macOS-arm64 wheel. + - pytorch=2.4 - pyg=2.* - - pytorch-scatter - pytorch-sparse - - # --------- interpretability --------- # - - plotly - - jinja2 - - # --------- dev --------- # - - pre-commit=4.* - - pytest=9.* - - # --------- loggers --------- # - # - wandb - # - neptune-client - # - mlflow - # - comet-ml - # - aim>=3.16.2 # no lower than 3.16.2, see https://github.com/aimhubio/aim/issues/2550 - - - pip>=23 + - pytorch-scatter + - pip - pip: - - hydra-colorlog - - rootutils - - umap-learn - - matplotlib - - statsmodels - # optional: required only for SingleCellExperiment/rds loading, and also - # needs an R runtime with Bioconductor - # - rpy2 + - -e .[io,test] diff --git a/src/grass_mil/configs/spatial_stats.yaml b/src/grass_mil/configs/spatial_stats.yaml index 03f2808..458e18d 100644 --- a/src/grass_mil/configs/spatial_stats.yaml +++ b/src/grass_mil/configs/spatial_stats.yaml @@ -15,6 +15,9 @@ cells_dir: ??? output_dir: ${paths.output_dir} units: all # all | held_out | list / file of unit ids positive_class: 1 +# Compare patients (the mean over their units; a patient is positive when any of +# its units is) or individual units. +compare_by: patient # patient | unit analyses: [enrichment, neighbourhoods, contacts, cross_l] # Optional fine -> coarse cell-type grouping ({fine: group}); identity by default. group_map: {} @@ -25,5 +28,5 @@ neighbourhoods: n: 10 # neighbourhood clusters sample_size: 200000 radii: [20.0, 50.0, 100.0] # cross-L radii, coordinate units -min_patients: 10 # per class, for a test +min_patients: 10 # per class (patients or units), for a test seed: 0 diff --git a/src/grass_mil/data/datamodule.py b/src/grass_mil/data/datamodule.py index c9c667f..92f7564 100644 --- a/src/grass_mil/data/datamodule.py +++ b/src/grass_mil/data/datamodule.py @@ -33,7 +33,11 @@ class InputSpec: edge_dim: int def as_dict(self) -> dict[str, Any]: - return {"feature_dim": self.feature_dim, "cell_types": list(self.cell_types), "edge_dim": self.edge_dim} + return { + "feature_dim": self.feature_dim, + "cell_types": list(self.cell_types), + "edge_dim": self.edge_dim, + } @classmethod def from_dict(cls, d: Mapping[str, Any]) -> InputSpec: @@ -173,24 +177,34 @@ def prepare_data(self) -> None: def _group_of(self, entry: UnitEntry) -> str: labels = self.read_labels() - label_id = labels.resolve(entry.unit_id, entry.region_id, entry.sample_id) if labels else None + label_id = ( + labels.resolve(entry.unit_id, entry.region_id, entry.sample_id) if labels else None + ) if labels is not None and label_id is not None and self.labels_cfg.get("patient_column"): return labels.patient(label_id) return entry.region_id def _random_assignment(self, units: Sequence[UnitEntry]) -> dict[str, str]: groups = {u.unit_id: self._group_of(u) for u in units} - # Stratify by the first target when it holds a few integer classes. + # Stratify by class when every group's units share one integer class. strata = None labels = self.read_labels() if labels is not None: - values = {} + classes: dict[str, set[float]] = {} for u in units: label_id = labels.resolve(u.unit_id, u.region_id, u.sample_id) if label_id is not None: - values[groups[u.unit_id]] = float(labels.targets.loc[label_id].iloc[0]) + value = float(labels.targets.loc[label_id].iloc[0]) + classes.setdefault(groups[u.unit_id], set()).add(value) + single = all(len(v) == 1 for v in classes.values()) + values = {g: next(iter(v)) for g, v in classes.items()} if single else {} discrete = all(float(v).is_integer() for v in values.values()) - if discrete and len(values) == len(set(groups.values())) and len(set(values.values())) <= 10: + if ( + values + and discrete + and len(values) == len(set(groups.values())) + and len(set(values.values())) <= 10 + ): strata = {g: str(v) for g, v in values.items()} ratios = self.split_cfg.get("ratios", (0.8, 0.1, 0.1)) return random_split(groups, ratios, int(self.split_cfg.get("seed", 0)), strata) @@ -198,7 +212,10 @@ def _random_assignment(self, units: Sequence[UnitEntry]) -> dict[str, str]: def _members(self, units: Sequence[UnitEntry]) -> dict[str, list[UnitEntry]]: if not self._explicit(): assignment = self._random_assignment(units) - return {s: [u for u in units if assignment[u.unit_id] == s] for s in ("train", "val", "test")} + return { + s: [u for u in units if assignment[u.unit_id] == s] + for s in ("train", "val", "test") + } out: dict[str, list[UnitEntry]] = {} for split in ("train", "val", "test"): ids = read_ids(self.split_cfg.get(split)) or [] @@ -240,7 +257,11 @@ def unit_table(self, dataset: GraphUnitDataset | None = None): dataset = dataset if dataset is not None else self.all_units() rows = [] for entry in dataset.entries: - row = {"unit_id": entry.unit_id, "region_id": entry.region_id, "sample_id": entry.sample_id} + row = { + "unit_id": entry.unit_id, + "region_id": entry.region_id, + "sample_id": entry.sample_id, + } label_id = dataset.label_id(entry) if label_id is not None and self.labels_cfg.get("patient_column"): row["patient_id"] = dataset.labels.patient(label_id) @@ -282,7 +303,10 @@ def loader( def train_dataloader(self) -> EgoGraphLoader: mode = "merged" if self.sampler_cfg.merge_units else "bags" return self.loader( - self.dataset_train, mode=mode, roots_per_unit=self.sampler_cfg.roots_per_unit, shuffle=True + self.dataset_train, + mode=mode, + roots_per_unit=self.sampler_cfg.roots_per_unit, + shuffle=True, ) def val_dataloader(self) -> EgoGraphLoader: diff --git a/src/grass_mil/niches/analysis.py b/src/grass_mil/niches/analysis.py index 08cbf91..bb1e783 100644 --- a/src/grass_mil/niches/analysis.py +++ b/src/grass_mil/niches/analysis.py @@ -72,7 +72,9 @@ def analyse_niches( cells.unit_of_cell, positive_class=positive_class, ) - model, niches = fit_niches(cells.space[clustering.space], margin, clustering) + model, niches = fit_niches( + cells.space[clustering.space], margin, cells.unit_of_cell, clustering + ) positive = cells.positive(positive_class) table = niche_table( niches, diff --git a/src/grass_mil/niches/discovery.py b/src/grass_mil/niches/discovery.py index 462ce3e..a1f0deb 100644 --- a/src/grass_mil/niches/discovery.py +++ b/src/grass_mil/niches/discovery.py @@ -6,8 +6,9 @@ assigned by a k-nearest-neighbour vote on the same coordinates. Cells HDBSCAN leaves unassigned form the ``Background``. -Niches are numbered by their attention-weighted margin, ``Niche 1`` being the -most positive; ``0`` is the Background. +Niches are numbered by their attention-weighted margin (the mean over units of +the niche's mean cell margin), ``Niche 1`` being the most positive; ``0`` is +the Background. """ from __future__ import annotations @@ -100,12 +101,15 @@ def settings(self) -> dict[str, Any]: def fit_niches( - x: np.ndarray, margin: np.ndarray, cfg: ClusteringConfig + x: np.ndarray, margin: np.ndarray, unit_of_cell: np.ndarray, cfg: ClusteringConfig ) -> tuple[NicheModel, np.ndarray]: """Fit the clustering on ``x`` and return the model and every cell's niche id. - ``margin`` (per cell) orders the niches. + Niches are ordered by their margin: the mean over units of the niche's mean + cell ``margin`` in the unit (the statistic the niche table reports). """ + import pandas as pd + from sklearn.cluster import HDBSCAN, KMeans from sklearn.decomposition import PCA from sklearn.neighbors import KNeighborsClassifier @@ -140,8 +144,9 @@ def fit_niches( model = NicheModel(cfg, pca, assigner, {}, sample) labels = assigner.predict(model.coordinates(x)) labels[sample] = raw - # Number the clusters by mean margin, most positive first; -1 is Background. - clusters = [c for c in np.unique(labels) if c >= 0] - order = sorted(clusters, key=lambda c: -float(np.mean(margin[labels == c]))) + # Number the clusters by margin, most positive first; -1 is Background. + per_unit = pd.Series(margin).groupby([labels, unit_of_cell]).mean() + score = per_unit.groupby(level=0).mean() + order = sorted((c for c in score.index if c >= 0), key=lambda c: -float(score[c])) model.mapping = {-1: BACKGROUND, **{int(c): rank + 1 for rank, c in enumerate(order)}} return model, np.vectorize(model.mapping.get)(labels).astype(int) diff --git a/src/grass_mil/niches/replication/stages.py b/src/grass_mil/niches/replication/stages.py index d50a64f..545230f 100644 --- a/src/grass_mil/niches/replication/stages.py +++ b/src/grass_mil/niches/replication/stages.py @@ -74,7 +74,7 @@ def cluster_folds( for name, cache in caches.items(): cells = load_cells(cache, sorted(cache.held_out)) margin = _margins(cells, positive_class) - model, niches = fit_niches(cells.space[cfg.space], margin, cfg) + model, niches = fit_niches(cells.space[cfg.space], margin, cells.unit_of_cell, cfg) positive = cells.positive(positive_class) table = niche_table( niches, diff --git a/src/grass_mil/spatial/cli.py b/src/grass_mil/spatial/cli.py index 720bd56..e596c65 100644 --- a/src/grass_mil/spatial/cli.py +++ b/src/grass_mil/spatial/cli.py @@ -44,8 +44,16 @@ def run_spatial_stats(cfg: DictConfig) -> Path: partners = [j for j in range(g) if j not in reference] rng = np.random.default_rng(int(cfg.seed)) positive = cells.positive(int(cfg.positive_class)) - patient_of = {u: cells.patients.get(u, u) for u in cells.units} - y = pd.Series({patient_of[u]: int(positive[u]) for u in cells.units}) + if cfg.compare_by not in ("patient", "unit"): + raise ValueError("compare_by must be 'patient' or 'unit'.") + if cfg.compare_by == "patient": + patient_of = {u: cells.patients.get(u, u) for u in cells.units} + else: + patient_of = {u: u for u in cells.units} + # A patient is positive when any of its units is. + y = pd.Series(dtype=int) + for u in cells.units: + y[patient_of[u]] = max(int(positive[u]), int(y.get(patient_of[u], 0))) per_unit: dict[str, dict[str, dict[str, float]]] = {a: {} for a in cfg.analyses} windows = [] diff --git a/tests/test_niches.py b/tests/test_niches.py index 445ee96..06a9f93 100644 --- a/tests/test_niches.py +++ b/tests/test_niches.py @@ -31,7 +31,10 @@ def test_niches_are_numbered_by_margin_with_background_zero(): rng = np.random.default_rng(0) x = np.concatenate([rng.normal(c, 0.05, size=(100, 2)) for c in ((0, 1), (1, 0), (-1, 0))]) margin = np.repeat([0.0, 2.0, -2.0], 100) - model, niches = fit_niches(x, margin, ClusteringConfig(min_cluster_size=20, sample_size=200)) + units = np.repeat(["u1", "u2"], 150) + model, niches = fit_niches( + x, margin, units, ClusteringConfig(min_cluster_size=20, sample_size=200) + ) assert model.n_niches == 3 assert np.mean(niches[100:200] == 1) > 0.9 and np.mean(niches[200:] == 3) > 0.9 assert np.array_equal(model.predict(x[:5]), niches[:5]) diff --git a/tools/update_repository_map.py b/tools/update_repository_map.py deleted file mode 100644 index f721126..0000000 --- a/tools/update_repository_map.py +++ /dev/null @@ -1,152 +0,0 @@ -"""Regenerate docs/repository_map.md from `git ls-files`. - - python tools/update_repository_map.py - -Descriptions come from a small table of known files, else from a module's first -docstring line, a YAML file's first comment, or a Markdown file's title. -""" - -from __future__ import annotations - -import ast -import re -import subprocess -from pathlib import Path - -ROOT = Path(__file__).resolve().parents[1] - -KNOWN = { - ".project-root": "Root marker used by `rootutils.setup_root(...)`.", - ".gitignore": "Ignore rules for generated artefacts, environments, logs, data and vignette outputs.", - ".pre-commit-config.yaml": "Code-quality hooks (ruff lint + format, file hygiene checks).", - "LICENSE": "MIT license.", - "README.md": "Landing page: method summary, install, typical usage, workflow table, citation.", - "CITATION.cff": "Machine-readable citation (preprint DOI 10.64898/2026.10.08.757722).", - "CHANGELOG.md": "Release notes.", - "Makefile": "Convenience targets: `test`, `format`, `clean`, `clean-logs`.", - "pyproject.toml": "Package metadata, dependencies, extras, console entry points (`grass-mil-*`), pytest/ruff configuration.", - "environment.yaml": "Conda environment alternative for local development.", - ".github/PULL_REQUEST_TEMPLATE.md": "Pull-request template.", - ".github/codecov.yml": "Codecov configuration.", - ".github/dependabot.yml": "Dependabot update policy.", - ".github/release-drafter.yml": "Release drafter configuration.", - ".github/workflows/test.yml": "CI: pytest on Linux and macOS, coverage, install-from-source smoke test.", - ".github/workflows/code-quality-main.yaml": "CI: pre-commit on every push to main.", - ".github/workflows/code-quality-pr.yaml": "CI: pre-commit on changed files of a pull request.", - ".github/workflows/release-drafter.yml": "CI: draft release notes from merged pull requests.", - ".vscode/settings.json": "Editor settings.", - "tools/update_repository_map.py": "Regenerates this map.", - "vignettes/README.md": "Index of the worked examples on the public MIBI-TOF dataset.", - "vignettes/01_anndata_to_spatial_graphs.ipynb": "Vignette 1: AnnData -> manifest -> cellular graphs.", - "vignettes/02_pretraining_and_niches.ipynb": "Vignette 2: BGRL pretraining, fine-tuning, instance export, interpretability report.", - "vignettes/03_nested_cv_niches_and_metaclusters.ipynb": "Vignette 3: nested CV, cell extraction, niches, metaclusters, scores, spatial statistics.", - "src/grass_mil/configs/local/.gitkeep": "Placeholder for untracked machine-local configs.", - "tests/conftest.py": "Shared pytest fixtures (Hydra configs on a synthetic datamodule).", - "tests/helpers/synthetic_cohort.py": "Synthetic cohort on disk (regions, labels, folds) with a planted niche, used by the phase tests.", - "tests/helpers/synthetic_datamodule.py": "In-memory bag datamodule for fast training tests.", -} - -GROUPS = [ - ("Root metadata and tooling", lambda f: "/" not in f), - ("CI and contribution files", lambda f: f.startswith((".github/", ".vscode/"))), - ("Maintenance tools", lambda f: f.startswith("tools/")), - ( - "Package: entry points and shared modules", - lambda f: f.startswith("src/grass_mil/") and f.count("/") == 2, - ), - ("Package: data pipeline", lambda f: f.startswith("src/grass_mil/data/")), - ("Package: models", lambda f: f.startswith("src/grass_mil/models/")), - ("Package: inference", lambda f: f.startswith("src/grass_mil/inference/")), - ("Package: evaluation statistics", lambda f: f.startswith("src/grass_mil/evaluation/")), - ( - "Package: interpretability (niches, metaclusters, report)", - lambda f: f.startswith("src/grass_mil/interpretability/"), - ), - ( - "Package: classical spatial statistics", - lambda f: f.startswith("src/grass_mil/spatial_stats/"), - ), - ("Package: utilities", lambda f: f.startswith("src/grass_mil/utils/")), - ("Hydra configs", lambda f: f.startswith("src/grass_mil/configs/")), - ("Documentation", lambda f: f.startswith("docs/")), - ("Vignettes", lambda f: f.startswith("vignettes/")), - ("Tests", lambda f: f.startswith("tests/")), -] - - -def describe(path: str) -> str: - if path in KNOWN: - return KNOWN[path] - p = ROOT / path - if p.suffix == ".py": - try: - doc = ast.get_docstring(ast.parse(p.read_text())) - except Exception: - doc = None - if doc: - first = doc.strip().splitlines()[0].strip() - return first if first.endswith((".", ":")) else first + "." - if p.name == "__init__.py": - return "Package marker / re-exports." - if path.startswith("tests/"): - return ( - "Tests for " - + re.sub(r"^tests?[_/]|test_|\.py$", "", p.stem).replace("_", " ") - + "." - ) - return "Module." - if p.suffix in {".yaml", ".yml"}: - for line in p.read_text().splitlines(): - line = line.strip() - if line.startswith("#") and not line.startswith("# @package"): - text = line.lstrip("# ").rstrip() - return text if text.endswith((".", ":")) else text + "." - if line and not line.startswith("#"): - break - return "Hydra config: `" + p.stem + "`." - if p.suffix == ".md": - return p.read_text().strip().splitlines()[0].lstrip("# ").strip() + "." - if p.suffix == ".ipynb": - return "Notebook." - return "Tracked file." - - -def main() -> None: - files = subprocess.run( - ["git", "ls-files"], capture_output=True, text=True, cwd=ROOT - ).stdout.split() - files = [f for f in files if (ROOT / f).exists()] - out = [ - "# Repository Map", - "", - "Every tracked file (`git ls-files`) with a one-line purpose. Regenerate with", - "`python tools/update_repository_map.py` after adding files.", - "", - ] - seen: set[str] = set() - for title, pred in GROUPS: - members = [f for f in files if f not in seen and pred(f)] - if not members: - continue - out += [f"## {title}", ""] - for f in sorted(members): - seen.add(f) - out.append(f"- `{f}`: {describe(f)}") - out.append("") - rest = [f for f in files if f not in seen] - if rest: - out += ["## Other", ""] + [f"- `{f}`: {describe(f)}" for f in sorted(rest)] + [""] - out += [ - "## Generated at runtime (untracked)", - "", - "- `logs/`: Hydra run directories (checkpoints, `metrics_per_epoch.jsonl`, predictions, caches, reports).", - "- `data/`: raw manifests and processed graphs (`processed_index.json`, `metadata.json`, `graphs/`).", - "- `vignettes/vignette_work/`: data, graphs and runs produced by the vignettes.", - "", - ] - (ROOT / "docs/repository_map.md").write_text("\n".join(out)) - print(len(files), "files mapped") - - -if __name__ == "__main__": - main() diff --git a/vignettes/01_anndata_to_cell_graphs.ipynb b/vignettes/01_anndata_to_cell_graphs.ipynb new file mode 100644 index 0000000..fd0e80d --- /dev/null +++ b/vignettes/01_anndata_to_cell_graphs.ipynb @@ -0,0 +1,492 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "d6bc8aef", + "metadata": {}, + "source": [ + "# 1. From AnnData to cell graphs\n", + "\n", + "This vignette turns a public spatial-proteomics dataset into the inputs GRASS-MIL\n", + "reads, builds the cell graphs, and defines the phenotype that vignettes 2 and 3\n", + "predict and explain.\n", + "\n", + "**Dataset.** `squidpy.datasets.mibitof()`: MIBI-TOF images of colorectal\n", + "carcinoma, 3 images, about 3,300 cells with 8 annotated cell types (Hartmann et\n", + "al.). `squidpy` is only used to download it." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "c5361a36", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-15T07:40:19.761536Z", + "iopub.status.busy": "2026-10-15T07:40:19.761386Z", + "iopub.status.idle": "2026-10-15T07:40:23.546546Z", + "shell.execute_reply": "2026-10-15T07:40:23.545913Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "3309 cells, 3 images\n" + ] + }, + { + "data": { + "text/plain": [ + "Cluster\n", + "Tcell_CD4 799\n", + "Epithelial 746\n", + "Imm_other 488\n", + "Tcell_CD8 464\n", + "Fibroblast 270\n", + "Myeloid_CD68 259\n", + "Myeloid_CD11c 168\n", + "Endothelial 115\n", + "Name: count, dtype: int64" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# pip install 'grass-mil[io]' squidpy\n", + "import collections\n", + "import json\n", + "import pathlib\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "import squidpy as sq\n", + "import torch\n", + "\n", + "WORK = pathlib.Path(\"vignette_work\").resolve()\n", + "DATA = WORK / \"data\"\n", + "DATA.mkdir(parents=True, exist_ok=True)\n", + "\n", + "adata = sq.datasets.mibitof(path=WORK / \"mibitof.h5ad\")\n", + "print(adata.n_obs, \"cells,\", adata.obs[\"library_id\"].nunique(), \"images\")\n", + "adata.obs[\"Cluster\"].value_counts()" + ] + }, + { + "cell_type": "markdown", + "id": "f65a0f1f", + "metadata": {}, + "source": [ + "## Cell tables and manifest\n", + "\n", + "GRASS-MIL reads one cell table per region, listed in a manifest. Here every image\n", + "is one region, saved as an h5ad file with the cell type in `obs[\"cell_type\"]` and\n", + "the coordinates in `obsm[\"spatial\"]`. Only the cell type enters the model\n", + "(`data.input.features=null`), so any assay with harmonised cell types can be\n", + "analysed the same way." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "f37e4870", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-15T07:40:23.548457Z", + "iopub.status.busy": "2026-10-15T07:40:23.548153Z", + "iopub.status.idle": "2026-10-15T07:40:23.662505Z", + "shell.execute_reply": "2026-10-15T07:40:23.661602Z" + } + }, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
sample_idinput_path
0point16point16.h5ad
1point23point23.h5ad
2point8point8.h5ad
\n", + "
" + ], + "text/plain": [ + " sample_id input_path\n", + "0 point16 point16.h5ad\n", + "1 point23 point23.h5ad\n", + "2 point8 point8.h5ad" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "rows = []\n", + "for image in sorted(adata.obs[\"library_id\"].unique()):\n", + " sub = adata[adata.obs[\"library_id\"] == image].copy()\n", + " sub.obs = pd.DataFrame({\"cell_type\": sub.obs[\"Cluster\"].astype(str)}, index=sub.obs_names)\n", + " sub.write_h5ad(DATA / f\"{image}.h5ad\")\n", + " rows.append({\"sample_id\": image, \"input_path\": f\"{image}.h5ad\"})\n", + "pd.DataFrame(rows).to_csv(DATA / \"manifest.csv\", index=False)\n", + "pd.DataFrame(rows)" + ] + }, + { + "cell_type": "markdown", + "id": "20e1ea32", + "metadata": {}, + "source": [ + "## Settings shared by every command\n", + "\n", + "Each image is about 1,024 pixels wide at 0.39 µm per pixel, so\n", + "`coord_scale_um=0.39` puts every distance in micrometres. The images are cut into\n", + "100 µm tiles (`data.unit=tile`); a tile is one bag, one label and one\n", + "cross-validation member. Cells within 20 µm of each other are flagged as\n", + "neighbours on the graph edges.\n", + "\n", + "The same overrides are passed to every command in vignettes 2 and 3, so they are\n", + "stored once in `cohort.json` (in your own project they belong in a cohort file,\n", + "`configs/experiment/example_cohort.yaml`)." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "27c51daf", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-15T07:40:23.663905Z", + "iopub.status.busy": "2026-10-15T07:40:23.663799Z", + "iopub.status.idle": "2026-10-15T07:40:23.666421Z", + "shell.execute_reply": "2026-10-15T07:40:23.666039Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "14 settings saved to cohort.json\n" + ] + } + ], + "source": [ + "overrides = [\n", + " f\"data.manifest_path={DATA / 'manifest.csv'}\",\n", + " f\"data.processed_dir={DATA / 'processed'}\",\n", + " \"data.unit=tile\",\n", + " \"data.tiling.size=100\",\n", + " \"data.tiling.stride=100\",\n", + " \"data.coord_scale_um=0.39\",\n", + " \"data.min_cells=20\",\n", + " \"data.input.features=null\",\n", + " \"data.graph.neighbor_cutoff=20\",\n", + " f\"data.labels.file={DATA / 'labels.csv'}\",\n", + " \"data.labels.id_column=unit_id\",\n", + " \"data.labels.patient_column=image\",\n", + " \"task.target_columns=[infiltrated]\",\n", + " f\"paths.log_dir={WORK / 'logs'}\",\n", + "]\n", + "(WORK / \"cohort.json\").write_text(json.dumps({\"overrides\": overrides}, indent=1))\n", + "print(len(overrides), \"settings saved to cohort.json\")" + ] + }, + { + "cell_type": "markdown", + "id": "f95c7a3a", + "metadata": {}, + "source": [ + "## Build the graphs\n", + "\n", + "Every command builds the graphs on first use; here we build them directly with\n", + "the datamodule the commands use. Every tile becomes a Delaunay graph of its\n", + "cells, with the bounded log edge length (`log_distance`) and the `neighbor`\n", + "flag as edge features." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "c08316d3", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-15T07:40:23.667392Z", + "iopub.status.busy": "2026-10-15T07:40:23.667315Z", + "iopub.status.idle": "2026-10-15T07:40:24.686247Z", + "shell.execute_reply": "2026-10-15T07:40:24.685259Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "44 tiles; cell types: ('Endothelial', 'Epithelial', 'Fibroblast', 'Imm_other', 'Myeloid_CD11c', 'Myeloid_CD68', 'Tcell_CD4', 'Tcell_CD8')\n", + "tiles per image: {'point16': 16, 'point23': 12, 'point8': 16}\n" + ] + } + ], + "source": [ + "from hydra import compose, initialize_config_module\n", + "from hydra.utils import instantiate\n", + "\n", + "with initialize_config_module(version_base=\"1.3\", config_module=\"grass_mil.configs\"):\n", + " cfg = compose(\"train\", overrides=[o for o in overrides if not o.startswith(\"data.labels.file\")])\n", + "datamodule = instantiate(cfg.data, _recursive_=False)\n", + "datamodule.prepare_data()\n", + "units = datamodule.all_units()\n", + "print(len(units), \"tiles; cell types:\", datamodule.input_spec.cell_types)\n", + "print(\"tiles per image:\", dict(collections.Counter(e.region_id for e in units.entries)))" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "284c1d8e", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-15T07:40:24.687820Z", + "iopub.status.busy": "2026-10-15T07:40:24.687708Z", + "iopub.status.idle": "2026-10-15T07:40:24.794387Z", + "shell.execute_reply": "2026-10-15T07:40:24.793924Z" + } + }, + "outputs": [ + { + "data": { + "image/png": "iVBORw0KGgoAAAANSUhEUgAAAdkAAAHnCAYAAAD0NyffAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOydeZgcZdX276qu3vfu2dcEkD0hkAAJJJCFnbCqLAqyKsomAiKIaEAR4cUNUFH0Iwi8IIogKPgSMYEQEkgIIWERRGB6tp6Z3vfuWp7vj7GK7pnunl6qepnU77pywXRXV1fX8pznnOec+1CEEAIVFRUVFRUV2aHrfQAqKioqKiqzFdXIqqioqKioKIRqZFVUVFRUVBRCNbIqKioqKioKoRpZFRUVFRUVhVCNrIqKioqKikKoRlZFRUVFRUUhVCOroqKioqKiEKqRVVFRUVFRUQjVyO6GXHjhhZgzZ05Fn3311VexZs0ahEKhae+98soruPTSS7Fw4ULo9XpQFIVPPvmk4L4GBgZw8cUXo6urC3q9Ht3d3TjjjDPKOp6RkRGsWbMGO3bsmPbemjVrQFFUzmvLly/H8uXLy/qObLZv345jjjkGFosFDocDZ555Jj766KO82957773Yd999odfrMXfuXNx6661gWTZnm7Vr14KiqLz/vF5vxce5fPnyvPs84YQTpm374Ycf4vzzz0dfXx+MRiP23HNPXHvttfD7/RV/f6VQFIU1a9ZIf2/YsAEURWHDhg2KfefPfvYznHnmmZg7dy4oiip4f5RzrTKZDL773e9i7ty50Ol06O/vx0033YRkMqnY78hmzpw5uPDCC2vyXSrFYep9ACq155ZbbsHXv/71ij776quv4tZbb8WFF14Ih8OR896LL76If/zjHzj44INhs9mKDoxvv/02li9fjj322AN33303enp6MDo6iv/7v/8r63hGRkZw6623Ys6cOViwYEHOe5deemleo1Ip//rXv7B8+XIsWLAATzzxBFKpFL773e9i2bJl2LFjB1pbW6Vtb7/9dtxyyy248cYbcdxxx2Hr1q34zne+g+HhYfzmN7+Ztu8HH3wQ++67b85rbre7quPdY4898Oijj+a8NvWaTUxMYPHixbDZbPj+97+Pvr4+vPnmm/je976H9evX44033gBNz+65+P333w+z2YyVK1fi2WefnXH7Uq7Vueeei+eeew7f/e53ceihh2Lz5s34wQ9+gHfeeQfPPPOMrMev0uAQFZUy+J//+R8CgHz88cfT3uN5vqTtBEEgCxYsIAsWLCCpVKqq49m6dSsBQB588MGStj/66KPJ0UcfXdF3ff7znyctLS0kHA5Lr33yySdEq9WSG264QXrN5/MRg8FAvvKVr+R8/vbbbycURZF33nlHeu3BBx8kAMjWrVsrOqZCHH300eSAAw6YcbsHHniAACD/+Mc/cl7/4Q9/SACQ7du3y3pcMwGAfO9735P+Xr9+PQFA1q9fr9h3Zt+3BxxwQMH7o9RrtXnzZgKA/PjHP855XTynL7zwQtXHPBP9/f3kggsuUPx7VGZmdk9RZwFiyPPNN9/EmWeeCZvNBrvdjvPOOw8TExM52wqCgLvuuksKUba1teFLX/oShoaGcrbLFy6mKApXXnklHn74Yey3334wmUw46KCD8Ne//jXnWL75zW8CgBRayw7llerxvPzyy9ixYweuueYa6PX6Ms/Ip2zYsAGHHnooAOCiiy6SjkcMN+YLF+cjk8ngBz/4gXTeWltbcdFFF+WcX47j8Ne//hWf/exnYbPZpNf7+/uxYsUKPPXUU9Jrf//735FKpXDRRRflfM9FF10EQgiefvrpin+z3Gi1WgCA3W7PeV30eA0Gw4z7+Ne//oVzzz0X7e3t0Ov16Ovrw5e+9CWk02lpG6/Xi8suuww9PT3Q6XRS+JzjuLKP+aOPPsI555wjLTO0t7dj1apVeZcMSkFuT33Tpk0AgJNOOinn9dWrVwMAnnzyyRn3Uco9CQAsy+KGG25AR0cHTCYTli5ditdffz3vPl955RUsWbIEBoMB3d3duOWWW/Db3/4277LOH/7wByxZsgRmsxkWiwXHH3883nzzzZxt5L4OsxXVyDYJZ5xxBvbaay/86U9/wpo1a/D000/j+OOPz1nj+9rXvoZvfetbOPbYY/HMM8/g+9//Pv7+97/jiCOOgM/nm/E7/va3v+G+++7DbbfdhieffBIulwtnnHGGtOZ46aWX4qqrrgIA/PnPf8bmzZuxefNmHHLIIWX9lpdffhkAYLVacdJJJ8FgMMBisWD16tX417/+VfJ+DjnkEDz44IMAgO985zvS8Vx66aUl70MQBJx22mn40Y9+hC984Qv429/+hh/96EdYt24dli9fLq2h/ec//0EymcT8+fOn7WP+/Pn48MMPkUqlAEyGwgFg3rx5Odt1dnaipaVFej+b1atXQ6PRwOVy4cwzz8y7jThpKHV98j//+Q9cLhcYhsGee+6Jm2++edqa4Omnn46+vj5cd911eOeddxCLxfDyyy/jRz/6EU455RTst99+Rb/jrbfewqGHHootW7bgtttuw/PPP4877rgD6XQamUwGwKSBPeyww/B///d/+O53v4vnn38el1xyCe644w58+ctfLum3ZHPSSSfhjTfewF133YV169bhV7/6FQ4++OC8eQJKMNO1En/31Amk+PfOnTuL7r/UexIAvvzlL+Puu+/Gl770JfzlL3/BZz/7WZx55pkIBoM5+9y5cyeOPfZYJBIJPPTQQ7j//vuxfft23H777dO+/4c//CHOPfdc7L///njiiSfw8MMPIxqNYtmyZXj33Xel7ep9HZqGervSKsX53ve+RwCQb3zjGzmvP/roowQAeeSRRwghhLz33nsEALn88stztnvttdcIAPLtb39beu2CCy4g/f39OdsBIO3t7SQSiUiveb1eQtM0ueOOO6TXioWBsym23WWXXUYAEJvNRi655BLyj3/8gzz88MOkv7+ftLS0kJGRkaL7zqZYuFg8d9lMDRc/9thjBAB58skn8+73l7/8JSGEkE2bNhEA5LHHHpv2PWIYUDzuL3/5y0Sv1+c93r333pscd9xx0t/PP/88ufnmm8mzzz5LXnrpJXLfffeRnp4eYjabyY4dO3I+e+uttxKNRkM2bNhQ+IT8l5tvvpn88pe/JP/85z/J3/72N3LllVcShmHIUUcdlRMeJYSQkZERsmTJEgJA+vf5z3++pFD+ypUricPhIOPj4wW3ueyyy4jFYiEDAwM5r999990EQE74HDOEi30+HwFAfvazn814bJVQLFxc6rV6+umnCQDy8MMP53z+d7/7HQFA9t5776LHUOo9KT7zhcaG7HDx5z//eWI2m8nExIT0Gs/zZP/99895Tj0eD2EYhlx11VU5+4xGo6Sjo4OcddZZhBDlr8NsQjWyDY5oKLZt25bzOsuyhGEYcskllxBCCPnlL39JAJDXX3992j72228/cvjhh0t/FzKy55xzzrTPdnR0kK9+9avS33IY2S9/+csEADn++ONzXn/zzTcJAHLzzTcX3Xc21RrZL37xi8ThcJBMJkNYls35lz2oiEb28ccfn/Y9opEdHR2Vfp/BYMh7vHvvvfe03z2Vjz/+mFgsFnLqqacW3a5cRKP25z//WXotEAiQQw89lBxwwAHk0UcfJS+//DL55S9/STo7O8lxxx1HWJYtuL94PE40Gs20teepdHd3k1NOOWXa+X3nnXdyjAYhMxtZQRDInnvuSbq7u8mPf/xjsn379mmThmooZmTzke9apdNpstdee5Guri7ywgsvkGAwSJ5//nnS3t5ONBoN2XfffYvus9R7UnzmC40N2Ua2ra2NnHLKKdO+a82aNTnPqbhGv3Xr1mnfffbZZ5O2tjZCiPLXYTahhoubhI6Ojpy/GYaB2+2WyizE/3Z2dk77bFdXV0nlGPmyWfV6vexlB+L3HH/88TmvL1iwAJ2dndi+fbus31eMsbExhEIh6HQ6aLXanH9er1cKs4vHnO88BgIBUBQlrWO63W6kUikkEom827pcrqLHNGfOHCxduhRbtmyp8tflct555wFAzn7vvPNO7NixA+vWrcMXvvAFLFu2DF/72tfw6KOP4oUXXpiWnZxNMBgEz/Po6ekp+r1jY2N49tlnp53fAw44AABKWsoQoSgKL774Io4//njcddddOOSQQ9Da2oqrr74a0Wi05P3IRb5rpdPp8Pzzz6Ovrw/HHXccnE4nPve5z+Hb3/42nE4nuru7i+6z1HtSvBcLjQ3Z+P1+tLe3T/uuqa+NjY0BAA499NBp3/2HP/xB+u5Guw6NjFrC0yR4vd6ch5PjOPj9fulhEv87Ojo6bdAbGRlBS0tL7Q52BvKta4oQQmpaMtLS0gK3242///3ved+3Wq0AgD333BNGoxG7du2ats2uXbuw1157SUlC4lrsrl27cPjhh0vbiQPkgQceOONxKXkesve7Y8cOdHd3T5uciQll+daGRVwuFzQazbTEuqm0tLRg/vz5edf/gMlJYDn09/fjd7/7HQDggw8+wBNPPIE1a9Ygk8ng/vvvL2tfcpDvWu21117YvHkzhoeHEQgEsOeeeyIcDuPrX/86jjrqqKL7K/WeFJ/5QmNDNm63WzKg2Uyt7xXHiT/96U/o7+8vepyNdh0aFdWTbRKmehRPPPEEOI6TCudXrlwJAHjkkUdyttu6dSvee+89rFq1SpbjEJM3qvFuTzzxRJhMJjz//PM5r2/fvh1erxeLFy+u2fGsXr0afr8fPM9j0aJF0/7ts88+ACa9g1NOOQV//vOfc2bqHo8H69evx5lnnim9dsIJJ8BgMGDt2rU53yWKGZx++ulFj+njjz/Gpk2byjoPpfDQQw8BQM5+u7q6MDQ0hOHh4ZxtN2/eDABFvVSj0Yijjz4af/zjH4t6o6tXr8bbb7+NPffcM+85LtfIZrP33nvjO9/5DubNm1fTCIjITNequ7sb8+bNg8lkwv/8z//AbDbjkksuKbrPUu9J8dkvNDZkc/TRR+Of//xnznUSBAF//OMfc7Y7/vjjwTAM/vOf/+T97kWLFuU95npfh0ZG9WSbhD//+c9gGAbHHnss3nnnHdxyyy046KCDcNZZZwEA9tlnH3zlK1/BvffeC5qmceKJJ+KTTz7BLbfcgt7eXnzjG9+Q5ThEL+3nP/85LrjgAmi1Wuyzzz6wWq2YmJjASy+9BACSx/f888+jtbUVra2tOProowFMlofcdtttuP7663HhhRfi3HPPhdfrxS233IK+vj5cfvnlJR+P6GE++uij2G+//WCxWNDV1VXywH3OOefg0UcfxUknnYSvf/3rOOyww6DVajE0NIT169fjtNNOk1Sobr31Vhx66KFYvXo1brzxRkmMoqWlBdddd520T5fLhe985zu45ZZb4HK5JDGKNWvW4NJLL8X+++8vbXvMMcfgqKOOwvz582Gz2bBr1y7cddddoCgK3//+93OO9bbbbsNtt92GF198UTqX+di4cSNuv/12nHHGGdhjjz2QSqXw/PPP4ze/+Q1WrlyJU045Rdr2iiuuwKOPPopjjz0WN954I3p7e/H222/jBz/4Adrb2/HFL36x6Pn7yU9+gqVLl+Lwww/HjTfeiL322gtjY2N45pln8Otf/xpWqxW33XYb1q1bhyOOOAJXX3019tlnH6RSKXzyySd47rnncP/9988YchbZuXMnrrzySnz+85/HZz7zGeh0Ovzzn//Ezp07ceONN0rbrV27FhdddBEefPDBGZWPtm3bJpWwRCIREELwpz/9CcCkRy96dOVcq7vuugsdHR3o6+vD2NgYnnjiCTz99NN4+OGHZwwXl3pP7rfffjjvvPPws5/9DFqtFscccwzefvtt3H333TllZgBw880349lnn8WqVatw8803w2g04v7770c8HgfwaXRjzpw5uO2223DzzTfjo48+wgknnACn04mxsTG8/vrrMJvNuPXWW0u+DipQs4sbHTF554033iCnnHIKsVgsxGq1knPPPZeMjY3lbMvzPLnzzjvJ3nvvTbRaLWlpaSHnnXceGRwczNmuUOLTFVdcMe378xW133TTTaSrq4vQNJ2TlCImqeT7ly+Z5IEHHiAHHngg0el0xO12ky9+8YvTjrUUHnvsMbLvvvsSrVabkzhTSuITIZOJInfffTc56KCDiMFgIBaLhey7777ksssuI//+979ztt22bRtZtWoVMZlMxGazkdNPP518+OGHeY/r5z//Odl7772JTqcjfX195Hvf+x7JZDI521xzzTVk//33J1arlTAMQ7q6ush5551H3n///Wn7E3/PTMIM//73v8lJJ51Euru7iV6vJwaDgcybN4/cfvvteTOGt2/fTs444wzS09ND9Ho92WOPPcill15KPB5P0e8Reffdd8nnP/954na7pd964YUX5nzXxMQEufrqq8ncuXOJVqslLpeLLFy4kNx8880kFotJ22GGxKexsTFy4YUXkn333ZeYzWZisVjI/PnzyU9/+lPCcZz0uXvvvZcAIH//+99nPP4LLrig4H2bnVBXzrW69dZbyZ577kn0ej1xOBzkhBNOIC+//HJJ55OQ0u/JdDpNrrvuOtLW1kYMBgNZvHgx2bx5c97nduPGjeTwww8ner2edHR0kG9+85vkzjvvJABIKBTK2fbpp58mK1asIDabjej1etLf308+97nPSaIlpV4HFUIoQgippVFXKY81a9bg1ltvxcTEREOtq6qoNDJnnXUWPv74Y2zdurXeh9LQHHfccfjkk0/wwQcf1PtQZi1quFhFRWVWQQjBhg0bpuUn7O5ce+21OPjgg9Hb24tAIIBHH30U69atk5KXVJRBNbIqDQkhBDzPF91Go9GUJJuosntBURTGx8frfRgNB8/z+O53vwuv1wuKorD//vvj4Ycflkq7VJRBDRerNCQbNmzAihUrim5TSlKLioqKSj1RjaxKQxKNRvH+++8X3Wbu3LlVt4NTUVFRURLVyKqoqKioqCiEKkahoqKioqKiEE2Z+CQIAkZGRmC1WtXEFxUVFRWVmkIIQTQaRVdX14zyp01pZEdGRtDb21vvw1BRUVFR2Y0ZHBycUa2sKY2sKJA9ODg4TT5MRUVFRUVFSSKRCHp7eyVbVIymNLJiiNhms6lGVkVFRUWlLpSyXKkmPqmoqKioqCiEamRVVFRUVFQUQjWyKioqKioqCqEaWRUVFRUVFYVQjayKioqKiopCqEZWRUVFRUVFIVQjq6KioqKiohCqkVVRUVFRUVEI1ciqqKioqKgohGpkVVRUVFRUFEI1sioqKioqKgqhGlkVFRUVFRWFKMvI3nHHHTj00ENhtVrR1taG008/He+//37ONoQQrFmzBl1dXTAajVi+fDneeeednG3S6TSuuuoqtLS0wGw249RTT8XQ0FD1v0ZFZQopLoWh6BBSXKreh6KiorIbUpaRfemll3DFFVdgy5YtWLduHTiOw3HHHYd4PC5tc9ddd+EnP/kJ7rvvPmzduhUdHR049thjEY1GpW2uueYaPPXUU3j88cfxyiuvIBaLYfXq1eB5Xr5fprJbwws87tl+D5Y9vgwn/vlELHt8Ge7Zfg94Qb3HVFRUagdFCCGVfnhiYgJtbW146aWXcNRRR4EQgq6uLlxzzTX41re+BWDSa21vb8edd96Jyy67DOFwGK2trXj44Ydx9tlnA/i0Cftzzz2H448/fsbvjUQisNvtCIfDaqs7lbzcs/0ePLDrgWmvf3nel3H1IVfX4YhUVFRmC+XYoKrWZMPhMADA5XIBAD7++GN4vV4cd9xx0jZ6vR5HH300Xn31VQDAG2+8AZZlc7bp6urCgQceKG0zlXQ6jUgkkvNPRaUQKS6Fh999OO97j7z7CMLxMDKZDFiWBc/zEAQBVcw1VVRUVApScdN2QgiuvfZaLF26FAceeCAAwOv1AgDa29tztm1vb8fAwIC0jU6ng9PpnLaN+Pmp3HHHHbj11lsrPVSV3Yzx+DhSfP412CSfxMfjH6PT1CkZV0IIBEGQ5bspipL+0TSd9/+LvVfsMyoqKs1HxUb2yiuvxM6dO/HKK69Me2/qgEAImXGQKLbNTTfdhGuvvVb6OxKJoLe3t4KjVpnNJJNJBAIBpNIp6Gk90kJ62jYGjQEWyoJMJgOz2QybzQaGqfgxmEa20c73/1P/5nl+xs+oEwAVAGBZFrFYDBaLBVqttt6Ho1IiFY0uV111FZ555hm8/PLL6OnpkV7v6OgAMOmtdnZ2Sq+Pj49L3m1HRwcymQyCwWCONzs+Po4jjjgi7/fp9Xro9fpKDlVllsPzPILBIOLxOAwGA9xuN7gxDud85hw89P5D07Y/f//zsdecvUAIQTwex/j4ODiOg06ng81mg9ForMpoZBuhRkOdADQngiBg/fr12LJlC1iWhVarxeLFi7FixYqGvM9UcinLyBJCcNVVV+Gpp57Chg0bMHfu3Jz3586di46ODqxbtw4HH3wwACCTyeCll17CnXfeCQBYuHAhtFot1q1bh7POOgsAMDo6irfffht33XWXHL9JZZYjGshgMAiKouB0OuF2u5FIJOD1etHd3Y1v9H4DOp0Oj7z3CJJcEkbGiPP2Ow9XLLgCwOSgbrFYYLFYAEzep+FwGD6fDxRFwWq1wmq1QqPR1POnyoo6AaiMek8A1q9fj40bN0p/sywr/b1q1SpZfqOKcpSVXXz55Zfjf//3f/GXv/wF++yzj/S63W6H0WgEANx5552444478OCDD+Izn/kMfvjDH2LDhg14//33YbVaAQBf+9rX8Ne//hVr166Fy+XC9ddfD7/fjzfeeKOkQU3NLt494TgOfr8fqVQKZrMZTqcTGo0GhBD4fD6wLIvOzs6cwSvFpbDro12Yt8c8GBhDSd8jCAJisRgikQgEQYDBYIDNZoPBUNrnVWYX5U4AZvp/QRByJgniEJz9t/iPZVk89dRTecsbtVotbrjhBjV0XAfKsUFlebK/+tWvAADLly/Pef3BBx/EhRdeCAC44YYbkEwmcfnllyMYDOLwww/HCy+8IBlYAPjpT38KhmFw1llnIZlMYtWqVVi7du2s8hpU5IEQgmg0ilAoBI1GA5fLlZNYx/M8hoeHYbPZ0NraOu3zBsaADmMH9JrSlxtomobNZpMenlQqhXA4jPHxcek9i8XSkB6hijwIgjDtn5iJXui1Uv0ViqKg1WpB07T0T6PR5PwtvhYOhwvqB7Asi3//+99wOp2wWCwwm83qPdmAVFUnWy9UT3b2k8lk4Pf7kclkYLPZYLfbpw0gyWQSXq8XXV1dRdfsR0dH0draKkuCE8/ziEQiiMViIITAZDLBZrNBp9NVvW+V0sn2DmcyftmvlYoY2i1k/PK9Jvf6LyEE4+PjeOCBB8Bx3LT3RU8WAGKxGOLxOARBAE3TMJvNaoKUgijmyaqoKIkgCAiHw4hEItDpdHC73QWNVyAQQCKRQH9//4yzd51Oh3Q6LYuR1Wg0cDqdcDqdIIQgmUzC7/eDZVkwDAObzQaz2Vx0wN1dskRFI1iuR1gqUw1dtgHU6XR5jWKzJEIlEgmMj4/D6XRiyZIlOWuyIosXL5buH/GeBCafo3g8Dp/PJxlng8EAi8UCg8HQNOdgtqAaWZW6I5be8DwPh8OBvr6+ggOBIAgYGRmB0WjMyWwvhl6vRzqdhtlslvOwQVEUTCYTTCYTgMk143A4jGAwCADTSoQaMUs02xCW6hFWYgjzeX9TQ6bNZgiVgOd5eL1eaDQa9PX1gaZprFixAgCk+4ZhGBx++OHS61OhaVpK3AMmr3EqlUIsFoPP5wMAMAyjhphrhBouVqkL+UpvZvI00+k0RkZG0NHRISXalQLLsvD5fDllZUojZkBHIhGpROitt97C5s2bp227bNmyolmiSq8PlhoOzTaCu7MhVAJCCAKBAGKxGDo6OvIuf4gREJqmkclk8uYglIq4LzXEXBlquFilIckuvQEm5TjdbndJA3YoFEIkEkFfX1/ZCXIMw+Rd01KSqSVC8XgcW7duzbvt5s2b0dvbm+PZTd1XIQPIMExeo6h6J81Ddmi4v7+/4HZarVZapvB4PFV9p7gvOUPMu8sySLmoRlZFcaaW3nR1dZVsKAkhGB0dhVarRW9vb0UeVL29Lp7nMTY2VtDQcxwnGWWO46TyDhGxLpTnebAsC41GI/3L9v5nkwDD7kC+0HApUBQl5RnIJdJTTYi5EZdBGgnVyKoowkylN6XAsiyGh4fR2toq+3qq0oi/X2yi4XA4oNVqwbLstG21Wi3MZjOSyWTOYJdvgMrOqBUNMs/zSKfT0v8Xy6KlKCrHQE812KJXrKIcpYSGZ8LlciEYDEoqe3JDURSMRmPOsozoqY6MjOSEmLdt24ZNmzblbKeKZXyKamRVZGVq6U1PT09Fg3Y0GkUgEEBPT48sWcEURZWkoV0t2UlcVqsV3d3d0u9fvHhxwSzRrq4uAJPeTTQalQYyo9EIu90uZVlnG8lKyoZErzjbQGcymRwDXSyxSQxZFzLQSpSyzCZKDQ3PhJjMV0vyhZhDoRBee+21vNtv2bIFRx111G4fOlaNrErVlFN6MxOEEIyNjQFA0SzjchHDa0qoNrEsO9mYIJWC0WhEW1tb3oElX5aoGFYT0Wg0cDgccDgcUsguEAggk8lAo9FIQhiVnheKosAwTEUTF7E2NdtAcxwntQwU/xVLuprJQDd7dnGhdclKQ8PFMBqNSCaTZSUByol4rQotg4jnYmrHtd0N1ciqVEw5pTelwHEchoeH4XK5chTC5ECv1yOTychmZHmeRygUQiwWg1arLSkcTtM0Vq1ahaOOOgqxWAwURSGTyRQccKeG7DiOQzQaxdDQkCSEYbfba+YpiOu91Uygso0xx3HgOE4KdXMcV9SLFicIhQw0wzB1M9CF1iWXL18u3SeVhoYL4XQ6MTExUTcjK04CCyUWarVaKfFvd0Y1siplMbX0pr29vTKviOXBR1lorFpQWg3i8TgmJibQ3d2tiNHQ6/WIRqNV7SPfOmslEwsx7AYAAwMD0vrWTDAMkyOEkUgkJM3mUoUw6kk1XjQAaS06O9ydbaBL8aKLrUdX40UXEvEPh8NYtWpVVaHhQohr/LVYBsmGEIJwOIxQKISWlpaSxDJ2Z1QjqzIj1ZTeTNuXQBBZN4DYK8MgrABKS4M+xAnhUDv6+/sVGyzEcHElFFtnrZaWlhb4fD60tbWV9TmKomA2m6WEMJZlEYlEEAgEpPfk7pVbb7JFLMqlUMJYJpPJ+bsQxRLGCCHYsmVL3s+99957OOWUU8o+3lKxWCyIx+M18xjFybDd/unzKi53bN68GRzH5WQXq6hGVgUAl+GRiGRgsunA6D4tramm9KYQkXUDiK4flP4mrAD+NT+sJhOoHuVm4zRNlyzQAJS+zlotZrMZPp+vZG+2EFqtFm63G263e1qvXK1WK3XKalQvV2mUSBhjWRbJZBLBYDBv1jig/Lqkw+HA6Oio4kY2nU5jbGwMBoNh2mRYXAbZY489YDKZwHEcuru7FT2eZkI1srsxgkDw+rMf4a0XB8FlBDA6Gget7MW+y1sQiYQrLr3JJrv/J5/mEH1lOO92sU3DsK3sBaVVrhMTx3EIBoMFi+UrWWeVg0q92ULk65UbiUSkXrkWiwU2m03telUGxULdxcqzGIaRwthKTHA0Go2k8KXE/jmOw9jYGCiKQnd3d9F7hmEYtLW1YWhoSPbjaGZUI7sb8/qzH+GN5wekv7mMgDf+PoBUKoWln9tHCrElEomc8o6p/z8TkgpRjAfY/IktJCNg6IMBwKaVPqPT6aDVaqX/VprYIial5AtnURQlyzprNcjlzRZCp9OhpaUFwKe9cr1eL3ieh16vl3rl7q5ebrWI91O+dcklS5aAYRgMDg5Co9HA7XbLnuFus9kkmT+5EAQB4+PjYFkW7e3tJXv/6j00HdXI7qZwGR5vvTiY971/bRrHHott0OpzSys0Go0k/VdJyQVheYxoB0DyGFpKR6N7737JkxXXy8SQXDgcLpjBmG2Ixe4r2RRKSolGozjkkENkX2etBLm92UJM7ZWbTqcRDocxMTExoxCGSmGWL1+OcDiM9957L6/qkcPhkJYgxsbGYDKZ4HK5ZIkm2Gw2DA8Py2Jks4UyWltbpeYXM5E9QaQoSrEJYzOiGtndlEQkAy6T36vkWQKHpQW2FnlLAyitBpal3TlrsiKWI7tzQsUajWaa4sxUCCFgWRYsyyKTySCZTIJl2ZwyEI7j8oryA8A777yDk08+uSEyIJX2Zguh1+slwy4IQo4QhsFggN1ul7XsZLYSDAZxzDHH4JRTTimo36vVatHe3i5lho+OjoIQAofDUVXtszjR5Xm+KqMtJs65XK6ys6HT6bTk7Yr1u82m0qYUqpHdTZlMcqLzGlpGR8NkU6YJue3YyYc3tmkYJCOA0tGwHNktvV4OYs2mTqcr+EAHg8GmKZavlTdbCJqmYbfbYbfbQQhBOp1GKBRCOp2GRqOB1WqFxWJRPZQpcByHRCIhheRnup+yM8NF1SSPx1OVkIvYEcblcpX92WQyifHxcZjN5ooz/LN1lE0mE2KxmGpk/4tqZHdTGJ0Gcw+149+bgtPeO2hlb06WsZxQNAX78XNgW9mL4X8Pon2PLjAGZQw6AMmjKJSU0kjJP/XyZvNBURQMBoO0fijKPQ4PD4MQMk3ucXfG6/VWrCFM0zRcLhdcLleOJKnVaoXD4Sj5PrBYLBgcHCzLyGYyGYyNjUnNN6q55zKZjBSuNhgMUlMBFdXI7raEw2EsOL4HNpstJ7t4v2Xt6F0kv/TgVCitBuZOO5JsGlYFjWyxpJTFixcjkUggFArB7XY3xMy73t5sIabKPSaTSfj9fqkrkN1ub2ghDKVIJBJgGEaWyYZOp0NnZycIIYjFYhgeHgZFUXA6nTCZTEXPrViixHHcjLXRYlcoQRDQ0dEhy3JJJpPJ0ddW+RTVyO6GcByHUCiEvr4+tJ3WikUnzsmpkw0EAvD5fFL4SylMJhNCoZDsEopTOfroo4smpQiCAL/fD5/PB6fTCavVWreBopG82UJQFAWTySQlxXAch0gkgmAwCEKIJITRCGvdSkIIwfj4OPr6+mTdL0VRUgIaz/PS82gwGOByuQqeV6fTiWAwWLCZOyEEPp8PyWQSbW1tsmY55yshqrUSVaNCkXIq9BuEcrrSq0xncHBwxrR8r9cLk8mk6PklhGBwcFD2QWoq4gCl1+sxMDCA9vb2vIadEIJgMIhoNAqbzQaHw1GXQSIejyMejzecN1sKYlKPmA2u1Wphs9lm9MSaEZ/PB51OV7MxKFt5zGazwW6355xTsZl7V1dXTvLVVBlEJYQrBgcH0dvbK/09MTEBi8VSN11lpSnHBqme7G5GOByG0WicMbzV3t6OoaEhaLVaxR6UWgy6YuhN9Mrb29sLJkJRFAWXywWn04lIJAKPxwOz2QyXy1VTr7IZvNlCFJJ79Pv9s0rucWqyUy0wGo3o7u6WjObU2ltCCHbs2IFHHnlEitgsXLgQe++9t9RaT4lnLp+fVu8OQY1Ec9/pKmWRHSaeCVHhxePxKCbaD0yu9VVbelCMcDgMh8Mh/W0ymeD1eotmgFIUJWXZxmIxDA0NQa/Xo6WlpWaJUo26NlsuU+UeY7EYxsbGwPM8dDod7HZ7UwphVJPsVC0URUnr49m1tzt37sSbb74pbceyLLZs2QKGYXDMMccodjxixCIbk8kkCbzs7jTXNFmlKkZHR9HZ2VnygEbTNHp6ejA0NFSSslMlmEwmJBIJRfZNCEEoFMop0hdl6ErFYrGgr68PNpsNIyMjGB0dLegJy4nZbEYymSzrWBsdca2xu7sbfX19cLlciMViGBwcxODgIEKhkGL3mZzImexULWLtbWdnJ956662827z22msFtZXlILt8R0TMdVBRPdndhlLDxFNhGAZdXV0YGhpSRG5QyeQncV0qX0JGuUkZRqMRvb29SKfT8Hq9oCgKra2tig60s8WbLYROp5OSdES5x9HRUQiCAL1eL3m5jYRSyU6VHEcqlUIsFkMqlUIkEqlbPXg6nS54ndTkJ9XI7haUEybOhxgqHRkZkb27RjUt6GYiEAjkJGNkfyfLshUZSL1ej56eHrAsi4mJCfA8j5aWFkXWnpp5bbZcpso9ioZjfHy8oeQe/X5/zdfoAUga4rFYDJlMRqpjNpvNaGlpkUK2+TxWpZunp9PpvJKO4vE0gsdfT1QjuxtQbpg4H2azGZlMBuPj47J6VkrNcpPJJAwGQ97BUAxRV/Pwa7VadHV1ged5TExMYGJiQpFa29nuzRYinxCGKPdoNBphs9lqLvdYy2QnjuMQj8cRi8XA8zxompb0jvPdt1qtFgcccAB27Ngx7T2lm6fzPJ83kU2O52w2oBrZWU6lYeJ8OJ1OjI2NIRQK5SQTVYsSyU8+nw9dXV153zMajfD5fLL8Bo1Gg46ODsVqbXcnb7YQU4UwUqkUgsEgMpkMNBoNbDZbVdq/paJUshMhBJlMBrFYDIlEAoQQMAwDs9mMjo6Okp6LSCSCJUuWwGq1YsuWLWBZFgzDYMmSJXVrnm4ymTA+Pi7rWNGMqEZ2FlNtmDgfbW1tGB4ehk6nK7lDx0yIM1651mXFwbfQ4KTVamVPXqJpGq2trVKtrcfjka3Wdnf1ZvNBUVRO4wie5xGJRDA0NARCCEwmE+x2u+yem5zJTqJilrieCkwuQ1gsFrhcrrLvF1Gwor+/H6tWrcJRRx2FWCyGQCCAPfbYo25romIv3d0d1cjOYuQIE08lu7Sns7NTlkFH7uSniYmJuhkkJWptVW+2MBqNBk6nE06nUxLC8Pl8kidns9mqlnusNtmJ53kp9MtxnDRRsNlsaG1trfr5nPqca7VaKclJ7j6zU1HvyZlRjewsRc4w8VQoikJPTw88Hg/6+vqqDvPKmfzEcRwIITN6MqXqvFaK3LW2qjc7M1OFMDiOQzgcRiAQADBZjlWOEIaYlZtOp8uaJImfi8fjEAQBGo0GZrMZra2tsnvYkUgEer0+7/q0w+GAx+NR1Mhmt7jLh9LPWTOw+/7yWYwSYeKpaDQadHd3Y3BwsGolGTk97YmJiYLardmIIepaSOJZLBZYLBYkk0mMjIxAo9Ggra2trIFH9WbLh2GYHCGMeDyO8fFxKRPXbrfDaDROu/8EQcD69evzrm1OPfdTS2nECZ7FYkFXV5ei1yo7TJwPsRVkvjpWuZhp3yaTCclkUnF98kZGNbKzECXCxPnQ6XTSGm13d3dV3ydH8pMgCGBZtqQBxWg0IhQK1VT7emqtLTC5xl1qtEH1ZiuHoihpsgNMrttHIhH4fD7pPZvNBo1Gg/Xr1+d0beI4Tvp7xYoVeUtpLBYLWlpaarr+Wcpz7na7iyYBVkt2i7t8iM+ZamRVZg1KhonzYTKZwLIsxsfH0d7eXtV+qk1+8vl8cLvdJW2r0+mQyWQq/q5qqLTWVvVm5UOn00mlOIIgIB6Pw+v1Ip1OY/PmzXk/s3nzZuy5556w2WwFS2lqRbEwcTY6nQ4cxyl2z2S3uCv0/fV6zhoF9UmdRYhh4lINjVzY7XbQNI1gcHoD+FKpVl5RzNgstU6Voqi8wua1RKy17erqQjgchsfjQTweL/oZ0ZtVKR9CCDiOQzqdRjweRyQSkdrIxWIxKfRbKPOc4zjY7fa6G1gxTFxqva7D4UAoFFLkWGZSdNrd1Z4A1ZOdVdQqTJyP1tZWDA8PV6wuU23yUzAYLFs2TtRXrbdXWE6trerNTg7sPM+D4zhwHCf9f/Z/C6HRaMAwjPRfvV4Pk8kkvdba2gqGYfIaWqWVk0ql3OfcarXC4/HA5XIpfGT5oShqt75fVSM7S6h1mDgfXV1d8Hg80Gq1ZSdaVFtiEYlECiaAFEL0nhth4ASm19oODAzAbrdPq7WdLWuzgiAUNZaFBOYpippmLA0GAxiGAcMwoGm6ovspGo3C7/fj0EMPzRsyVlo5qRRKDRNnI64by916rtRIkNj2Tm41tGZBNbKzgFpkE5dCdmlPb29v2Wn7lSY/RSIR2Gy2sgdWo9EoNRFoJGaqtRW9WT7NgsR5aKxaUNratOCbCiEkx1jmM5qFBmOKoiTDKBpLnU4nvUZRVE2iMjzPY3R0FDqdDv39/ejr60MymcQ777zTEMpJ2cdZLJu4GG63G16vFz09PbIdT74Wd/kwmUyIxWKqkVVpXuoZJp6KRqOR2uP19fWVFSKqNPmp0gmGwWBo6PXNqbW2g4ODMBgMcLvcMO/KYPSB1wCOgNLSsCzthu3YflB0ZfeAGILNF3oV/1vIWNI0nWMotVqt5FlqNJqGDhOK67KdnZ2Sd0hRFA455BCcfPLJOdrB9f4d1TznDMNI11gu+dJSS4Ma/TlTGtXINjmNECaeitjjcnh4GD09PSUPCpUoP8XjcZhMpooGnkZIfiqV7Frb0affBfXGpw2xCSsgun4QAGA7rl/W9crs1xphEicXHMdhdHQURqNxWp03y7LQarWSchIhRBJ1kFNfuxwqCRNPxeVyIRAIlFRHXgrFWtxlM5vum0pQjWwT0yhh4nwYjUbY7XZ4vV50dnaW9JlKkp/8fn9VITDR0DbLQGBgdKB3RpFvahDZOIToPjQ0eq3s65WziWAwiEgkUlAWVFx+EKEoCu3t7fB6vbK3eiyFasLE2YjLDHLd74Va3BWimZ4zOVGNbBPTSGHifNhsNmQyGfj9/pLKisr9HalUClqttqownpiUIVezAyVJp9MIfDIOwhZICOIITDDA0d5S99BmI8KyLEZHR2GxWNDX11fwfkskEtMycQ0GAyiKkj15qBTkfM7NZjPi8bgseQiFWtzlw2AwIJVK1fzcNQLqk9ikNGKYOB8tLS1Ip9OIRqMlbS8mP5WCz+erOvQlGtlGRRT68Hg8k2VK3S2gtPkfW0pHg7ZoMTw8DI/Hg0AgoHZBwaQH5ff7JWNVrNMNIaSgx9XR0YGxsbGaLjHIESbORgwZ15pGf86URPVkm5BGDhPno7OzE4ODg1JCTDFKTX5iWVbKTq0Go9FYlYiGEojXNx6PS+uC2eU6lqXd0hpsNpYju2FvccKByXXEWCyG0dFRCIIAs9kMu92+2wm1ZzIZjI6Owm63l/S8FPNUaZqGw+FAMBisSc2pXGHibGiahkajkdada4XJZEI4HJ55w1nI7vXEzRIaPUw8FbG0Z2BgAD09PUUf7lKTn0ptBDAToiBFveF5XjKsYvs2t9ud9xrbjp0cdGObhkEyAghDwbasR3odmDznVqsVVqtVEscfGxsDz/MwGo1wOBx1r/lUEkIIfD4fUqkUuru7S55cRCKRoqImdru9ZklQSj3nbrcbfr+/qgb05YpLNMpzVg9UI9tkNEuYeCo0TaO3t1fq2lPoAS0l+UksNZHrHNQr+UkQBITDYUSjUclLKqVpN0VTsB8/B7aVveCjLEZCXtjmFu6ElC2OL8pPTkxMgOM4GAwGOByOprufipFKpeD1euFyucqeiGUymaKh2VolQckdJs7GYDAgnU5Xdc/P1OKuELtj8pNqZJuIZgsTT4VhGHR2dmJoaAi9vb15H7ZSHkCfz1eybmsp6PX6kssRqkVUp4pEIgAmPaNC52ImKK0GjEsDC2yIxWIllT5RFAWTySQleiWTSfj9frAsC51OB6fTqVhbNKURm6tzHIfe3t6yPc1SvTOlk6CUCBNPxWazVdXQvZL2eVqtVrrPdidUI9vgEJYHH2WhsWox6m2uMHE+DAYDnE4nRkdHC7bfKqb8JAgCUqlUVR1/piImZShlZMX10VAoBEIIbDYburu7ZcsAFkulKulgZDQaJUORTqcRDAaRyWSg1WrhcDiaJhs0mUxibGwMLS0tFWfOlqP+1dHRAY/HU3Uv5XzUYjmo2obuM7W4y4eYb6EaWZWGgAgEkXUDiL0yPFmyoaXALHRBe6p8smj1wmq1gmXZgh5pseSnQCAge5chk8kEr9dbdoOBYojroKFQCIIgSE28lVjHEycl1Ybi9Hq9tE6XyWQQCoUwMTEBhmEkg9toEzxBEKSM33IVxqYSjUZLXqdUKglKyTBxNtU2dJ+pxV0+TCYTxsfH4XA4yv6+ZkY1sg1KZN1AbgYpS8Bv8SNiHID9+Dl1Oy65cLlc8Hq90wr/gcLJT6JHKGeoGCivbKgY4npnMBgEz/Mwm83o7OysiUqQuM4mlzeu0+mkjGaWZREKheDz+aDRaGC322E2m+tucOPxOCYmJtDW1iZLnXO5koNyJ0HVIkycTTUN3SuZ0DEMs1uWlKlGtgEhLI/YK8N534ttGoZtZW/dBOHlpL29HUNDQ9BqtTlhyULJT+FwWNFZcKWeYCqVksKsJpMJ7e3tNS+VsdvtCAaDVWWMFkKr1UoJRBzHIRwOIxAIgKZp2O12WCyWmhpcQRDg9XpB07Rs4dpKSlrkToKqddWA0g3dVSZRjWwDwkfZgqo+JCNg7ONR6FvN0Ov10Ol0TasrS1EUuru74fF40N3dLQ1y+X4LIQShUEixWb5OpysrKSN7/VKv18Ptdtd1rUmv1yOTySj+PQzDwO12w+12g+d5hMNhDA4OgqIo2Gw2WK1WRQdssR1de3u7rOvF+SIqpSBXElStwsRTERu6lxPyrkaMQ6PRgOO43apee/f5pU3EZOsyOq+hpXQ0HF1usIRDMplEOBye1mBaq9VCp9NBp9NBr9dLbcMaEZqm0dPTg8HBQfT19Ulht6nJT2JSilK/w2g0IhGJg6aFgq3jWJZFMBhEMpmEXq9vuExchmFqKjKg0WjgcrngcrkgCAIikQiGhycjMFarFTabTTaDO7Udndz3QT4pxVKpNgmq1mHibCpp6F5qi7t8mEwmJJPJipL0mhXVyDYglFZTVNXHYDGi0MobIQQcxyGdTiOTySAWi4Fl2ZxtxL6d2f/qaYQZhkFXV5fUHk8sM8lOfgoEAujt7VXk+4lAILzqR3zTCGJTWsfxAo9gMIhEIpFXfamRsNvtiEQisieGlYKYCORwOCAIAqLRKIaHh0EIgcViKWndMjuTPnuSk68dnZwUk1IshWqToOopLlNJQ/dKk6WAyclsuZ22mh3VyDYotmP7IQgC4ptGJnuG6mhYjuzOUfXJB0VRUpuufIg9JTOZDNLpNBKJBFiWzQkB0TQthaLFf0qv2ej1erS0tGBkZATd3d05yU9ieY1SxxBZN4D4S5+ugYut4yKRMOgjWuBwONDS0tKw0QARk8lUcjMGJRHXau12e0nyjlMz6cVJjmlFN7xj3rzt6OREjnrXSpOg6hUmzqbchu7VJNjpdLqaLGs0EqqRbVAomgI53AHb4W7ExsNo7W+H1lj9gyjq/TIMUzAjUzTCmUwGkUgEmUwmxwiL6f9iOFpcF64Ws9mMTCaD8fFxtLa2Ih1Pgguk4AuPo6tPmdIlPs0iWiDJjN4VQ9fpBzZNkhlFUaAoqqESWUqRd0z8czgnaiNOcsLhMDpO30/xte6ZpBRLoZIkqHqGibMpt6F7uS3usmn0iaoSqEa2gYnH4+jp6QEHHmmehRa1me1qNJockYKpCIIgGeF4PI5AIJCjSyp609mGuNTkLKfTCe+oF+PPvg/qdR+83CAohkJsGYHt2H5QdOkPqSAIYFkWmUwm57/ZUFEOKJJkxkdZMK7mMLLA5BpbNBqteBBUknzyjuOjYxA2DiHfVdW8HYP2DOXP/UxSiqVSbhJUI2mQl9PQvZwWd/lotImg0qhGtkHJXicyGAxIJBKy9ICUA5qmYTAYCoaMCCGSEU6lUohEIkWTs3Q6HbRarTTYGHemEH114tOBlyOSpyPWCJdkQLOMvV6vh8ViyfkeYHIdcEQ7kj+bm6EwGh6DnXHCarU2xGA4EzabDcPDww1pZLMR1911Fhpe7uO829RikiP3YF9qElQjhImzkbuhezFEhTWz2azo9zQKqpFtUOLxuHQT6vX6hmvHVgyKoooOIEWTszgBKODZRDYOIbI3BTB0SQa0pGMtkmRmXdYDa3+vVKYidsdpROUjEZqmq07kqSUZhgdhKFDc9LIQSkdDY1U2U7ocKcVSKCUJqlHCxFORs6F7MUwmE2KxmGpkVepLJBKRQjcajWZWtYkqlpzFBVLwcp78n+MIOuztYFzyagxPbR0HLQXr0h4pPO10OuF0OsFxHILBICYmJqDX6+FyuRpSh1XMzG7UQUxskhAKhWAwGGBd2o3YhqFp21mO7FZ8PbwcKcVSmSkJqpHCxNm4XC4MDw8XNbJyeP4GgwE+n6+qfTQTqpFtUKqpRWtmNFYtoKXzrpMq5dlkt47jIhmMhLyw7zln2nYMw0gTn3Q6Db/fj0wmA7PZDKfTWRP5xFKw2+2YmJhoOCMrCAICgQDi8ThsNptUrhU6PARNMgnhzSBIRgClo8EfaIHh6E7Fj6lcKcVSKJYE1Whh4mxomgZN00XFIiptcZdNo00ulEY1sg1IIUGBZgkBVgNPEfAHmqF5M5rzOgFgPLxdUc+G0mqgdRuhZ40z1gLq9Xp0dnaCEIJEIiGVqNhsNllFGCpBq9VOWwOvJ2IzCJZl4XK5crSnWZZFOBJG3+n7AasFDH0wgO69+yHQkNY2lTqXSgp3ZCdBGRgd+CgLmOiGDBNnI+oZF/Luq6mRncruMJ4BqpFtSMLh8LTEFbG+rBFnwHLBcRwGBwfRffr+SDpGEH1lCGD/WyN8RBfC+1AwJhKyiMEXw+12Y2JioiThdIqiYDabYTabpTDo0NCQtDZXLyF98X6pZzg7mUzC5/OBpmm0tLRMu3cJIRgeHkZPT8/kOdJqANukEIUGk0lEIyMjJddvlkulUoql0t7WjuE/vw3NrthkYh1DwXZEJ9AH5E06aABmauheSYu7Qt+TSqWappViNahGtgFJJBLTBAXEm3K2GlnRwPb09EwmNB0/B/zBVmjSBLYOJyitBnZCMDg4CLfbrWgoVPQEy51pUxQliTDwPI9QKIRAICApRdWiKbyI3W5HKBSquTrV1PXWYu39xsbG4Ha7C4YmjUajogIb1UgplkLsxUHQ2yOQUro4gsTLI5OdjBq4k1axhu5yTdzEDOPdwcjuHoVKTYQgCJKoQDbiDHM2MtXASjAUNE6DFCKmKAq9vb0IBAKIxWKKHpM40FSKRqOB2+1GX18f3G43QqEQBgYGMD4+XpNQrjgpqxWCIMDn88Hj8UAQBPT19aG9vb2ggY3FYiCEzCiv53K5kEqlkEgkZD1epTOwZ+qkRdjGbfnmcDgQDofzvifXOROT83YHVCPbYBQqKSjU/q3ZKWhgMflAT12PoygKPT09CIVCiEZz123lxG63FxxoykWn06GjowP9/f2wWCwYGxuDx+NBMBhULGucoijQNK14/06WZTE6OoqhoSEYDAb09/fD6XQWHYg5jiu67jeVzs5O2ScnSntRM3XSio1HqupmoyRi9r+S4w1N07OqYqIYari4wShUUjAbEwSKGVig8KxZbJE3MjICQogi62oURUGj0cieHGMymWAymUAIkUT0gUnvQe4uQ2LDgGolA/Mx03prIQghGBkZQVdXV8m/laZpdHV1YXh4WMpIrhalzosIr0fB+l9oaaRoFqGhIRBCoNFopHX9RqkoyNfQXe5JAUVRu0Xyk2pkGwwlSgoakZkMLFA8NEVRFLq6uqSsXiWaubvdbgQCAbS3t8u+b7H/qs1mgyAICIVCGBwcBMMwkuBFtVgsFgwNDclmTMpZby2E3++HzWYre11Pp9PB5XJhbGxMlrpWpZIICSGYmJhAOp2G5ciunMYTItal3bB3frpWzvM84vG4lIENTIb7zWYzjEZjXTLV8zV0l7usUKvVltXDuVkp++q9/PLLOOWUU6SZ6NNPP53zfiwWw5VXXomenh4YjUbst99++NWvfpWzTTqdxlVXXYWWlhaYzWaceuqpGBqaXoy+uzFTeryo+dnslGJggZnXfyiKQmdnJ5LJpCKKWOK6ptJhPZqm4XK50NfXh9bWVkSjUQwMDGBsbGyaVGQ5iJ5Ctcdf7nprIZLJJNLpdMUTIlHaspq1ckB+KUWReDyOgYEBmEwm9Pb2wnH8XFhX9ILS/fe7tDTIQvu0TloajQY2mw2dnZ3o6+tDb28vbDYbUqkURkZG4PF4MDQ0hEAgMK1Zh5KIDd1F5CzfAXafddmyPdl4PI6DDjoIF110ET772c9Oe/8b3/gG1q9fj0ceeQRz5szBCy+8gMsvvxxdXV047bTTAADXXHMNnn32WTz++ONwu9247rrrsHr1arzxxhu7hRdXiHylO9mIg77SJSxKUqqBBUpLshAN7djYmCJZqLWSmhPRarVSRnAqlcLExARYloXFYoHD4Sj7+bBYLIjFYhX17yxW31ougiDA6/VWXSPa1tYGj8cDg8FQsQckt5Si2FBeq9XmaBZni5yIfXL94SAi0fyZuyKiXnl2NrogCEgkEpKhBSa9TbPZDJPJpMi4ObWhezUt7vJhMpkwPj6uSBSqoSBVAIA89dRTOa8dcMAB5Lbbbst57ZBDDiHf+c53CCGEhEIhotVqyeOPPy69Pzw8TGiaJn//+99L+t5wOEwAkHA4XM3hNxwDAwNEEISC70ciERIIBGp4RPLCsiz56KOPSCaTKWn70dHRkrclhJCxsTEyMTFR6eHlheM44vF4ZN1nuQiCQKLRKBkcHCQej4eEw+Gi90k2HMeRoaGhsr4vkUgQj8dDhoaGSCqVquSQpzE4OEiSyeSM25VyrjmOIx999BHheb6iYxkaGiIcx1X02WwEQSB+v5988sknJZ8nQRDIJ598QliWrfr70+k0CQQCZHBwkAwMDBCPx0N8Ph9JJpMl3x8z4fV6SSKRIIRMjtNyHHc29X62KqUcGyR7zGTp0qV45plnMDw8DEII1q9fjw8++ADHH388AOCNN94Ay7I47rjjpM90dXXhwAMPxKuvvpp3n+l0GpFIJOffbIPnedA0XdRzq3VZhpyU48GKkDKTItra2qQ1MbkQPQSls3SLIbaI6+npQXd3N3iex+DgIIaHh5FIJIqGDzUaTUnHTghBOBzGwMAAIpEIurq60N3dLUt4MBgMFu3aVC4ajUYSqqgEOfIe0uk0PB4PKIpCX19fyedJjLyMjo5W9f3ApCfrdDrR09ODvr4+dHd3w2AwSIIoHo8HIyMjebtglYrb7Ybf7wdQfYu73RXZz9g999yDL3/5y+jp6QHDMKBpGr/97W+xdOlSAIDX65VujmxErc983HHHHbj11lvlPtSGIhqNzhjSYximoeTySqUSAwtUVpPX2toKv9+PsbEx2RKWXC4XgsFgVeFSuaBpuuyGBQaDAcloAlp2Uvs5W5qykJ6wXGQyGUSjUfT29sq2T6ByoYpqs8UFQcD4+Dh4nkdPT09Fxlqn08FkMiEUCskaKqVpWspSFmFZFvF4HGNjY+B5XgpFi0lVM11rsaE7l8oAERaE5WWVNtVoNEW1kmcDihjZLVu24JlnnkF/fz9efvllXH755ejs7MQxxxxT8HPFBtSbbroJ1157rfR3JBKR/aGtN9FodJqY+FSaMdW9UgMLfCrMUS5iVrDX65UlE9VkMsHn8zWEkc2mWMMCh8MxOUAKBMy2KPyvfgBwBJSWhmVpN4zLu+AP+GVZby0E+W+5Tm9vryL3rtg1JlGG1GY1UorRaBR+vx9tbW1V50W4XC54PB5YLBZFDYxWq4XD4ZCMOSEEqVQK8Xgcfr9/xhIiIhAY30rBu+V1UBzBiHYUlqXdUoeqajGZTEgmkxXlDDQLsl7dZDKJb3/723jqqadw8sknAwDmz5+PHTt24O6778YxxxyDjo4OZDIZBIPBHG92fHwcRxxxRN79NmrXCrkg/80ALTXjsRIPrx5UY2CB6n6n6H2Ojo6io6OjqvMlzv4bWQZuasMCr9cLQRCg35EA++qn4XPCCoisH0Q0GkXrKXsr+lx5vV60trYqmszY2dkJj8cjRc5mohIpRZZl4fV6JbENOZ49sQRtZGQEfX19Ve+vnO81Go0593GxEiJ24zhSr3waYSSsIPVelkMa0mg0IhQKzWojK+uaLMuyYFl2mrHI7oe6cOFCaLVarFu3Tnp/dHQUb7/9dkEjO9spRyibYZi6rg+WSrUGVqSaAc3pdMJkMkmiFdUgeseNjtiwoKenBz0dXWC3Tu/bSQGgdkago5XzoKLRqBS+VJJsoYqZrrE4mS31niKEwOfzSRO11tZWWSe3Wq0WFotFkfKzcihUQpSMJhSXhhQbWcxmyn7KYrEYPvzwQ+nvjz/+GDt27JDq/I4++mh885vfhNFoRH9/P1566SX8/ve/x09+8hMAkyo0l1xyCa677jq43W64XC5cf/31mDdvXtFw8mwmHA6XPLvW6/VIpVI1KympBLkMrBzY7XZQFIXh4WF0d3dXPEiKkxulaiyVQIhxAJvf8JCMAD7KgnHJ72VyHAe/31+zlm6lClWUE4lIJpMYGxtTLJQu4nQ6MTg4CIvFUvdnRUSM3DA6wJtPsQry3T/NEJGrlrKN7LZt27BixQrpb3Gt9IILLsDatWvx+OOP46abbsIXv/hFqXfi7bffjq9+9avSZ37605+CYRicddZZSCaTWLVqFdauXbvb1siW09nCYDAgkUg0rJFtJAMrYrPZQFEUhoaGPm2rVgGinrGScnxyQoxUQWk/SjeZBCX7d/63fV01E5pKsFqtSCQSRddcS5FSFOt5xaxhpSdUYraxGDZuJKMzmSRH59VglvP+EUV2mmXyWi4UqTaOVgfENkzhcFjRfpC1gOM4jI2NzZj0JMLzPLxeb8nb1xK5Dezg4KCsCW6xWAyBQKDiRBxCiNRIvNGJxWLw+Xywvycg9tLwtPal1hW9irRbGx8fh8FgqPi5rOaai9ens7Mz76TV4/EUXf8Mh8MIBoPo6OioaVtCAFKzCCVa+lVD+P8+kdZgs5Hz/vH7/dIacLNQjg2anVOHJqLcbMfs9e1GohE92KlYLBa0tLTA4/FUtEYrdidp5DUkQgjGxsYQjUbR398P88oeUIscWdJ+FHRHtE2T9pODRCIBjuPqNvEVOzQNDw9Pe0aKeUqZTAYDAwPgOA79/f01N7DAZNg4Ho833L1lO7Yf1hW9IMx/p2laGvRhLlnvn9kurzh7i5OahHg8jp6ennofRlU0g4EVMZlMkjRfb29v2SEqsTi/s7NToSOsHI7jMDw8DKfTKRk6n9+HtlP3gfY0DfgoC9rCwDMyhGmubZXwPI+xsTHMmTNH3h2XSbZQRfZzlU9KURQuSafT6O7urnutptydhuRAlIaM7E2hw94OjVWLIe8IeIEHI1PinMFggM83PUFvtqB6snWk3GzHqZ9tBJrJwIoYjUa0t7dLgvfloNfrayrSXiqxWAxDQ0Po6uqSDKwgCOA4DjqdDpRWA8ZlAK2brK2VUxULQNnt65QkW6gCmKx6GB4ezvFQp4r519vAApPJdQ6HQzruRoHjOGj0WjAuAyitpqhwUCWIjSxmK/W/s3ZjEolEResQYtp7vWuHm9HAihgMBqnGsre3t6ykO6vVimg02hD5AIQQjI+PQxCEaTWchdSQzGYzAoGAbEo7gUAAZrO57vdjNi6XC4ODg9i2bRu2bdsmKT0dfvjh2HvvyfpguWpe5cRut2NwcLCsZEilSaVSORMUnU4HjUYz7fVqmM29ZVVPto5UmrjVCBrGzWxgRUQBh8HBwbJqj6e2AKsXHMfB4/HAaDSis7MzZ4ASRSkKTeLk8kbS6TTi8XjZAg+14P3338fmzZslgQWWZfHKK6/gX//6F9rb2xt2QBdFKhrFu8vXfaetrQ1jY2OyfUcjjGlKoRrZOpDiUhiKDiGejldkoAwGA9LptAJHVhqzwcCK6PV6dHV1wePxlKwLTdM0aJquq450vvBwNjO1TRS9kWQyWfExCIIghYkbDZZl8dprr+V9b+vWrVX16VUajUYDl8vVMOuUqVRqWpRClGKUq1mL0Wis6l5sZFQjW0N4gcc92+/BsseX4cQ/n4hzXzkX92y/B7xQnnKKTqerm5GdTQZWRKfToaenB4ODgyUPvi6Xqy4KUFOzhwtdg5mMLDDpzY6Pj1d8LKOjoxU1b68FsVis4LVkWRaxWKzGR1QeNpsN6XS6rpNpEUEQ8l5jUQVNDo97NmcYq0a2hvxixy/wwK4HkOInwyJpIY0Hdj2AX+z4RVn7qVeYazYaWBGtVove3l4MDQ2VZGhFYfNaUiw8nE08HofJZJrxPqFpGlarFeFwuOxjCYfDUjeZRiOdTiMUChVcbxblDBsdsSVeo4SNp0JRFNxutyweN03TDVmaKAeqka0RKS6Fh999OO97j7z3CFJcY69HzGYDK8IwjGRoS6lXrOXse6bwcDbltH9zOp0IBoNlDeQsyyIUCjVcV6J0Oo3BwUH4/X709vZiyZIlebebN28evF5vQ4eMgcmQrNvtlj0TvBxmui9EpS059NRna5axamRrhC/pkzzYqSS5JHzJ8maDohRZLdgdDKwIwzDo6+vDyMjIjKG6WoSMSw0Pi6TTaWi12pLrfymKKqukp16yicXINq6dnZ3o6uoCwzBYsWIFli1bJp0zrVaLZcuWYfXq1ejo6MDY2JjUrahRsVqtYFm2bklBpWQ5t7e3y5IEpdVqG37iUwlqCU+NaDG2wKAx5DW0RsaIFmN5XoGYjad0uG53MrAiGo0Gvb29GBwcLCqxp9FoQAhRTHc1n7jETPh8PrS1tZX1PeWU9IyNjcHtdjdEXWk6ncb4+Dg0Gg06OzunHRNN01i1ahWOOuoohEIhJJNJSVaRpmn09PQgmUxKAv0ul6thJg7ZdHR0YHBwsC4lR6WU6RgMBhBCkE6nqyrjEiNDjVK6JBeqJ1sjDIwB5+9/ft73ztvvPBiY8urN9Hq94kkRu6OBFdFoNOjr68PY2FjRtVeHw6FIq7JywsMiHMdBEISKrlUpJT2xWAyEkLr3/izkuRZCq9VKze2nYjQa0dfXB4ZhMDAw0JAJURqNBi0tLVUlqVVKvvKdfMjhzc7W5CfVyNaQKxZcgS/P+zL09ORsz8gYceG+F+KM9jPK3pfSdWW7s4EVoWkavb29GB8fL/jwWywWWQfmcsPD2fh8voLGZCZmKunhOA4+n69oKzmlKde4TsVgMOT9fRRFwW63o7+/H8lkEh6PpyGyerOxWCzgOK7myXbpdLokz5JhGBgMhqqehWbplV0u9Y/57EZoaA2uPuRqHK0/Gs4eJ9rMbTAwBvj9fgQCgbIK+hmGUaxOUzWwn0LTNPr6+jA4OAi32z1N3IGiKCmqUK3iUSXhYRFBEEr2OgrR3t4uhSWzIYTURDaRZVlEo1FJnUlkprBwqYiShYV6yorr06IOMyEE7e3tDREaByAplNU6bFzqd7W2tmJgYABms7khw+71ojHunt0IQgi0lBZ99k9bbrndbgwPDyORSJS8xqrUTawa2OlQFCVlHRNCppV/iMIB1YgyiK3puru7KzrvwWCwatWl7JIem8kCPspCY9XCHw7CZrMptlYmCALWr1+PLVu2SAZ28eLFOOKII+Dz+ao2riI6na6kxBqNRoOuri6k02mMjIzAYDCgtbW17oaDpmm0trbO2Jy+XlAUJS2fVHovajQa2eQ+G4XZ80uahHQ6nbewu6urCwMDA2UPsnLqfaoGtjDZbdSmrkvqdDpwHFfRtSimPVzOPsTwcrU47A4MPbkLsV2xyWbdWho4yIruM+dVve9CrF+/Hhs3bpT+ZlkWGzduRCwWw8knnyzrgKvX60vW3NXr9ejr60MsFsPAwIAUYainsTWbzQiHw2VNyCuF5/myE/rsdjsGBgbgcDgqSgYU68/rve4vJ+qabI2Jx/NLKRbrhVkIOdcwGs3ANmK9HEVR6O7uRiQSmSYnZ7PZypaYK1VcYibEZgVyDP7Rf3hAb49MGlgAYAVgWxiRdQNV71vsCpROp5FIJBCLxeD3+7F58+a827/99tuy3weV6E5bLBb09/eD53l4PJ66J+d0dHRIEzMlqWQJhKIotLW1VZykZTQa635+5Ub1ZGtMLBYruCbEMAza29un9cIshDgrr1a9ptEMrEi9w3P5oCgKXV1dGB0dhSAIcDgcAD7tnjKTlKFIteHhbILBoFSaUg2E5RF7ZTjve9FXhsEsbgHRTHo4PM9DEISc/5/JIFIUBY1GA5qmodFooNFopEbv+RDlD51OZ9W/TURsVVguFEXB5XLB4XBgfHwcfr8fHR0ddXleaJqWBPqV7GtcaZcdsc3g1LX1UhA7jM0mVCNbY1iWLToQG41GmM1mTExMzJgpajAYkEgkqjKyjWpgBUFoSCMLTA64nZ2d8Hq9IITA6XRKBmSmgUWO8HA2yWQSBoNBlnPFR9lPPdipsALiExFo3UZoNBopGznbaFZyDBaLpaAIgVLyh6L2dyWJajRNo6OjAyzLYmxsDAzDoK2tTZE66WKYTCaEw2HE4/GK2mWWQiqVqridoyj2UYqzkE2jPvPVoIaLa4goXDDT7NDpdILjOESj0aLbVVsr26gGFpB3rVkJREObyWSkJtuiYHoh5AoPZ+Pz+aqWNxQEAcFgECMhL8DkPyZKR6O1v11al7RYLDAajdDpdGAYpuLfIiY55WPx4sWK3JdytCrUarXo6emRIhh+v7/mSxxKh415nq94PVyr1YJhmIpKjmqpZlcLVCNbQ0SDWEqWZkdHBwKBQFEjqtFoKr4ZG9nAAo1vZEXa29vB8zx8Pp9Uu5zJZBAMBnO8s0rEJWaCZVnJiywXQggikQgGBwcxPDw8qXI1tx/WZfk9D8uR3aC0ynTbEeUPxQFdNLz77LOPIt8nZ6vIeopZUBSFjo4OWfoCK4G4Nlvu5GO2tb1Tw8U1JB6Pg2GYkkJLYiKUWBcnZziq0Q0s0DxGFpgcTCYmJjA2NoZdu3bh4YcfBsdxkrHYf//9AUD2+sZSlhSyERu5h0Ih8DwPq9WK7u7unHvLduxkhnJs0zBIRgClo2E5slt6XQlomsaSJUtwwAEHQK/XSyHkiYkJBINBWddkRbRabUm6vKUgilnYbDb4fD4EAgG0t7dXXTddCkajEeFwGLFYTNbQuhxeuVgSFolESs5VACZD4bFYTLEweK1RjWwNSSaTZYVfxHq9oaEh9Pb2FhygyzFIzWBggeYyssBkIf5f//pXbNu2TXpNLEXJZDI48cQTZf0+MVO3FCORSqUkz9pkMhUVWKBoCvbj58C6ogfD//age+9+xTzYbCKRCJxOZ45hamlpwdDQEIxGY1UiG/kQQ8bl6jwXo15iFu3t7RgYGIDJZJJtMi5OEqvF6XRiYGCgrOx3g8HQMA3r5UANF9eQSurO9Ho9HA5HQV3QcrLxmsXAApNGttbJJNXAsizeeuutvO9t375d9u4iPp+vaDs7MTFnYGAA4XAYbrcbfX19aGlpKWnQp3UMYNPWxMACk91epnp+YsnU6Oio7HJ7RqNRMVlScXLc0tKCkZGRikKm5SCGjUdHR2XbZyqVksUTpygKLS0tZbXrm20t75pnFGtyxKSnSm5cm80GmqbzNtcuVcO4mQws0HyebCwWK2hIRYPHsmzVgwdhebD+JJLR6Vml4trwwMAAJiYmJD3e9vb2hu5sUuxa0zSNrq4uSQREThiGUbS1mihmYTKZpMmOUsbDYDCAYZgZkyVLpdLynXxYLBakUqmyZGBnk6FVw8U1IpVKgabpimeHra2tGBoagl6vz7n5DQZDXuObTbMZWKD5jKzZbC6oJ63VaqHT6aTaQRGKoqDT6aDX66X/FvLeiUAQWTeA2CvDIKwADUMhvIyCZVUvItEIotEoaJqG0+mE2+1uqnM3UxmKXq+H3W7HxMSErOFdMWRcaVOFUrFYLDCbzQgGg/B4PGhtbVVEramtrU0KG1eSDJdNOp0uGikpF7FLT3d3d0nbi85DIU2BZkI1sjUikUiAoqiKZ4di6GxgYAC9vb1SyE+s+StEMxpYoLmMbDwex8TEBA455BC8/vrr095fvHhxXuNACEEmk0E6nZbUj7Jn7xqNBnq9Hnq9HuzGccRfyhKK4Agi6wcnk0qOn1N0zb7RiUajM5Yh2e12jI6OyprgYzQaa7b2VwsxC7GsbHR0tOz61KnIvVyj1+tBUVTJHrKYYawaWZWSEVPSq3moaJpGd3c3hoaGpEzVYgNrsxpYoDmMLMuy8Hq90Ol06OvrQ19fH7RaLV577bWc7OIVK1bk/bzYwadQdEOUIEzFkkhuGpn+eQDUriispzd315NSlYE6OjowMDAAvV4vy/0sCojUUpBeaTELMSoSiURkKxWTi/b2dgwPD5ekTmYymRAKhapuetEIqEa2RogKRtUOhjqdDq2trRgdHS3a9aWZDSzQ2EZWEARMTEwgk8lM80YWLlyIRYsWAfhUzahSGIYBwzDQpzVIcvnXp0hGAB9lwbhqk6AkN4IglGxgxLK27ElmtYgh42oFPcpFFLNIJpMYHByExWKBy+WS5Tdlt5yrJGyslNqaRqOB0WhENBqdsQEATdOzZk1WTXyqAXIbDLPZDL1eL6kLTVVIaXYDCzSmrCIhBKFQCB6PB1arFb29vdPObzweh8PhgNPplO3ca6xaUNr8jyqlo6GxNuc1BlB2PSTDMNIkUw5MJlNdBemVELMQw8YjI9OjH6WQL9NbLlpaWspSx5oNhlY1sjUglUpJWq9y4Xa7kUwmJwUuNDpMDIfAZfhZYWCBxvNkU6kUPB6PpDmcL3GFEKLIcVNaDUxHdCLfcKOkElMtiMViZbc1M5vN0Gq1CAaDVX8/RVGgaVr2EqFyj0HMBE8mk/B4PFUrUun1ekmoolzkzCyeirg2XUx+VKSQpnWzoYaLa4Co9CT3uk9HRyf+8egOfLw1Ap4VwOho9C+0YsU585rawAKNY2Q5jsPY2BhomkZPT0/RiZKSiRrR/TSwohuJzaMgGQGEoWBb1qOoElMtqHQ9VE6hCofDgXA4XPf1P7nFLNxuNzwej5T5XiqpVErRc2Gz2aSes8WeJzHK0MjlZ6WgerI1IJVKgRAi++xw618/xoevhsD/t3MKlxHwn81hvPn3QVm/px7U28gSQjAxMYGRkRG0tLSgs7NzxkhEufJxpRKNRmE0m+A8cQ903bIYHTccClzSC+uxfaDo+k9EKoXn+YqjO3IKVZjN5ppqDs+EXGIW2dnG5VBJi7pyKaXnbL1D+XKhGtkaIJZqyGlkuQyPt17Mb0zfenEQXKZ+4S85qKeRjUajGBgYgMFgQF9fX8nrU3Jp4WZDCIHf75cScyitBozLAJvLUXaT+Eaj2nIcuYQqGiFknA85xCx0Op2UqVsOSj97RqMRPM8XVatjGKbhrkklqEZWYcSHoppZe/a+UqkUJiYm8OG/PgGXyd+Bh8sICIw39wBcD1nFTCYDj8eDVCqF/v7+stYKK5HMLIXx8XG0trZOG/SsVqts6j71Qo6a12yhimqw2WwNO2mxWCzo7+8Hz/PweDxle3culwvhcLgsxaVaIApUzHZUI6sw1azTZRtVj8eDwcFJ8QGLxYI99+kHo8t/+RgdDZ5KY2BgAF6vV7a2XrWklp6sIAgYHR3FxMQEurq68hq1mYhGo7LXJWYyGbAsmzf7VjTozZx9KcfEE5gUquB5vqqQb6NPWsSEod7eXqlFYalJQRRFoaurq6RsY47jZE3QLIaohBaPxwtuI9YxNzNq4pPCJBIJGBg9kpEICMsXzQQlhCCdTiMajUriFQaDAVarFS0tLdMG/oNW9eKN5wem7eegVb3o7O4AMCmPFggEkMlkYDKZ4HQ6a1Z4Xw21MLJiSU44HEZ7e3tVSUuxWKxo3XIleL3eovu0WCyKGPdslLoOcgtAVCtUIdawl1O3Ww8qFbPQarWwWCwztg5UMrM4H62trfB4PDCZTHnvM5PJhGQyWXYGeiPR+KNtE0MEguT6EaR2RkGxAka0o7AsnezNSdFU2UZ1Koedsgci4Qg+3hYGl5nMLp6zyIbDTtlD2kav16OzsxOEECSTSYyNjYHneVgsFtjt9prNWstFaSObSCQwPj4Oh8NRtbCB2PxBzsE5EonAZDIVNUQ2mw2jo6OKGVlRpF2J61CKIEE5yCFUIYaMHQ6HbMelFJWIWTidTmnbQhORdDpdUylDmqZht9sRCoXyGn+j0YhQKKQaWZX8RNYNAG98WqdGWAHR9YNIJBJgF06uRZVjVKeSyaQx/7hOrDxnPhKRDEw2HUKRIOLx6bWHFEXBZDLBZDKBEIJYLIaRkREQQqSG041QMiOipAfl9XrBMAz6+vpkMYzpdFrW4n1BEBAIBNDfX7w8R6PRQBAExc6Vkqo78Xhcds8/W6iikn1brVYMDw83hZEVEcUsIpEIBgYG0NLSUnCdO1ukoq+vL+89k0qlav77HQ4HBgYGYLfbpz2P5bTybFRUI6sQhOURe2UY+YY+fnsQPSfvPdmzswp8Ph86OjrAMBrYWiZnn263GwMDA7BYLAUHXoqiYLVaYbVaIQiCtMZD0zQcDgfM5vpr4cptOMSSnHQ6LXvrN7lLd8bHx9HW1lbS7zebzYjH47KJ5mdD0zQEQVAk2qFUWNZsNiORSMwYFs2HeDyNHjKeiihmYbPZ4PP5EAgE0N7ennfip9VqYbPZEAgE8nbZUep6F0OsD56YmEB7e/u095qd5rmTmgw+yoKw+bN/wQoQYtUt5rMsC4qipoUTKYqCzWYrWelFNKx9fX3o7OxEOp3G4OAgRkZGpBB2PZBTVlGc5ZtMJvT29speZiNXg2tg0ivmOK7kVmh2u70iVZ9SUKqnp9J1mC0tLYjFYhU1ZW/0BKhiiMaqu7sbfr8fw8PDeZOGnE4n4vF4Q3mIZrMZmUwm7/FOlY1tNlQjqxAaqxZgCniSMujNTkxMFBQ1dzqdCIVCZQ+QGo0GbrcbfX19aG1tlepF65GhLIcnm05PZlhnMhn09/cr4u3J1fhBxOv1oqOjo+TtxVpCJYyh6MnKjdzrsVOpRqiikUt5SqUUMYuuri6Mjo7mvF7vTPX29nZ4vd5pr4tt75oVNVysEJRWA3KQDdQb072MavVmeX5So7iQ90RRFJxOZ8GQUClotVqpB2o9MpSrMbKiLB2AGaUQS91foXKJWCwGvV5fkdc0lWg0Cp1OB47jyipb0Gq1CIVCsiescBynyOAWDofR0dEhyznLh16vzxGqKKfXrrgOXW/FMTkQxSxisRgGBgbgdDql3AuGYeBwOHKETpQQUykHUd99aoazyWQqu5FEI6EaWQUhh9phMBiQem0cFEdA6WhYjuyuWm82+8EohKgP6nQ6q15fqleGcrmDHCEEgUAAsVgM7e3tspQixGIxDA0NFZzli3WFfr+/qu8hhIDneTAMU7Y6j5jdLPd1EL1YudcnOY7DJ598Ius+s7HZbOju7s4RqhAnjKUghowbrR9rpVgsFpjNZgSDQXg8HrS2tsJkMsFut2NwcBCZTAYURWFsbKzuhqytrU3KEBcxGAzw+Xx1PKrqUI2sQgiCAJrRgFtkhbC/Ft3Ozv+2LKte9SmZTM44aFAUJbWVam1treo7s/fZqBnK8XgcExMTcDqdM2bllgrP8xgaGoLJZCooUCHX7D+TyYBhmIoNmhJeCMdxUmNzuRAEATzPK7YmG4lE4Pf7JWF9u92O0dHRstSllC6NqgeimIXD4cD4+Dj8fj86OjrQ0dGBZ599Fu+99560Vr548WKsWLGiLslfGo0GZrM5p+m8UrkBtUI1sgohKj1FIhGYLCYwLnkKvIPBYMkp9haLBX6/X5FsyUbJUGZZFl6vFzqdTrZG3tn7JoSgtbU1byhWPK/VGjc59qPRaKDRaGS9zkoYWZZlodfrFYt+iFrP2WIX5QpVKF0aVU+mills374dO3fulN5nWRYbN24EAKxataouxyhWSFitVun8K1mzrTRq4pNCJBIJSQRbLg+DEJIzwyuFUrpdVEs9MpQFQcDY2JiUKNTe3q7YA1hov3LpUcuRbauE/JwSHoTSmtT5rpUoVFFOIwGLxVJU7q/Z0Wq1aG9vzzGw2WzZsqVuvVwpioLb7c4JERsMBsXW8JVGNbIKkUqlJM9CrvIOscF1OcbEaDSCZdma6X8qnaEsSiF6PB5YrVb09vbWrXeuHBEC0VBXO0FQUjhCLuqZUJQtVFEKogrRbCYWixUcF1iWrWspk9VqRSKRkLLDmznDWDWyCkEIQSKRkCWcKBIIBCpqplwLbzYfYoZyf3+/lO08MDCAiYmJGY0+x3EIBoM5s+lUKgWPxwNBENDf319yLakSiMaiWjnGmbxhhmGwYMEC6d+jjz5acNupJTcbNmzA5z73uaLHsHbt2px7Y86cOWUJ7V944YX461//CgA46aSTig6EH3/8MY444oiS9y03ZrMZWq0WwWBwxm2zQ8azFaPRWLBKQEzA83g88Hq9iMViNa9Vze7S08y9ZdU1WQUQPZx4PA6apmVZf0omkzAYDBUN6nq9HoIg1KQZc7FjKCVDWRAErF+/Hps3bwbHcdBqtTjssMOw7777gmEYWUpy5KCYF5tieUxE02i16mEokugmXo9i19ThcGDHjh0lHZNGoym75d7atWuxaNGivIl05YoAPPfcc0Xf53m+7mtqLS0tGBoagtFonDH7XBzY651xKzcsy2JiYgI8z+Owww7Dq6++Om2bJUuWYI899pCWM2KxGEKhkHTfm0wmmM1m6HQ6xa6pwWCQ9N31en3TTnhUT1YBxKQnObwdEZ/PN2PZTjEapXejmKHc3d0thXpHRkbg8XgQDoexfv16bNy4UfJ0WZbFpk2b8N5776Gzs7MhDCyQv38sLxDc/X/v4+Db1mHZXetx8G3rcPf/vQ9emD44VFse09LSguuvvx7z5s3DqlWrpAndq6++igMOOABHHnkknnzySWn7jz76CMuXL8f8+fNx6qmnIhAI4KmnnsK2bdvwuc99Docffri07d13342DDz4Yhx9+uBReHRsbw+mnn45FixZh6dKl+Ne//jXtmLK94FNOOQULFy7EgQceiD//+c85v7melCNU4XA4ZlXIOJ1OY2hoCOPj42hpaUFvby9WrlyJefPmSZNvrVaLZcuWYcWKFQAmz5dOp4PL5UJPTw/6+vrQ1dUFnU6HUCiEwcFBRb3dqeNWMxpa1cgqQDwelzWUmclkqvaItVotaJpuqN6yYoZyb28venp6wLIsNm/enHfbrVu31i0RIx/5Enh+uu4D3Lf+QyTZycE7yfK4b/2H+Om6D6Z9ttSoQigUygkXi5mffr8fJ510Enbt2oXu7m7JkH3ta1/Dgw8+iE2bNuWEga+++mpcfvnl2LlzJ4488kisWbMGZ5xxBhYtWoQ//elPeO2116Rt+/r68Oabb+KEE07A//t//w8AcM011+CWW27Btm3b8NOf/hTXXHNN0eN+6KGH8MYbb2DTpk349re/3RAGViRbqKLYoK2kmlYtSSQS8Hg8CAQC6OjoQHd3t7SENT4+jpNOOgnXXXcdzj77bNxwww1YtWpV0ckfTdOwWCxob29HX18fent74XK5kMlkpAnz0NAQAoEA0ul0VeePYRgYDAbEYjFotdqGGgNKRQ0XK0A6nYbJZJJtsd7n88lS69rW1obR0VH09vZWvS+5EScRxRIxYrFY2aLvSpAveSfF8vjdKx/n3f7/bfoYV67cSwodl5PsVChcbLFYsHLlSgDAwoUL8cknnyAUCoHjOBx88MEAgC9+8Yv4/e9/D2BykvLss88CAM4//3ycfPLJBb/ztNNOk/b7l7/8BQDwz3/+E++9996Mxyvy05/+FM888wwAwOPxYGRkpKFE90sVqhCf4Xqu/1dKLBaD3++HwWBAd3f3tEl6MpkERVFScqTD4ai4F6/o7Yo5I4IgIJFIIBQKSRN7nU4Hi8UCk8lU1r3Q2toqCeskEom6qlJVgmpkFSKRSMBkMlU98+J5XrYyIIZhoNVqpXB2oyH2ucx3zsSm041AvmSliWha8mCnkshMrtH2ukxSslO11zM7Y11cixWNfz6vMdugz5ThK+5b3K/IG2+8UVI0Zf369di0aRO2bNkCo9GIfffdF8lksmFC/SKlCFU4HA5MTEw0jZEVy/yCwSAsFgt6e3vzGjRCCMbGxtDX1wdA/s5DorebfV4zmUxFa7sURcHhcCCTySCZTIIQUrQnbqPROFPLWYIoGJ9Op6UZXjVUuxY7lXplGpeCqDaTj8WLFzfMQ5VvQGq16mEskORk0mnQap00XKUkO1WK0+mEVqvFjh07QAjBY489Jr23aNEiaY32f//3f7Fs2TIApXedOfroo/HrX/8awOTv37VrV8FtI5EI3G43jEYjXn/9dXzwwQeydlWSk46ODvh8voKTYa1WW7Pyt2oQJUUHBgak7PuWlpaChlNsMJLd3k/pSVA1a7tWqxUbN27E73//e/z85z/HXXfdhRdffLGhliEKoRpZmRE9WEJI1VJ3hBCkUilZvU5x9lhOmUYtWbFiBQ4//HDJoDIMk5OIUW8K1XoatBpcsnRu3s9cfORcGLSaipKdpq7J3nPPPUW3f+CBB3DxxRfjiCOOkLwUALjnnntw7733Yv78+Xj55Zfxve99D8BkCc6FF16Yk/gkkv0b7733Xqxbtw7z58/HgQceiL/97W8Fj+H4449HOBzGggUL8Itf/ALz5s2TtVORnJQiVGEwGBq2RlMQBExMTMDj8YBhGKlcrti5TqfTyGQyOV5mPXroTl3b7evrK7i2++KLL2Lnzp05CZEbN27E+vXra3rMlUCRJlzVF5tkh8PhhtMXHR8fh9lsRjgchkajgcPhqFiMwu/3Q6fTyd4WTBAEeDwe2WUI5WJiYgJ6vR48zyORSMDpdNYlXJdKpfDxxx9j7ty5UrkHz/MQBCGvV80LBD9d9wH+36aPkcjwMOk0uPjIufjGsXuDpj7VF1b6nIsTPDlEUOTQRBZD2Up3bgLyX7NSiMfjCIfD6OrqmvZeJpOB3+9HZ2ennIdaFTzPY2JiAplMBm63u+QyI0IIPB7PtFK4aDQKnudLlmytFYIgIBwO4xe/+EXeiIJWq8UNN9xQ8yhXOTZIXZOVmXQ6DZ1OB7PZLLUuqwRRgD/bG5ELmqZhtVqlG6XRSKfTcLvdoGla6hQil+h/tRQLq2loCtcfvw+uXLnXtDpZsVtPLSY1otfYKFqvtQhFVovZbEYikUAwGJyWXKfT6RomqzW7xrW1tbXsTlOBQCBv56x6eLKlIB5TMyREFqLxzmqTIyo9mc3mqgY5sbG1UoOky+VCMBhsyPKE7PIYmqalCUsjUMpgZNBq0OsySQY2u41drZiatFRPGnUAn0pLSwtisVhejVy5egZXSr4a13INLMuyiMfjeb3VRr5GxZKcGikhshCNeVabFDHrNLsDSKXkm1HLiZixV4rEXL1xu93w+/11nxBUOmmqh9IWTdMNYWQbxZsuhWJCFfUSpihW41ouo6OjBUPejWxktVot5s+fn/e9RkqILERjntUmRa7SGLGDj9KDk91uRyQSqbvxyibfsVAUJa1/1JNKuu6Ig3WtB7DskHE9aeTBOx+FhCr0ej0ymUzNjiMWi2FgYADRaBTd3d3o7OysauIeDodhMpkKGqRGvk6pVAoHHXQQli1bJp2DqcpUjYy6Jisj8Xhc0kSt5qb1+/15EzDkRmwp5ff7ZS0TqoZCCTsOhwMDAwOw2+1184wKJTwVghACjuPqVjwvhoxrGaaeSr2/vxIKCVXodDpJR1cJSq1xLRee5xEMBovmNTSqkSWEwOv1ore3F319fZg7dy5SqRQ+85nPNLwHK9J4Z7WJEVPjzWZzxVmZmUxGasBdC6xWK+LxeMPUm6VSqbyDWPaEoB5U0qZNrjZ2lTK1K089ULp/rFLY7XbwPJ9T6qZUyLjcGtdy8Xq96OzsLHoflttYolaIa9DieKjVamG1WpvGwAKqkZUdsa61UiM7MTEhi4RiObS2tmJiYqKm31mIdDpdMKGjnhOCco2FXMlO5bS6Az5tNxcKhfDAAw9Ix1JK27upiNGNbdu24Zvf/GbRbdesWYP77rsv57VmWo/Nx1ShCoPBIKv2dyU1ruUSj8fBMMyM3ncjerKpVAocxzV8YtNMNFccp4ERvRbxZhU92nLgOK7skKQcmEwm+Hy+itYc5UYsgSpES0sLJiYm0N7eXsOjKnM9lk2CDQ5Ba+8CUF1osZxWd8Cn7ebGxsbwm9/8BhdffHHVCVCLFi3CokWLyv5cIw7c5SAKVQwNDUk15Vqttura4ak1rkpNqkUjXkr5W6NNiMQwsRIljLWmeZ+ABmNqstJMxiIfcjUCqIRGklss9rCbzWak0+maZ86WJAso8MCL3we5ay50v1wE+u69gBe/P/m6zLS0tODKK6/EggULcPrppyMejwP4tN3czTffjHfffReLFi3CD3/4QwCTyS+nn3469t57b1x77bXSvp577jksXrwYCxYswFe+8pVpkYJsL3jLli044ogjcPDBB2PlypVSK7x8NMKkrVoYhkFra6v0O6sJGbMsi5GREYyMjMDhcKCvr0/RXrXj4+Noa2sr2Xg2kpEdGxtDa2tr3klauX2O641qZGUikUhIvVKB8gvwBUFAJpMpu/ZNLgwGA3ieb5ii+2LUekJQcl/g9T8E2fhjUOx/JfjYBLDxx8D6H1b83cVa3a1cuRI7duxAX18ffvGLX+R87vbbb8f++++PHTt24Fvf+hYIIXjzzTfxwAMP4O2338azzz4Lj8cDn8+Hn/zkJ9iwYQN27NgBnU6HJ554ouDx7L///ti4cSPefPNNXHrppbjrrrsKbtto3lGlmM1maLVaBINBGI3Gsutl5ahxLZdUKgVBEJqmsUE2yWQSgiAUnIA0W8s7NVwsE2LfxEprWwOBgNQmql6Ixqu7u7su31+qiHz2hKAWofWSQvhsEtjyS1CYWjJDgC2/BI66HtCWX95VKFys1+txxhlnAADOOeecosZOXMZYsmSJFCk58MADMTAwgFAohJ07d0qNGZLJZNHrHwwGcd555+Gjjz4Cx3EFQ5GzxcCKtLS0YGhoCEajEQzDlHTvJRIJ+Hw+aLVadHR01CzLuplDrVO7A019D/jUyCqV5S03ZXuyL7/8Mk455RR0dXWBoig8/fTT07Z57733cOqpp8Jut8NqtWLx4sXweDzS++l0GldddRVaWlpgNptx6qmnYmhoqKofUm8oiqp4rUaUUFQydFQK4rHXq7F7OeURtfRmSwp7xsYmPdd8sInJ9xWkmEETjWyh9nirV6/Gjh07sGPHDrz//vu46aabCu7ru9/9Lk4++WS8/fbbWLt2bcF7ZTaEirPJFqqwWq3w+XwIBoN5PSq5a1zLxefzweVyNeV6uNfrRVtbW8H2fDRNN5TMZSmUfRXi8TgOOuigaZmEIv/5z3+wdOlS7LvvvtiwYQPeeust3HLLLTnhkWuuuQZPPfUUHn/8cbzyyiuIxWJYvXp1QyjUVIKoSwugIgEAUUO4EWb+7e3tdVubLcfIikL7tZoQzHhtLO0ghTxVrQmwyJuolU6npYbqTzzxBI444oic97Nb2BW7JxcvXoz169djcHAQwGQYeuqEN/uzkUgEPT09AIBHHnmk4PE1e9JTPmiaRkdHB9atW4ff/e53OS3XeJ5HOBzGwMAAUqkUent70d7eXvOJRiaTQSqVarjGKaWQSExOUguFuMWJm5h81iyU/RSceOKJ+MEPfoAzzzwz7/s333wzTjrpJNx11104+OCDsccee+Dkk0+WirrD4TB+97vf4cc//jGOOeYYHHzwwXjkkUewa9cu/OMf/6ju19SJRCIBnU4nhY/KDWOGQqGG6X6R3di91qRSqbLWqmrhzZY6YeJpHYTDvgZgqjGmgMWXVxQqBgq3unO73Xj11VexYMECfPzxx7j88stzPud2u3HIIYdg3rx5uP3220HTdN7f0tbWhl/96lc4/fTTMX/+fBx33HFFz+n111+Pa665BkuXLi263jfbwsUir776at6Wa3/5y18gCAL6+vpkrXEtB0JIUenERkYQBIyNjRWtGhAnbs22JltVqzuKovDUU0/h9NNPBzB5Eux2O2644Qa88sorePPNNzF37lzcdNNN0jb//Oc/sWrVKgQCgZz1y4MOOginn346br311mnfk06nczyWSCSC3t7ehml15/V6wTAMdDodbDYbYrEYMplMSWus8XgciUSiblnF+eB5HsPDwzVf0xkcHERPT09Zg7PX64Xdbpe1565IKpXChx9+iLlz5xYN5Uu9gxkNqA13TK7BsolJD3bx5cCKbwO0vB5NS0sLfD5fydsTQsCybNnLGZlMpuwm84Ig1E3pqtJWd6XAsizuuuuuvAN8vVquZSPqkFeSFzI4OIje3l65D6lkxIzrYhO3RCKBZDIJt9td9+Mtp9WdrNOt8fFxxGIx/OhHP8IJJ5yAF154AWeccQbOPPNMvPTSSwAmB0WdTjftRmhvb4fX68273zvuuAN2u136V8+Tmw+WZZFKpaSBuJy1Wb/fD7fbreThlY1Go4HBYJDKQmpJud5Pa2urot5sKSIUYkMISsMAq24BbvgI+Ppbk/9ddYvsBrYSxJBxuXPqSpY/mqG1XSXEYrGCHpTYcq1ecByHSCTSMBGxcojH46BpesZM6GZdgpD1iMXapdNOOw3f+MY3sGDBAtx4441YvXo17r///qKfLRZeuummmxAOh6V/4vpRI5E9sJRqZNPpNLRabUPeOOV6SvVCo9HAZDIpMsCJxqWY4RcEYbpR0RoB55yKQ8SlUMm1KRQylptmHQxnwmw2F0xgqnfLNTFM3GwhelEwoxRxmWZNppP1SWhpaQHDMNh///1zXt9vv/2k7OKOjg5kMplpLdbGx8cLnmi9Xg+bzZbzr1HgOA40Tefc3KWuydZDQrFUaJqGxWJBJBKpyfdVo50qTgjkNiCJRKIkL7ZZdFTFNoxKMxvXYxOJBEZGRgoqX9Wz5VokEoHBYKhbI4pqGB0dRXt7e0n3y9TJW707TJWKrEZWp9Ph0EMPxfvvv5/z+gcffCDV0y1cuBBarRbr1q2T3h8dHcXbb789LUOyGUgkEtBqtdPWf2a6acTBrpE7lLhcLgQCgZrczNV0N6EoCjabTfYJQTwen1FUnaKopvHaKvFkyw0XzzYvVkwmCoVC6Ovrw3HHHYdly5ZJBpVhmLq2XBMEAYFAoKouWqXWp8tNLBYDwzAl51Nke7JiCVozUPYIH4vF8OGHH0p/f/zxx9ixYwdcLhf6+vrwzW9+E2effTaOOuoorFixAn//+9/x7LPPYsOGDQAmu1tccskluO666+B2u+FyuXD99ddj3rx5OOaYY2T7YbVCTDsv17tuZC9WROzjGgqFFG0gDxRvDFAKTqcTAwMDsNlssg0YHMcV3Fe929hVitiZRylDOJuMbCKRwPj4OFpbW3MS31atWoWjjjoKsVhMagJRr9/s9XrR0dFR1T1fj2smCAJ8Pl9JusrZnxGPU8wwbmQnRaTsI9y2bVvOrE3UQL3ggguwdu1anHHGGbj//vtxxx134Oqrr8Y+++yDJ598EkuXLpU+89Of/hQMw+Css85CMpnEqlWrsHbt2qaMt4uJEOVktwqC0DSKJWIfV4fDoehsN5VKVbWmRVEUnE4ngsGgLMpZYjJToTpcKdmpycKiogegpJFtlvB5IUTFJEII+vr68p4rrVYLp9MJQggGBwfrknAkSrlWm0VdDyM7Ojpa9uQg25PV6XTIZDKKVBXITdlndvny5VKWYva/tWvXSttcfPHF+Pe//41kMokdO3bgtNNOy9mHwWDAvffeC7/fj0QigWeffbbhMoZLZWqf0VIW5xsxo7gQFEVJYWMlEY1WNYghYznC25FIpGDZjiAIIIQUvM4pLoWh6BBSXHkat/mY2uru9ddfl9rO5WsvNxPZPWaXL1+Ot99+e9o2F154If76178CKC9c/OCDD2JsbKzpJh7ZJBIJKSLS1dU1o/GhKAoMw9RcHEGUH5SjG1WtjWw0Gs27xDYT+TzZZqDxfe0GJlvpSWSmzGJCSMPVxc6EzWbDwMAAnE6nog9jtYMzRVFSElS15zcejxdc5yqU2MYLPH6x4xd4+N2HkeJTMGgMOH//83HFgiugqbCMJ5928WGHHVby5/NN+pTqYvLQQw9hwYIFTTlhLsV7LYTb7UYgEEBHR4eCR5jL2NhYQfnBcqllyRXP8/D7/WWFiUWy1451Oh3C4bDch6cIs2PxpE7kq++aqcVdOByG3W6vxeHJSrOU9FgsFiQSiaqMyNToRDZiqDXf4PaLHb/AA7seQIqf9GBTfAoP7HoAv9jxi2nbVsrU5utbt27F8uXLsffee0s64mvXrsW5556Lk08+Geeeey4++ugjLF++HPPnz8epp56KSCQiJY088MADWLx4MRYuXIgPPvhg2vetWbMGS5YswQEHHIBvf/vb0uvf/OY3sc8++2D+/Pn4wQ9+gKeeegrbtm3Dueeei8MPP1y231sLyvVep6LX66UGIbUglUqB53nZtM5r6clWW2okfq5WmfJyoBrZKhAH8+ybfSZPtlmNrNlslh5uuckXEaiGauUWk8lk3sJ4MdkpX1g7xaXw8LsP593fI+89UnHoOFtW8atf/eq0999991288MILePHFF3H11VdLcpivv/46HnvsMTzxxBO4+uqrcfnll2Pnzp048sgjcdttt0mTEEIItmzZgptvvhnXXHPNtP1//etfx6uvvopdu3Zh+/bteOutt+D3+/GHP/wB7733Hnbu3ImrrroKZ5xxBg455BD88Y9/xGuvvVbRb601UzOHqzFa2VrRSiJ63HJ6zbUyspFIBHq9XpZclGZaklCNbBWwLAue53OMajEjK3baaaYbJJvW1lZMTEzIvt9qynfyYTQawbJsxTPdSCSSN1u8WLKTL+mTPNipJLkkfMnKogBiuHjHjh15BV3OOOMM6HQ69Pb2Yt9995XK544//njpN2zduhWf//znAQDnn38+Nm7cKP2Gc845R9rP9u3bp+3/xRdfxJIlS7BgwQJs374d7777rqS8dvHFF+Ppp5+G2WwuSbijkajWe51KNc3cy8Hv98PpdMo6Ka2FkeV5vupSo2ZFNbIyU0yGrxF6xlZDtcarEOU2BiiFarzZfBOlmZKdWowtMGjy/wYjY0SLUZnBZapRE//O9sSztxHD4BqNZkbRiFQqhWuvvRZPPvkkdu7ciTPPPBPpdBoMw2Dbtm343Oc+h8cffxznnHNO0whQyOm9ZkPTtOIhTJZlkUgkZI+E1cLIjoyMSO1RdzdUI1shLMtKvQ1LIZVKQafTNX0NYXt7O8bG5O2NKrcnC0yukwmCUHbWZ6HylplUvAzMZJJTPs7b7zwYGHknESJPP/00MpkMhoaG8P7772Pvvfeets2iRYvw5JNPAgD+93//F8uWLZOEKf7whz8AAP7yl79g4cKFOZ9LpVKgKEpK7HnuuecATEZkwuEwVq9ejR//+MfYsWMHeJ6vWci0UuT2XqeidBb+6Ogourq6ZN+v0kY2HA7DaDTKXlNeia52PVCziytErFErdSbs8/lqmn2oFOKDUmmD+nyIBf1yIzadKCfbNRqNTgsVi0sCM83Cr1hwBYDJNdgkl4SRMeK8/c6TXleCgw8+GMcddxxGRkbw85//PG/d4D333IOLLroIt912G/r7+/HQQw+BoiipxnLx4sVgWRaPPfZYzuccDge++MUv4uCDD8aee+6JxYsXA5g8R6eddppUQ3znnXdCEARceOGFuPDCC+FwOBpqXbaazOFyMBqNGB8fV8SrD4VCRbWTq0FJI8txnBQ1kBuGYSrqLFVrqmp1Vy/KaTOkFCMjI+A4Dj09PdINKggCRkZGpKbWIizLYnx8HN3d3fU4VNlhWRZjY2PTfmelKNm2anR0FE6ns+Rw9NDQkOTlpFIpfPTRR+jp6YHVai154ExxKfiSvskQskIerBzwPA9CSNGBu5QWeZW20VOCqa3uCqk2KYXP54PRaJT1u3iex+DgIPr7+xUJt3q9XrjdbkVERAYHB9He3i7LvTF1nAgEAtDpdHVpzFC3Vne7E/l0awuFFCcmJmbVgr9Wq4VGo0EqVb3YgtJzvHLWZgkh02b1Yg1hOYObgTGgx9rT0AYWmFxHlCNbvF7at8VQau11JpxOp+whY6U77CilABYKhWAymRSbfDWLIIVqZCskn3HIZDLT1hZ5ngfHcU0hoVgO1ZbJiMih9FQMjUYDvV5fUm/cqWvDmUymaLJTsyOGjItNdErtjtJI50iMKCm19loM8TzIVeoWjUah0+kUHT+UCBdzHIdwOCxbome+ELxOp1ON7GxFTKaZWkuZb53S7/fPKi9WRDReYoOESqm2MUAptLa2liSkIYaARCYmJhrKeCiBHN1MGqUpACFEmhh1dnbWzHudisvlmtbKsxIEQYDf76+JOpzcXrLc2cT5vG2tVltzOctKqP+T0YSIhmXqQzzVyBJCCgobzAbkqJtNpVKKe/k0TcNsNs/YCi/7WMLhMEwmU8OFQeWm2pBxo5Tu8DyPTCYDjUYDjUZTV6NvMplKipzMhNfrLbnXaiMRDAZhsVhkXePNFy2ppHVjPVCNbAUkEgkQQqYZh6nKRcFgsC7dOWqFaLyqKdtQonwnH2IZSqGHUvTGRF3f2X7tREoJGRej3l6s6L2KGeCNEHmgKApGo1FS36qEZDIp7aeZYFkW0WhU9taYSnaOUprmPOo6IyY95Zthiq8RQgoqB80m3G43/H5/VYN0LQbG7N64+YjFYlKWotjdpNk8iEqpxput5+CX7b2WUmJVS1wuF/x+f0WflbPDTi0hhGBkZESRJK1GW/cvB9XIlomYgTpTsk4sFiur7KNZEY1XM3TEcDgcCIfDeScE0WgUVqsV6XQagiBU5UEIqRQyQ0MQZMi+rsV6vkajKdhQYcOGDdi6dav0d3YbPKC4wplSNKL3OhWGYSAIQkWNKsRqhGbz3AKBAGw2myKlQKonuxvBsiwIIdPWY6cO3M0uoVgOom5rud5srddTRPWifB6G+BBXI75OeB7jP/sZPli8BP855lh8sHgJxn/2MxAFmirIiRguznc9phrZbMpdj5Uj47aRvdepOJ3OsvWM0+k0MplMXWo/qyGTySAej8seJhYp5MnKkbinNKqRLRNxPXaqkc0uRUkmkzAYDA09AMgJRVFwOp1lZ1TWQ8DAarUiHo/neBhiwlo4HIbVaq3YM5q491747/81yH89WJJKwX//rzFx771VH/eGDRtwzDHH4LOf/Sz22msv/PCHP8T999+PQw45BIcddpiUPb18+XJcf/31WLp0KebPn4/t27fj5JNPxl577VW0wXsmk8H555+PefPm4bDDDsOOHTswODiI+++/H3fffTcWLFiAXbt2AQDWrVuHxYsXY++998Yrr7wCYPL+//rXv47DDjsMCxYsKNh2r1JEwYtG9l6nYrFYyspXEGt7Ozs7FTwq+RGPWwnJR5FCnmwzZBirRrZMRDnFqeHi7Mxin883K8t2imGz2RCJRMryTmuRWZyPlpaWnKzoSCQCi8WCUChU8UxcSKUQWPtQ3vcCD/1eltCx2Innrbfews9+9jOkUils374dK1euxMMPf9pmz2w245VXXsHZZ5+Ns88+G48++ihef/113H777QX3ff/998NsNmPXrl245557cMEFF6C3txdf/epXcf3112PHjh2YN28egEmRgS1btuC+++6T9vnb3/4Wc+bMweuvv46XX34ZN910kyS7mN12rxJE71XUCm+WyStFUVKv2VIIBAJwOBxNMYHIxu/3w263K1rvXsiTbQZBCtXIlkmh3qeikRUHg2Z7UKqFoqiyG7vXKrN4KmazGel0GjzPg82kMTHkQcDnqyrZifP5JA92KiSZBCdDw/slS5ZI8oA9PT048cQTAQDz58/HJ598Im136qmnAgDmzZuHRYsWweFwwOVywWq1Fow2bNq0Ceeeey4IIVi8eDGSyWTBdXZx/wcffLD0vevWrcOvf/1rLFiwAEcddRTi8TiGh4cB5LbdK4dm9F6nUmh5YiosyyIejzddRnsmk0EymVT8uAt5ss0gSKE2CCgDQgh4ns9b5J7JZGAymeDz+WpSPN6IWCwW+P3+kss60uk03G53DY5sOi0tbqx78Nd4/6UXwWXS0Gh1WLT6DBxx1hdA0+UP5kxLCyiDIa+hpYxGMDJENrInJDRNS39PzQ7Ofn3qZ4qtX9E0LXkMxdZb9Xq9pHks7o8Qgt/85jc46qijcrZ9+eWXK6oTF5XSGIZpSuMqotVqwXHcjOvXzRomzqfVrgTFPNlGT7pUPdkyKJT0JL4nLsI3glB6vShXK7heGYPbn/kz3ln3HLjMZCiPZzN47ak/4NUn/rei/dEGA1wXXpD3PdcFXwKtsKpVtSxduhR/+MMfwPM8Xn/9dZjNZthsNlitVsRisWnbT/UsjjnmGNx///2S0d2xY0dFxzEbvNepzJR9LwqfKJGVOxPVJB/6fD44nU5Fw8QihSbuDMMo2sNXDlQjWwaiisvU8g4hlQI/MgLfyMhutxY7FaUau8sJm0njjeeezvve9uf+AjZT2hraVFqvugrur14G6r/3B2U0wv3Vy9B61VWVHmrNuOKKKxAOh7Fw4UJceeWVePDBBwEAp5xyCv7whz/kJD4B0z2Lyy67DF1dXViwYAEOPPBA3HbbbWUfQ7Ouvc6EmK+QD57nEQwG6xbRqVRMJJ1OI51Oy95AvhCFIgHNcI+ore7KYHh4GBzHob+/H8BkycbEvfcisPahyTChXg/3RRei9aqrQM2CGXilpNNp+P3+otmGhBAMDQ0p1uKuGKGxUfzu6i8XfP/in/8Gzo6uaW3TSkVIpcD5fGBaWhreg50Ky7LTcgoymQy0Wm3OgCbnejohRAqpTv2ecqn0minNyH8n4FOjXMPDw2hpaalbAxGWZeHz+coKVRNCMDAwgN7e3ppFGoq1w1SyVWYhyrFB6ppsGfA8n3NTiSUbEum09HfbNdfU+OgaB71eD0EQCrb+A/J3LFKaVCoFv9+PTCoJRqcDlyf1n9HpEU6mEPV4JDGBcgUFaIMBuhqsUymBRqMpmNwnIqde8WxZe50JUdYzuwY7Ho+DYZi6duiqxJOdmJiA2+2e1ddLTtRwcYmIa0ViEkctSjaamfb2doyNjRV8v1blO2Ink4GBAYTDYbS1tWHunnth4cln5N1+4cmnY87cPdDX1yclsHEcJ4kEsCwLQRCaQpg8H7t27cKCBQty/n3uc5+T3i9FdH3qZLMSZuPaazHEUh4hw4ELpMCnWUxMTKCtra2ux1WukU2lUshkMrBarQoeVfk08vOoerIlIrbQEpOeZirZ8P7rXzDPnQu9Xg+dTte0kmCVotVqQdN0wbCi0us5iURCagjgdDrhcrlyvK8jzvoCgP+uwaZT0OoNOOSk06TXAUjdXMR+nqIqEs/z0v9TFCVt1wzrQ/PmzZsxKUnMMhbvWVERSvx9giBUlaSzu3iv2RCBQL8jgZHfbgFYAjAULIvbgT4AdbxtytEEJoTA6/XWZYmnGGKtbKMmnKpGtkTE9nbihZypZMM5dy5YikIkEpEMtIhOp4PBYJj1BritrQ2jo6N5H8p0Oi37Q8HzPAKBABKJBEwmEzo7OwsOIDStwdJzzsfhZ56FD999F3vutx90+uJreGJTiOzrJYaTxcxz8X3xXzMihowLXZ9Kw8XZa6+zKbGpFCLrBsBtzqqV5ghSr3gR0WphP35O3Y6rHE92fHwcLS0tDTcxUo3sLCGRSOQkZYglGzlrsv/FdcGXYLTbYQSmLYqL4ubpdHrWG2CGYaDVapFMJvMK7ssxyBJCEI/HEQwGQVEUXC5XWXXKWp0eLd094IXKwk1TjanYQILnealIPtvoNoNhKRYyFgShot+wO3qvIoTlEXtlOO97sU3DsK3sBaWtzzkp1cimUilwHNeQmsqNLkihGtkSYVl2Wtiz9aqrkElnEHvssUmP1mCA80vnFy3ZEKXW9Hr9bmGA29raMDg4KGVkywXLsggEAkilUrBYLOju7q74fFitVkQiEVl6d4rhY9GQZIeYxTCz6O1qNJqGNbpTQ8bifVhuy7Hd0XsVBAGZTEZ6ltMTcYDNnzxHMgL4KAvGVT8jO1PoXwwT9/X11eioykOcyDcqqpEtAXGgmNpVh9JoQL5wLva66koIgQA4iwXhZLLi8p3ZaIBpmobJZMrp11qpN0QIQTQaRSgUgkajgdvtlqXnpl6vR0qhRLV8IWbR6OYLMRfqU1xrRGGVqd53Oeuxs9V7FZO2JCP6X4lOEYqipHV8i8UCl80Br3YYJJ+hZSjw+voNxKV4smNjY2htba3bmDJTUlOj6xc3xkjc4Ig9RqcqPYmlHYzJBF1PD0wOh5QxKTeiAbbZbGhra0NPTw96e3vR29uLnp4eOBwOUP9dAx4ZGcHg4KD0b2xsDOFwGKlUqqL+ltXS0tKS09i93BrLTCaD0dFRDA4Ogud59PT0oLu7W7Y6SNGIyJWhyGV4RHxJcJn894HYYEIciMVlCFGMIZ1OS/cRRVG4/PLLpc+Ojo5Co9FgzZo1FR3bnDlz8io4iWzbtg3f/OY3JU82m5aWlqLrsU899RQOOuggHHDAATjggAPwi1/8AjqdDpdccgn22GMPHHTQQdhnn31w2WWX5ej5nnHGGXA6nTlZzsCkQEZbWxsWLVpU0W+tBo7jkEgkEAwG4fV6c56noaEhBAIBsCwLg8GQ93lsa2uD3W6H0WiERq+F+cgu5Lu7TEd0wh8OwOPxSGI3tWQmI5tMJvOOfbVkJmW4Yv2QGwHVky2BRCKREwIUERt9ZyMalFqm5je6B0xRlBSStdvtJRlZQghCoRCi0Si0Wi3cbreiiQ0GgwHpdLoqwy0IBK8/+xHeenEQXEYAo6Nx0KpeHHbKHqDpwt5poRCzIAhwuVzYvHkzkskktFot/vjHP+KAAw6o+BhnYtGiRZJRoyhq2uBV6P5488038a1vfQvPPfccent7wbIs/vSnP0kG+Z577sHq1avBcRy+853v4NRTT8WmTZsAAFdffTUuvvhiPPRQbkncF77wBVx88cW47LLL5P6ZUkhX9ESnPhcajUZ6Jsxmc9UiGdH9NDALXUhu8YJkBFA6GpYju2E7th8UTUmlZj6fD1arFU6nsyYRjWJGVgwTy73UUy5ylIzVE9XIloCY9DSVSCSC7u7unNdMJhMmJiYqliuTm0YxwE6nEwMDA7AazUiOReHqyZ+cJApG8DwPu92O3t7emgw2FosFsVisKiP7+rMf4Y3nB6S/uYwg/b34tD1L3k92iJmiKBx11FF49dVXsXz5cjz11FM49dRTwfM8/H4/jjjiCLz77rvQaDT48MMPcf7552Pz5s14/fXXcd111yEej2Pu3Ll46KGHpiWt3HnnnXj44YdBURRuvPFGfPGLX8SGDRtw33334U9/+hP8fj/OPfdcxGIxHHfccQAKG9kf//jH+Na3voW+vj5otVro9XpceOGF07ZjGAa33347PvOZz2DHjh1YsGABVqxYgQ0bNkzb9sgjj8zpLiTy6quv4hvf+AZSqRRaWlrw4osvTtuGECIZT/G/2REmsTRLp9PBZrMpusTi9/thsVnhPMkJcuwc8FEWGqs2J9mJpmm0trZKSyIejwd6vR4tLS2KagMX6m4DAF6vF21tbXUfxxplLK0U1ciWAMuy0wYocYaf7+K7XC4EAoGG1zGu1gDr9XoYDIbSBigCGHemJusEOQKf1gPL0smZPAFBMBiUjFxbW1vNxdKNRmNJLckKwWV4vPXiYN733npxEItOnANGV9ls/KyzzsLDDz+M/fbbD3q9Hh0dHZiYmIDD4cDChQvxwgsvYOXKlXjwwQdx3nnnIZ1O4/rrr8czzzwDp9OJ//mf/8F9992HG2+8Udrntm3b8MQTT2Dbtm1IJBI49NBDsWLFipzvve2223DSSSfhuuuuwy9/+UsA+e93nufxzjvv4Nprry0p2qDRaHDQQQfh/fffx4IFC8o6F+l0Gueddx7+8pe/4IADDoDf7wfHcdLzmMlkwHEcvF4vTCYT9Ho9TCYTnE5nXbyhVCqFZDIpdaqhtJqiSU4URcFms8FmsyGVSsHr9YIQgpaWFlkS86ZSyICJJYv1DBOLqJ7sLEdMcph6s4XD4YJiCmLLN7fb3RBJLJUgtwGOrBsA++qnjdIJKyC6fhDRaATkcGdewYhaIgouVEoikgGXyb8uxGUEJCIZ2FoqGySPOOIIXHXVVXj88cfxuc99DqlUSgoxX3LJJXjooYdwwgkn4Mknn8TGjRvxzjvvYOfOnVi+fDmAyUmi+P8ir7zyCj772c/CYDDAYDBg1apV2Lp1a849vXnzZnzrW98CMBm6vfXWW3OuT3bmMICyPK5SzrUYMhefQUII3n77bcyZMwf77LMPOI6TwqoMw0jhbYZh0NnZWXftYkEQMDo6WnG41WAwoKenBxzHwe/3Y3x8HHa7HXa7XbbnJN8auyAIGBsbq3uYWKQUT1Zs49iIxlg1sjOQTqdBCJk2i4xGowWVTyiKgsPhQCgUgtPprMVh1pSyDTDLAxuH8grbUG9F0X3aAXWrE8xGp9Mhk8lUtPZrsunA6Oi8hpbR0TDZKl9PFkPGP/rRj/Dee+/hsccek95buXIlrr76avztb3/DvHnz0NraitHRURxyyCF48cUXpZpdYNLLA5BXFjLfYCv2jOV4Dhk+V+d5aubw/vvvjzfffBPz58+f8fcIgoBdu3ZhzZo10rGI3qh4jNnfI04oKIqSlMQaVXggm9HRUXR0dFQd6mQYBu3t7SCEIBwOw+PxwGg0yqYfPPW6e71etLe3N0yIthTjKWYYN6KRbYyz2MBMFaEAkLe0YSpih4ZG1tRUgnxZ0J32dlBc/vNA2Mk6wUZAXJetBEanwUGr8k+6DlrVW3GoWOSKK67AnXfeOa0lGkVROPPMM/GVr3wFF1ww2c923333xcDAAN566y1oNBqwLAuPxyOF4Hmex+LFi/Hkk08iGo3C5/Nh/fr1OPTQQ3P2fcQRR+DB/30Q/4n+B7/6/a/AEx5j8TFpfTNbc/i6667DD3/4Q2kNNZVK4Ve/+pVkRMXM6UQigZtuugmdnZ3Yf//9pdphsYRJXCcV/4nHLD5v++67Lz755BO8/fbbAIBAIFDVeVWKcDgMnU4na4hXnLz39/fDYrFgZGQEw8PDSKcra82Yj3g8LpXdNQqleLKNLEihGtkZiMfj02bN4XAYDoej6OfEtZVoNKrg0TUHk0ke+W81SkdDY619s+p8mEymqsooDjtlDyw8sR+MbvK3MjoaC0/sx2Gn7FH1sX3mM5/BRRddlPe9c845BzzP44QTTgAwOeA8/vjjuPzyyzF//nwsWbIEH374oTQp1Gq1WLJkCc466ywsWbIEy5cvxy233AK32y2FZQVBwBXfvALPPvMsPr/y8wgHJ5uO+5I+BDIBSVhCNKDz58/HbbfdhlNOOQX77bcfDj300BwBjm984xtYtGgRFixYgHA4jGeeeUYyoieffDLOPvtsPPfcc+jp6cHWrVsBAJdeeimWLFmCnTt3oqenB0899RR0Oh0eeeQRXHzxxTjooINw9tlnV31u5SaTySAcDiuak2EymdDb24u2tjYEApMlQNFotKpJvSAIGB8fl6X2XE64FAtEuMmIWAG0Wu20KEijoPaTnYH//Oc/cLlcOWFfj8dTUtYrIQQej6dh1jbqBcuyGPnzO9C8OX3CYV3RW1ft1ql4PB5J2abS3qRchkcikvlvCFn58NX/+3//D++++y7uvvvuqvclGliO5/Cf6H/yDtoURWEPyx6gqcnJRLaIRr3FNOrdT1bstdrT06NoVvBUBEGQkgfNZjNcLldJ4d7sXqzDw8NwuVyKJFhVAhEIIusGEN04BHAElJaWkiWpKSVxmUxmWitBJVH7ycpEvqQnsddmKQMJRVGSd9QIWXr1IJ1OY3R0FN1nHICEY/jTB0ZHQ3OIC9zCxmqZxTCMtNZY8T50moqTnMrlK1/5Cl5++WX885//lGV/4vonD76gV0QIAc3Q0Gkaf1201oyNjSledpMPmqbhdrvhcrkQj8cxNDQEhmHyNorPRywWA8MwDWNggcmmCtH1n2bsi8mSAKZNzBtZ9Uk1skUQszizb9JQKDRjqDgbt9uNoaGh3dLIJhIJTExMoK+vDzRNw378HET2ptBpb5fqBAcHB2FOmeueCSoirsuWc43ryW9+8xtF9svQTN6M603/3ISffP8nMGg+vV5HHXUU7rnnHkWOo5kQ1/PrKaJPURQsFgssFgsymQx8Ph84joPb7Z6uWJdKQRgdBed2w+fz1SXiJggCOI4Dx3FSMh3HcWBTGQgbh/J+Jl9ThUau4lCNbBHi8fi0GWk8Hp+WfFIMMZmjUCea2UokEkE4HEZfX1/uA8DQYFyfDtDd3d0YGBhAf39/Q2QzigklzWJklYKmaLgNbviSvpzXj1x5JE47+TQ4GAdompZKZ3Z3OI6rm6EqhE6nQ1dX1zQ1KYfNBt999yGw9iGQVAr/1uthP/984BvXABVm54pr8/kMZnYd81TE8ivxn1arhcFgAEU4+AolS9a5qUK5qEa2CPF4PCfLjmXZiuTVWltbMTIy0nDNjpUiGAwilUqhp6cnb1Z2NjRNo7OzE8PDww1xfvJp9u6uOBgHYAT8qUndaZqi4TK40GZqy9FaFiUhd1djSwjB8PAwurq6GvIcTFWT+s/tt4N77PFPN0inEf7tb8EwGrRdc01B71L8u9AygkajyTGYRqNReq3cNo+E4UFp6bxNFRopWbIUVCNbhEwmk9N5JxgMVlT3Kg5CldZgNhMTExMghKCzs3Pae4W8eYPBALPZLAl41BuxA83ujJh70K5vh52xg2ZoMDQjJTsBk+dJFAHIZDKzrttOqfh8PjgcjoZ/timKgkWnw8hTT+d937/2IaRWr4Ymyzhme5fi37WYSFBaDSxLu3PWZEUsR3YXrKsv1sCiXqhGtgBigXz2OkYymSyrIXg2ra2tGBsbk+TVZhuimLher5/WElAklUoVrL9zuVwYGhpqiLC6uC5bTqeg2YQgCBAEATqdTvJgCyU5ieE+jUYjeTuiYMTuQDKZRCaTqXhcqCWEEEQ9nsne1/lIpdBhMEA3RY+9XtiOnQy9xzYNT2uqkA8xabHWkqwzoRrZAqRSKalAHoDkhVY6SxIvvBhynk0QQjA0NAS73V40nT2VShU0wADQ1dWFgYEB9PX11dUjMpvNGBsb2y2NrJhRL3plpXoGohqTODkVM7Rns7EVBKEhutTMBMuyCAQCk5NcsxmUwZDX0FJGI5gG0lunaAr24+fAtrIXgSEftHYDbK78UrbAp4IUjTa+zt4noEri8XjOxao0VJxNa2srJiYmZt6wiRAEAR6PB263e8Z6sZn6QtI0ja6uLoyMjMh9mGXx/9t79zBH6jrf/12VezqXzqWTdLrTAV3k4gyMgAosCMNt1IVROS5447LrunhcxjOCLOrqWdfHB3V/K+iBRdz9saKwLsfjAyO7x5866AgzjqPAMAKKArtDkk5f0t1J536pSn1/fzTfIp3OPXVLd72eJ89MpypJ5Zuq+ny/n9vbaDQO5S7malVkUwvgasN34lFaT/bIkSP4zGc+IxrWxm473ZorPPLII9ixYwd27NiBM888E3fffTdqtRpuuOGGjnqyx48fx86dO3Haaadh+/btYkOQhx56CNu3b8e2bdvwvve9T9LORlIwNzeHyclJTU4kqFRkPB4Xex5Ho1FMTE3Be8P1LV/jeN/7wGoky78RxmSAa9qLYrXUcT+tNqTQ3tmhEfL5/DqjUalUhi4zsVgsYheczQDP84jH4wgGg5K1YbNYLHA4HFheXu6+s4wMkgAlCHUceugB3PMXH8D/u+cvcM9ffACHHnoAgjD47+31enHkyBHxnPn+978vm55svV7H2Wefja9+9avic73KjFE92X379uG3v/0tnnzySbhcLnHCcscdd+DYsWP47W9/C4/Hg927d4uvveGGG/CFL3wBv/vd7/D444/DYrGAEIKbb74ZP//5z8UWig8//LD0X3pAMpkMbDabZkrPKOVyGclkEonEWiwzEolgampq3XFO7NkDz19+BHj1OcZmg+/GGyFcc7WovqM1eqmD1WqtrG5k21Cr1cR6t15ExnvF7/dvitVsrVZDIpHYcAG3o5NuZTMejwfVahXlcnnYwxyYsbGxvj//8Pe+i1898r/Bv7ri4qtV/OqR/43D3/vuwMfBMAwuuOACPP744wDWVotXXXUVgLUyqZNPPlmcDLz88ss499xzAQC//vWvccEFF+DMM8/Ef/tv/63l6vUrX/kKtm3bhu3bt+PBBx8Ez/M4dOgQ3vve9wJYa6xw+eWX481vfjM++9nPdjzOr371q/ibv/kbnHjiiQDWpANvuOEGsCwrZh7TVcYXv/hFzM/Pi0bXZDLhggsuALA2qaBlc4QQlEol1Ot1lEqllsl0alCtVpHP5zWRpAesXVtLS0uIxWLI5XIIBAKYmZnB+Ph4S1c/YzDAfMMNmHpsP17/2H684ZeHEfjEXkSiUaRSKU2uBoHuSlla7V+sG9kW0JovaljT6bRkajo2mw21Wm2ky0QqlQrm5uZEge5e6DehaXJyEgsLC6qt+h0OR199jLlaFU//cF/LbUd/+IOhXMdXX301vve972Fubg5ms1l027pcLrz5zW8Wuz195zvfwXXXXYdarSbqyR49ehTnnHMO7r777nXv2agn+/Of/xz/83/+TywtLa27Kf/d3/0d3vnOd+Kpp55COBzueIy/+93vOmrDGgwGWCwWsfTn9NNPx+9//3u89NJLcDgc2L17N84880zcfvvtANZuqHfffTe2bduGyclJOByODXJ9akAIWetgpnJyEC3HSSQSmJ+fh91ux8zMDILBYE/XZLFYhNPng3l6WnQRMwyDSCSCZDIJnufl/gp90623uFarAnQj24Jyubyu7q9Wq0maBOPz+YYSCFeTYrGIVCrVd3JSv+52Gp9NJpOqKBmZTKa+bjSl1Yy4gm2Gq1ZQWs0MfCznnXcefv3rX4t6so382Z/9Gb7zne+AEILvfe97eN/73oc//OEPePbZZ7Fz507s2LED3/rWtxCPx9e9rlFP1uVy4eKLL8bTTz+9bp/Dhw/j6quvBgB86EMf6niMvSZIGQwGMWuZ53lUq1UcPHgQ//iP/4hf/vKX2L9/P/bv3w+O4/BP//RPeO655zA/Pw9CCB588MFehktW5ufnMTExoVpiHm1TGo/HwXEcpqamMD09jbGxsZ6TMgkhbX8vg8GA6elpzM7Oam4h4HQ6B1bJUhPdyLYgn8+LBkGOkpKxsTGUSqWRk8HLZrPIZDKIRCJ9J3sMEtOmknlqxWf7ySS3j3tgbDMRM1mssI8P7glp1JN9z3ves24bNY5UT9bj8YAQgjPPPBPHjh3DsWPH8Lvf/Q733HNPy/emUnStGvt3S1RrhOrJ9gIVX9+2bRvC4TDOOussTE1NwWKx4J3vfKd43EajUZzMXXXVVTh8+HBP7y8XuVwOBoNB8RapgiAgnU4jFoshnU7D5/MhGo32LALQTKlU6phDYTKZEAwGMTs7q6l7FNV7HjV0I9uCQqEAt3stVVyKrOJWeDweZDKDr26UZmVlBeVyGVNTUwOVMfVzw25kfHxc1CFVGqvV2vNs3mS24Kx3vrvltjPf+S6YzMN5QgbRkz127BiANe/Dyy+/vO51559/Ph5++GEUCgUUCoWWerLnnnsuvv/97wMAvvvdznHlVnqy995774b9eJ7HZz/7WYTDYezYsQPnnnsulpaWsLS0hEqlgscffxynnnoqpqam8Oyzz4rXyE9/+lOcfPLJvQ2WDHAch0wmg0AgoMjnEULERv/JZBImkwkzMzOYnJwcuukFVZDphM1mg9frxfz8/FCfJTXd4rJaRDeyLaBJT7RmUI66K6fTiVwuNxInzOLiIgghCIVCqnRTCYfDSKVSisdbxsbG+vp9zrv6A3jre66BybK2YjdZrHjre67BeVd/YOhjkUJPtpGzzz4bV111Fc477zxceOGF+Lu/+7sNiUV/8zd/g3//93/HWWed1VUcncZTd+/ejdNOOw1vectb1m3/+Mc/jtNPPx2nnXYa0uk0fvCDHwBYK5e6/fbbcemll+Lss8/G61//euzatQuTk5P41Kc+hfPOOw/bt29HNpvFjTfe2NeYSQUhBHNzc4q0TeQ4DouLi4jH4yiXy5icnEQkEoHT6ZTss3u9pzkcDthsNqRSKUk+VwpsNlvHhESGYTTn5tb1ZJsQBAG///3vcdppp6FUKqFYLMrWzSWTyayp03SZVaoFvbmMjY0N1TC/Xq9jYWFhqGSRWq2GhYWFnnR8paJSqeCll17CSSed1Jerm6tVUVrNwD7uGXoF2wuD6MnyPA9CSMebrVptQAVBAMdxAwkQyKEnu7i4CJvNJpt2NSEE2WxWdEf7fD7ZSoN4nkcqleqayNZIKpWCyWSSxaPXL9VqFaurq22F5WlNsNyNZHQ92SEol8tiCUEmk2n7Y0rB+Pg44vG4Jo0sIQSJRAIejwdO53Car1LEtc1mM9xuN5aWlhRz2VH6nYeazBa4A8qIRw+iJysIAur1umZ77bIsK9aUqy1AUCwWUa/XZTGw5XIZ6XRafH8lJpC5XK7v7xIIBESXtZoyfsDafaBTUxLakEJL3dp0I9tENpuF3W4Xy3jkFF9mGAZjY2PI5/NDGzIpqdfrSCQSCAaDkiR9depZ3A9utxtzc3MoFouKJZ9oPQbUr55sY9vETjf0dt/5xz/+MW677bZ1z8mlJ6u2AAGtP5WybWK9Xkc6nUapVILVakUgEFC0DWCxWBxoRRoOh5FIJGA0GlVtwEET9NplR2ux65NuZJsoFouYnJxU7Ebu9XqRSCQ0Y2R5nkcikUA4HJZsNtitZ3E/TE5OIhaLwWKxyDoBotCb/GaB9hTutmJq1+lp165d2LVrl1yHtwE1BQho28RhV5eEEBQKBayuroJhGHg8Hvj9fsVX5p1Kd7rBMAymp6cRj8cxNTWlan9gq9WKSqXScgFgNpv7qm9XAj3xqQmO42C327G6uqqIcDfLsrBarZpoZ1atVpFIJDA9PS2pu2XQzOJWMAyjaP2s1ley/UC1QHtZDWpNMowKEJjNZvA8L3tDl5WVFTgcjqGuAylqWqWkW+lON1iWFWto1Zx4Op1O5PP5ltu02FpRN7IN0IuWZqgp5Zry+/2q9+otl8uYn5/vq4uTWpjNZng8HkXbU466oaXNH3r9bXvtWaw0DMPAbDaLsma1Wk3y36ZSqaBcLg/kVpW6plVKeind6YbRaBRdx2pdExaLBZU2cn1anBTr7uIGSqUSTCaT4jFSg8EAk8kkiQjBIOTzeWQyGczMzEh+I+inZ3E/uFwuFItFFAoF2ZMxOsWARgFCCGq1GkwmU1/fQcvfl2VZmM1mCIKAWq0mZiIPiyAImJ+f7ysOS3ssZzIZEEIwPj6OmZkZzY2fVOWIFosFgUAAs7OzmJ6eVvx70s8blWtSe1NVFVldXYXT6UQ2m1U843diYkKV1ezq6ipyudxAXZx6QU4R9lAohOXlZdn7rI56XJbneTGJaLNBM5FZlkWtVhv6XJifn0coFOpprJSoaZUKGouXCrvdDrfbjYWFBcnesx+sVqvmpA/bsfmuuiEoFotwuVySxhB7hSajKJkZt7y8jGq1OnAXp16Qc3XOMAympqZkj8+yLNvX+xOuDj5dAeGGM8z1el3UZw2FQpiensaOHTtExZpu3HDDDXj00UdBCMEf/dEfdez7+vzzz+OSSy7BSSedhFNPPRV79uwBx3G4//77EQgEsGPHDpx00knYvXs3nn322Q2vf+9734uzzz574O86LLQnMrBmUAZJfslmszCbzR0nhe10Wv1+v2r9jHthkNKdbrhcLpjNZlUWB53isoC2wju6u7gBWiogV9F5N6iou9wKH4QQLC4uij1K5UTKzOJWmEwmeL1epFIp2b5Lr3EeIhDk9sdQOJQE4QQwJhaO86fguiwKhu1/EmMwGMTWiJ///Ofh9/tx00039fUe9Xq9q4uwVCrhXe96F+69915cdtll4Hke3/3ud8WVwnXXXSc2unjkkUdw6aWX4rnnnhPHe//+/ZowMDQT2Wg0olarIRaLIRAI9ORJqdVqWF1dxczMTMvtjTWtbrdb0aYoUjBo6U43fD4fFhYWFPf+Wa3WtsadahgrUX3QC/pK9lVo0lOzWLuS0BiTnO5PQgiSySRsNpsiephKeAWcTqdYJiEXvYi45/bHkD+QAOHW9iOcgPyBBHL7Y5Idx+HDh/HWt74VZ5xxBi655BIAax2J3v3ud+Pss8/G+eefj9///vcAXkte6mYMvvvd7+Jtb3sbLrvssrXjJgQf+tCHWsa63/Oe9+Ad73gH/u3f/g3Amsv09ttv76o3qzQejweRSATZbBaJRKKlh6jCVzCbn0WZK2Nubm6DR6edTqvb7R4pAztM6U4vBINB5PN5RSsk6MS31eRXa7Wy2jD1GiCfz4vp+mrGruhqVg6BakEQMDs7C6/Xq3rnFqkJBoNi/awc2dHUyLY7NwhXR+FQsuW2wi+ScF0cAWMabrVXr9fxoQ99CI8++ii2bdsm9hPeu3cvPve5z+Gss87Ck08+ib179+I//uM/AKCnFWazFmy3G/KZZ54pGvI77rgD119/vWbqvBthWRahUAg8z2NxcRHA2nnCsAz+8dg/4oHfPYBKvQILa8H73/B+7J3ZK07WaMtTtWpapWTY0p1u0LBNPB5HKBRSrNuSxWJp2d1Ja2U8upF9ldXVVRiNRtVbHFqtVnAcJ3kJhdRdnHr9TKUmLI3x2Wg0KvlNkWXZjhduPc+JK9hmSE1APc/B6B3OyDIMg9e97nXYtm0bAIhu+J/97Gd44YUX1h9PH2Pfyqj20g0qmUziJz/5CR577DHEYtKt1qXGaDRiampKrFt98JUH8d3/ek1VqCpUcf/v70e1WsUHox+E0+nE9PT0pkkUy+Vy8Pv9sn5GY7OKSCSiiKvW4XCsWxxRtNaQou+z6IknnsCVV14pKlLs27ev7b433ngjGIbB1772tXXPV6tV7NmzB36/H2NjY9i9ezdmZ2f7PRRJKZVKEARBEzNyqetmOY5DPB5HOBxWzMAC8mYWt8JkMsHv94urFinpFpc1OE1gTK0vJ8bMwuCUt/b46aefxrFjx/DMM8/gV7/6VV99ifvRggWAY8eO4ZRTThG1ak888UScf/75eO655/DOd75zkMNXBIvFgonJCTwce7jl9n3xfQhOBTVR0yolcimJNaO04Lvdbm+pyKM1d3HfZ1KxWMQZZ5yBu+++u+N++/btw69+9auWag979+7FI488goceegiHDh1CoVDAFVdcoWqZhCAIfdcRygU9eaQ4UavVKpLJJCKRiOIN4dWo+3U4HGAYpmPm4aB0MrSMyQDH+a0T1hx/PDW0q5jyyiuv4PnnnwcA0V184YUX4pvf/CaAtQSeF154oa/z+IMf/CAef/xx/PSnPxXjXPfdd1/LGPcPfvAD/OhHP8L73/9+/Mmf/Anm5+fxyiuv4NChQ9i+fTt++MMfSvAt5WO5vIxKvXUjg3K9jOWyuk1hpEbq0p1umEwmhEIhRZpVtLseDQaDpkru+jay73jHO/DFL34RV111Vdt9kskkbrrpJvzrv/7rhhlUNpvFfffdh69+9au49NJL8aY3vQkPPvggnnvuOTz22GP9fwMJqNfrEARBdVdxI16vt6uGZzdKpRIWFhYwMzOjSqadWs01AoEA0um05HGZbhev67IonDsjYMxrlxVjZuHcGYHrMmkazLMsiwcffBB//ud/jjPOOAPXXHMNAOCuu+7C/v37cfrpp2PHjh340Y9+1Nf72u127Nu3D1/84hdx8skn401vehOeeeYZ0Q33ne98Bzt27MAb3vAG3Hfffdi/f7/iSkhS4TF7YGFbxwxtRhv8NnndqkojR+lON6xWK3w+H+bm5mQ3tGazecOqVQsLpUYkv/MKgoBrr70Wt956K974xjdu2P7000+D4zhcfvnl4nPhcBjbtm3D4cOHWzYfr1ar6wqPc7mcZMdLuDrS8RSYOjSVDORwOLCysgKfzzfQSZPL5ZDNZlXtPKNGvTHwWnx2dnZW0vhst7gswzJw7zoBrosjqOe5V13I0qxgP//5z4v///Wvf71uWzAYxMMPPyzqv9Lve//994v7vPLKKx3f//TTT8eBAwfAcdy6xhU33HADbrjhhq7Hd8IJJ+Cpp57q6buoRTabRSaTwQdP+SD+5Xf/smH7+9/wfliN6inMyIFcpTvdcDgconatnGWCDocDhUJB1jLBYZH8DviVr3wFRqMRH//4x1tuX1hYEHvPNhIMBtt2D/nSl74Et9stPiKRyNDHSQSC7I9fwdwXjqD6zy/D/f0scj+JgQjaKGJmGAbj4+NYXV3t+7WZTAaFQkGVlmdawWg0wu/3S9qRptd6WcZkgNFrlczA9gKNu0mhGLPZzhmqLMVxHKLRKD5+1sfxke0fgc24li9gM9rwkW0fwbsD75YlzKAWcpfudGN8fBwsyw7tkeuE3W7XVJJTKyRdyT799NP4+te/jqNHj/b9w3Y6GT796U/j5ptvFv+mbQCHgdY0UhieiH+7d50w1HtLhdvtRiwWw/j4eM/jmUqlAKBlLFxJlMwsbofD4UCpVJLUZabFPsY8z4Nl2a7jvbKyItbWUux2Ow4fPrzuOS19t2EghCCTySCfzyMcDouhKwNjwMfP/Dj+8vS/xPPHn8dp0dNgN69pSM/Pz4PjOE2vjHpF7tKdXpiYmMDc3Jxs/eDbdWOjIi9q34MAiVeyBw8eRCqVEmOARqMRsVgMt9xyC0444QQAa/1ma7UaMpnMutd2citYLBa4XK51j2HoVtM4bDs8qWAYBi6Xq6fZNb1BGI1GTcTLlM4sbsfExAQymYxk2YZa62MsCELP3W18Ph+OHTu27tFoYLXUim5YaEY9wzBtlaWsRitO9J4IvNr7hcooCoKAhYWFkR8PKVR3pGBychKZTKZlJrAUtKqL1VKtrKRG9tprr8Wzzz677iIOh8O49dZb8eMf/xgAcNZZZ8FkMmH//v3i6+bn5/H888/jvPPOk/Jw2tKtpnE5nkKhUADHcapfaB6PZ8OEpBlCCGZnZzE2NqaZGbhaSU/N0PisVEkYBoNBkfKEXiCEgOM4ybLGtbZCHwRCCOr1OpaXlzE1NQWPx9PxO9lstg2divx+P+x2O2ZnZ1W//odBqdKdbtAa2sXFRVlKa2hcthEtGdm+3cWFQgEvv/yy+Pfx48dx7NgxeL1ezMzMbGjVR1O6Tz75ZABrLtAPf/jDuOWWW+Dz+eD1evHJT34S27dvx6WXXjrk1+kNWtPYytAyJhZ2vxPVWg35fH5di0MqsG6xWGCxWMSm/nLCMIwYdxgbG9uwXRAEJBIJseZYK8jds7gfjEYjJiYmMD8/P7QbXUt6lRzHSXoOjrqRFQQBHMeBYRhMTk72tLq3Wq0tY4Yul0v0xEUiEU30Zu4HpUt3usGyLCKRCOLxOGZmZiQdz7GxMczNza3L8zGbzaNrZJ966ins3LlT/JvGSq+//vp12YyduPPOO2E0GnH11VejXC7jkksuwf3336/YiUxrGhtjshTH+VMYczvQylzV63VUq1VUKhVFDbDP58Ps7CxsBgP45WUY/X6wVit4nsfs7CxCoZAmVo2NqJVZ3I6xsTGUSiVJGplrIS5br9fBMIyk14wgCCNnTIDXBOkJITCZTH2dd50Ulux2O8LhMOLxOKamphSvMx8GNUp3umEwGDA1NYVEIiGpdnWrvuImk0kzCVF9G9mLLrqor5l8q9IBq9WKu+66C3fddVe/Hy8ZtHax8IskSE0AMTJwXTDdsabRYDDAbre3TCaQ0wAzhKD+nQfw4ve/D1KpgLFa4b72WpTftRuRaFQTLqFRwO/3I5FIwGazDXXDpBd1O4PEcZwoJj/Mb1Ov13HWWWcBWMvKpxnTTqcTjz32WNfvcMMNN+C9730vrrjiCpxwwgl4/vnn25apPf/889izZw9mZ2dhNBqxc+dOfP3rX8e//uu/4q//+q8RDodRLBZx6qmn4otf/CJOP/10AGstHW+55RYIgoBgMIiHHnpIUQ8GXb0ajUYYDAZUKq0bTXSj3aTJbDZjZmYGiUQCExMTmvIWdUKt0p1umM1mBINBJJNJSasfjEbjutX7SLuLNwuNNY2zL8bAOEwYm5ocSJIMkNcAL911FyoPPij+TSoVrP7zP8PLMDDd/ImBjldOtJBZ3Aoan00kEkPVz7IsKwqhNyIIAg4cOIAjR46I8bBzzjkHO3fuHGg8Wknd/dVf/RVqtZqk3cmo1N1dd92Fd77znSCE9Cx19z/+x//A97//fZx88sm47bbb8M1vfhOf/vSnJTmuTjSuXhtrgweBNjRo19jeYDAgGo0imUyC4ziMj48P/FlKoHbpTjdsNhvGx8clCd9QaFyW/jb9akDLifbuhArDmAyAywTvhK9rgtGgUAPs9XoxOTmJSCQiPkKhEGw2G2q1GpaXlzE7O4tEIoFEIoFkMonluTmk7/92y/fNPPAAhAFn7nKilcziVhgMBgQCAczPzw/8Hu0u4AMHDuDgwYPiDJrjOBw8eBAHDhwY+LOaOXjwIC644AK86U1v6ip11ytU6o7mRDAMgw9+8IM9Sd01trAsFAqyqEc1Q3WfWZYd2sACa27hbjJtdIJWrVaxtLQ01OfJjRZKd7rhdDphtVolG8uxsTFZpS6HYcsbWUqrLEMl6GaATcUiSBtDSspl8BIKCUiFVjKL22G322EymQZq9EFpToDiOA5HjhxpuS9d2Q4Lz/O4/vrr8S//8i/4zW9+g//zf/4PgNek7p566inceeed2Lt3b1/v+7vf/Q5nnHFGz8bqzDPPxB/+8AcAwDe+8Q28/e1vRzgcxrPPPotrr722r8/uB0IIarUa6vU6zGazZPFjm83WU3kJwzAIBoMwGo1IJpOaWSk1o5XSnW54vV4IgjDUdUjRUtZ/M7qRfRWGYVr2wVQTg8EAZyQCpo3BYmw2GGWWsBoErRtZYC0+m8vl1rXr7IfmZAta8tUKGqMdBkIIBEHA61//+pZSdx/+8IexY8cOfOQjH+m7y1W/7sVG43LnnXdi//79mJubw7nnnosvfelLfX12r9DVq9FolGT12ojRaOyr9tnj8cDtdiMej2vyxq6V0p1eCAQCKBQKkiQpteotroWJkG5kG+ilJlVpOIYB++53tdxmv+ZqsBo0ZlrLLG4Fdf/Nz88PdKNsNrKdkpxMJtPQfbG7Zf5SqTv66IfTTjsNx44d6/k3o1J3S0tLeOGFF/CmN70JAPCnf/qnG7pHDQtdvQqCALPZrJnzyuFwIBQKIRaLaSbBBtBe6U436HW4vLw8cNIapbleVitqPNo4YzWC1WodeGUjB6urq1hcXMTrPvMZ+D56I5hX45yMzQbfR2+E7c/+TBGli82KwWBAMBgcKD7bbGRpklMrzjnnnKFWFjRRjmGYrlJ3giDgueee6+v9P/jBD+KJJ57Az372MwDoWerO4/FgaWkJx48fBwD89Kc/FevhpaBx9Sq3DKXBYFiXkNgLFosFkUgEyWRyaAMhFVos3ekGwzCIRCJiS8tBaTayWskwHp0pj0JYLBbV3Z20RaLJZEIkEgHDMAjs3Qvz+9+PejoNz4kniivYYrGIWCyG6elpTcxgtZpZ3A6bzQaLxYJMJtN3yUNzvSytH2+VXTwotG0i7U1Mpe6q1SoCgQD279+Pu+66Cx/96Edx7733gud5XHfdddi+fXvPn2G32/H9738fn/zkJ/Gxj30MRqMRF198Ma677joAa1J3jz32GEqlEk455ZR1Unf33HMPrrzySrEG8tvfbp2k1w+0kxUN4SiRJUs1nPvtr2s0GjEzM4PZ2Vl4PB5Z+vP2g1ZLd7pBm1XQzP9Bs/EbV660IYXaSZgMGcFlEA3sZ7NZSWZtiURCFByo1WpYWVlRJEuyFbVaDXNzcy1r8tLpNMxm8wbXI8dxSCaTCAQCqmcVFgoF1Go1zXR76pVEIoFAILCujKNSqeD48eM48cQTW066mmXhGp+Xok6WukrlNjRSt2ccBp7nUa/X+24qQen2m7WjWq1idXV1YFk2OjG2Wq2qnfuEELHRw6hSrVZFDexBzvnZ2VlMTk7CYDCgVCqhVCrBL0PeSj82aHSWHApBk5/UmHvkcjnMz89jenq6ZdF7vV5vGZczmUyIRqPIZDJYWVlR4lDborYXYFAGic+2i/mYTCZ4PJ6hk09ofE3ulZwWaioFQRBDNWrEXodNetSCuMAolO50w2KxwO/3D5y97XA4xCQqrbRW1I1sC2gLPqUghGBhYQHlcllUMGpFOyMLvJZAwDAMZmdnVct6HFUjy7IsgsEg5ubmen4NldOSg3q9DkKIpG0OV1ZWsGPHjnWP8847T1VJMNpUgrrXlZhUtEKqz1RTXGBUSne6MTY2BqfTicXFxb5f2xiXHSTOLgfqB/E0yPj4OBYXFxVpoUb7D/v9/q4ZqJ2MLMXr9cJqtYr9VpVO5R+FzOJ22Gw22Gw2pNPpnlx+9MYs9UqQGh6p3bdU6q4ZmlykNLQlosFgaNttSUmk0iBVS1xglEp3uuF2u8FxHFZWVjaIznSCtlcEtKOLPJp3Q5mhdXNyz0QLhQJmZ2cxNTXVU4lHrzcAu92O6elpJJNJzXZB0So+nw+lUmldtmin80CO9m30ZqnkTULJiRGNAdOJhJQGfpjfwmazSZYl3CguoETt/aiV7vSC3+9HrVZDLpfr63WtBAPUZHP9KhJC3Q5yZAsSQpBKpSAIQt89dHvd12g0IhqNim5ov98v+0171DKL2xEOh8WMbYZhsLS0hImJiZbjV6/XIQiCZCsIOgtX8iZB61CVQBAEse+zwWCQvGSOtngcxODQrm9SxTUbxQXkTkocxdKdXgiFQqJoRa/jNzY2hmKxqHqmN0U3sm0YHx9HMpmU/Ieq1+tiur/cFwXV1VxdXRVXzHIaQS33LO4HlmUxOTmJhYUFTE9PY3Z2tqWaFEWqVQQVHFd6RaLUKqgxziznhI+6a/ulnbbsMDSKC9RqNdnEBUa1dKcbVPA9FoshHA73FEJxOBxYXl6G0+nUhCylbmTbQI2RlEkhpVIJqVSq55NFKsbHx2G1WhGLxWTVxaxUKiOf3UixWq0YGxtDtVrFSSed1DFLcW5ubmg1EUII5ubmEAqFFNV0rVarKBaLspadVCoVrKysYHx8XPY8B1pbOwhyKbfQpMRUKiV6RaRE66o7w0KbVcTjcUQika4TqMYmFPT/apan6Ua2Ay6XC7lcbujZJyEEKysrqNVqQ0msDYPVal1XNC/HKrpSqYxcfWwnvF4vZmdnYbfbO67Q7Xa7qAgzKHNzc5icnFR8kkIbMMiRES4IAhYXF0EIwYknnjgyoQQ5DBYVF8hkMkgmkwiHw5JKFW6WyW07DAaD6FXqRfCdrmBNJpNYa64Wo3HWqwQ1ssNA3cNGo1HSC2sQDAYDZmZmUCqVxJuflIxyZnE7wuEwFhYWOvZAbW7n1i+5XE5UY1IauVz8xWIR8Xgcbrcb4XB4ZM4LuUVC5BAX2CylO90wmUwIBoM9lUfZ7XYUi0VNtFYcjTNfJRiGAcuyAzeZrlQqiMfjCAQCmhF6ZhhG1LBNJBKaaKCtZViWRTgc7lg/Sy/oQeB5Hul0WmxTqDRSx4AFQUAymUQ+n0c0Gh25FVYv2rLDIrW4wGYq3emGzWaD1+vt2m/c6XSiUChooiGFbmS7QFtn9Us6ncby8jKi0agkNYBSrzpdLheCwSDi8bgkZQubJbO4FRaLRUymaMWgsTxCCJLJpNhEZNTJ5/OIx+Pw+XwIhUIj+Z161ZYdFqnEBTZj6U43HA4HbDYbUqlU232oR4K6i9Vkc94VJaRfV6AgCKI7Y3p6WjLDI0dXHovFgmg0ilQqNdBEopHNklncDo/Hg2q12vYG3FgE3yvLy8sYHx9XbRUiCIIkhpCGRMrlMqLR6Eh2/KL0qy077GfNzMwglUqJpUf9sllLd7pBM6k7SZMyDCPGZtVEN7JdYBim5/Zc1WoVsVgMPp+vry4lvdBLt6dBoOoX1WoV8/PzA5+Qo9pOsR9oWU+rm3C/k7FyuYxaraZqLE2K3yybzSKRSGBiYgKBQGAkV69qQq+/fD4/UPlQsVhUpDOdFgkEAiiVSm2vO6W8Et3QjWwP9CLmTrVfZ2ZmZFnRyWVkgbWJRCAQgMPhQDweH2gmvxWMLI3PtmpePjY21rORpU3k1VJ6olQqlYHPVZ7nkUgkwHGcZCERraB0z9tBxQU2e+lOL4TDYaTT6ZYud6fTObCHQEp0I9sDtBNMK2h9I8dxsvYpldPIUpxOJyYnJxGPx/ueAW7GzOJWWCwWuFyuDfFZg8HQc7YoLddRe7zK5XLfEyNCCNLpNJLJJEKhkCKdxJSGassqTb/iAluhdKcbtFnFwsLChgQns9kseUexQdj8d0UJoAXu1WIFfLoCwq2t9Gq1GmKxGMbHx9u23ZMKJYwssHZiRqNRrKysdF29b1XGx8fBcdyGiVc76btGVldXYbVaNbHqFwShr3OK4zjE43EwDIOZmZlNm9HaaVItNy6XCz6fD7FYrOu5tFVKd7rBsqxYQ9s4ZvR+PEyFiBRsrbS0ASECgfloEUu//E+AJ2BMLExv9qO83YLpyLRiLemUKqimJ+3y8rK46uo0gdjMmcXtmJyc3KCyQuOy7W58tNl5JBJR8lCHhjZTKZfLmJqa2vTZrHLXynajUVygU4e2rVS60w3ah4A2q6D3K5vNhlqthlqtplpi5ta6Mw5Ibn8M5YPzAL/mwiGcgNrhFMZfHKwR+SDwFQ7I8eIqWgn8fj/cbjdisVjHGNVmzyxuBY2jNcZnO8VlaVhB7YYklF6z1Wkyn8lk6qml3WZAC78PFReYm5truareiqU73bBYLJiYmFjnbnc6naJesVrov1IXCFdH4VCy5bb8oSTsb5uE0WqW7cIkAkFufwy1g7Oo8QRZEwvH+VNwXRYFw8p/MxgbG4PZbMbs7GxbJZHN1LO4H8xmM9xuN5aWlhAIBDqWfywuLsLv92vmxtgtHksIwdLSEmq1mqKaqFpBKm3ZYWgnLkC4OlYTy3B6ustjbjXsdjt4nsfi4iJCoRAsFotuZLVOPc+BcG0SWjgBS7FFCI7emhGwLAuDwdD238b/U6Od2x9D/kBCfA/CCeLf7l0nDP39esFkMiEajWJubg7lcnlDedJm61ncD263G3Nzc2IpBdWybLw509VtL5rBStEps7hSqWBhYQE+n0+1TlRqQ7Vl1Z48NooLpBZTsBwroXAoCcIJqJpYCApOuEcFl8sFjuOwvLwMv98Pg8GgqvtfN7JdMDhNYExsS0PLmFlMvn4KjKn7LJ+m29frdVGDlP5L9TwbtxFCAF4ADs6i1eWTPzSL2ul2mGxmmEwmGI1GGI1GsCwry6qaXuzpdBqzs7NiP1rC1UGyNTBBsmWDDzQ+a7FYNmhZ1ut1sfOXliiXyxuk0QghWFxchCAIPTVh38xIrS07DFRcYPEHv0f+l0vi82pMuEcFn8+HhYUFZLPZtXrZfBF8uvLq/VxZr4xuZLvAmAxwnD+1bjVJcfxxbwYWeK37CMuyPScr8OkKFvh4640cgVUwgbCsKFfG83xbd6XBYBANcaNRNhqNfRllr9cLq9WK+CsxuP9AUDo8D4YTMGeaV9SNrSUa47OTk5NYnpuDJZuFwedDcmlJM3FYCsdxyOVyCAaDoiGlMowTExNbtrlBI3Joyw4D4ergn1ppuS1/KAnLHwfBmg3ifabxAWgjzqw0wWAQycQsjM8UYHhyGQt8EozC4TZAN7I94bpsbRWSOzgLhidgzCwcfzwlPi8X3VbRjoCr51V0vV4Hz/Pio1qtiv9v5epmGGadIW40zDabDe4/EBQeT4qr7K0+qzabzRh3uTD3//wDyt/7HrLVKmC1Yux918B0661qHx6AtWSnAwcO4MiRI2Jm6lvf+laccsopokLTVl69NiKXtuwgVKtVZBMrHcNW2fkVwGUSPWbND7loZdTbPaiXbdAH/bx+js35Qh25Xy6rep/SjWwPMCwD964TkHsDg5A7qJjLgTEZMPbHYeR/vtFl3O8qmhrIXhEEYYNRFlfLVQ74xVxLN3bu4Cxyb2AAY+cLqtMF1+82rVD99rdRfuCB156oVFC8/9tYsloR2LtXteOiHDhwAAcPHhT/5jgOhw4dQr1ex65du1Q8Mu2iRkelarWKfD6PcrkMQsiaQMWEC1yHCXfghEnF3aDtDHovDxoS6/fRFx3CbYVfJOG6OKLImOlGth+MLIxeZZsIVHfYYeNDqP46BVITFFtFUxHyVjV6ndzYDE8Qcgdh9Fr7vrgEQRBj0/28TkoGnRigVkP6/m+3fM/0t78D/0c/ClbFBhQcx+HIkSMttz311FO4+OKL9ZrLJmi9rNwtI1sZVafTCZ/Pt87A8xKEraREa5PcZorzWWTa3KdITUA9z8Ho1Y2sZpBKsaQfCCGoVCuIXnESyK7XYfH4PNyTXtic6iZjdHNjG5xrN2utX4StGHRiwC0ugrSRLCPlMhLPPgtWxV7F+Xy+bRkDx3F46aWXxGStTqECuRLrtAjVlpXayPZqVJtxXRYFz/MoH1kAOOXCVqNGtVpFKpWCAWxP9ym50Y1sj3Acp1jHJUomkxEzQBmTAc4pD0qVMmxQuaxAomQwLTLoxECwWrFitbY0tIzNhsjpp6u+kjWZTC0NrclkwkknnSSuZAkhYm3hhlBBm8S6RqPcaJgby9FGDapZ2pyF3S+DGtVmGJaB8BY3vBdMwswbVMmU1TLUuBqNRkxOTsJoNCJ7PofcgcRQ4bZh0Y1sj9RqNUWNLCEEuVxuXemH3W5HOp2WXEZvEDYkg5lYCKc7tuysmrVa4b3heqzc+80N27zXX6eqgQXWDOk555yzLiZLOeecc9a5ihmGgclk6tl9TBPrqFHmeR6VSkX8fysas92bV8taMcpGoxH1Ktd36YdURrUVlUplU4oyDEOtVkMqlQLLsqJxpYxdPI18PgdyLKdo0mojupHtEaV7XxYKBTidznUXkxYEiCkMy8B52QzyJ7MIuQIwOE1YyqygWCpqqumCkkzs2QMiEKS/8x2gUgGsVnivvw4Te/aofWgAgJ07dwIAfvnLX4pt+c4991zx+UHpN7GOutkbjTJdKbfLdmdZtqMLW2popzUcnMUCH+9Y+iGnUd1wXFtc2q6RRuMaDAZbTgpz+Rzcu07E6llpRZNWG9GNbI/UajW4XC7FPi+dTmNmZmbD80ajUTN9S6vVKqwOu5gMNjExgVgshrGxsS15I2AMBtj+4sMIfeD9GON5VKxW1F59XguwLItLLrkEr3vd6zA+Po5KpQKn06l42Q7DMGKHs16h2e7UMJfLZdEot5IYbI4rNxrmXuLKtNNaq9IP60WTihnVRmq1mp6chrXQRyqVAoC2xpVSKBQQiUSwurqqeNIqRf079YigpGGjfWVbXbC0o5AWJK6ahdoZhoHP58Py8jImJiZUPDL1KBQKYh9jM4BXXnlF9ptvvxiNRng8HgiCgNnZ2ZHwPHTKdm9Fc1y5VquhVCr1FFc2EBalQ7Mt98kdnEV5mwUuj1vx37VY3LpeIuA140oIQSAQ6HouUHUwta893cj2gVI/1vLyMsLhcMttDocDCwsLmjGyNCOV4nQ6kclkNLPaVprm7z0+Po5sNis2d9cStPtYtVqVvUxFaYaJK9eWiwDXOizD8AR+mwdGu/KromKx2Pa+sJnheR6pVAqCIPRkXCnN151arna9vYvGqNVqHV1pvQiDK0U7QxoMBrG4uKjCEalLK9UWt9uNbDar0hF1x+/3Y3l5We3DUB3qXrbZbHCFPGBMrW+NSpZ+NKO2KpDS8DyP+fl5zM/Pw+fzYXp6uq/k00KhILYIVXM1u3V+sSFQMtloaWmpq6uVynBpFYvFApZlUS6X1T4URWnVUJ5hGNjt9rY6s0rTXO9tMpnEBiA6a9AStVZXvVolavV6fcvIDdbrddG4er1eRCKRvj0tdLwa2zGqlTSqG9keUOoEpwo83VxcdrtddQPW7YQNBAJi/GSrkM/nW8bMfD6fZprNt/I+6KvZjdTPdsF0rh+Mee0WyZhYCGe64Lx0YzKiElApxc1MvV7HwsICkskkPB7PQMaVsrq6ui6kpqaR3XpBswFQqkaW6h92w+FwYHV1VdWLrtuYGAwGOJ1OzcYj5aBdwxKatNOcKKYGPM9vmMTRpgt6ecgaPM8jX8hjZvcpwDvX2u8ZnCYUKiUsLS8hGAwqfkyFQmHTJhMKgoClpSVUq1VMTExIUipZLBY3aFzrK1kNo4SRFQSho5B2I2azGdVqVdbj6UYvBsPj8WB1dVXTrm2p6GagtLJabBdH93g8yGQyKhyR9pifnxflCRmTAUavda3jmtOJer2OUqmk+DG1mhyNOoIgIJVKYXZ2Fk6nEzMzM5IYWJ7nN3Qa02OyGkcJI5tOpzfMvNpBTxg1XbG9GFmGYUS38WanXC53vEHQTkbt+gcrBcdxLY2s0+lELpfbUu79VlCR73YGbXJyUhS2V4rN5mGgxjWRSMDhcGBmZmZDLsMwrK6ubvCe6TFZjSO3kSWE9F0DZ7FYVF3N9jomdrtdrFPczOTz+Q3lTM34/X4sLS0pdEStabeSZRgGDodDMwlaalCv15HJZDq2LWUYBpOTk5ibm1PsuLpN4EYF6hZOJBIYGxtDNBqV1LhSWiUgArq7WNMQQmRNnc/lcnC5XH3NVh0OB4rFomzH1Au9Hu9WKOmpVqtdJx0Wi6VjMwQl6OR29Hq9W9plPD8/j8nJya7ntdVqhdVqxerqqiLHVSgURroJBSEEy8vLSCQSsNlsiEajsuWT0Elk82+ou4u3OK3cG92gMlyjgMlkgsVi2bSrJDpD7uVC9vv9WFlZkfuQ2tKp1pI2p6i0kezbzOTzeZjN5p6zWX0+H7LZrCLu/0qlMpLNQgghWFlZQSwWg8ViQTQalX2ykMlkWt5LdXfxFqZYLMJut/c901LzpBmkm9PExASWl5c3ZcyvWq32nDVMy6+0Og70d9pKCIKAlZWVvrJ3GYbB1NQU5ubmZP0t+5nAaQVCCNLpNGKxGEwmE6LRaNdQilS0c63rRlbDyP3DrKysDCxdR8UClGaQUhSGYeD1elVdxclFL/HYRrScyUsnT2qcV2qxsLCAUCjUtyGjPaDljLPXarWRWcUSQpDJZBCLxWA0GhGNRvsOgw0DTepr93m6kdUocoq1V6tVmEymgeO9VCxAaaiAQb+4XC4Ui8VN112o30mH1jN5tVJupATFYhEsyw5cv+xyucBxnGzNYUYhHksIwerqKmKxGFiWVdy4UjKZDDweT8ttekxWw8iZWdxLC8VOqJUNOkxD+VAotOmSoPotsWAYBi6XC/l8XsajGhyr1Yparbbp65tpKcmwzSUmJyexsLAgy3i1y5TVAoQQZLNZxGIxAEA0GoXb7VbNoHWa7OruYg0jl5Gl7rhhlGrUEgsYJtuaGufNklzTS1ZxK9RwGTf3Le6Ell3aUrG4uIhgMDi0UWBZFqFQCPPz8xId2WtosUa20bgKgoBoNIrx8XFVj5Nq7XY6Bt3IahS5jOywq1iK1sUCWrGZSnoKhcJASR0Mw8BqtSqaId5PwprD4UA+n9esS3tYaPKZVKtE2sAil8tJ8n6AshrWvUAIQS6XQywWQ71eRzQahcfj0cQkoFuFhu4u1jDtOuQMgyAI4DhOkoQGpcUCpJDbMhgMcDgcmpaA65Vh3HlKl/P005pP6y7tYSCEYHFxEaFQSNL3nZiYQDqdlixpTEvx2Hw+j1gsBo7jEI1G4fV6NWFcKd1a0uruYo0j9ck0TEZxM0rHZaWq2aOND0ZtFd7MMO48qhusVDesfldGm9VlnEql4Pf7JW8wI3VZjxaUdwqFAmKxGKrVKqLRKHw+n6aMK/Caq7gTupHdQhBCUCqVJLt4lBYLkEpJhmEYTExMqN5mcBg4jhu6abuSY9CvV4ZhGFgsFtVlFaWkUqmA53nZVogmkwkul0sSD4WaGrLFYhGxWAzlchmRSAR+v19zxpXSKau4Ed3IbhEG6e7UCaXFAqSUaxsbG0OtVlO9af6gSOHOo6LpStSlDhLj20zlPIQQsSZWTsbHx1GpVIaa/EoRlhmEUqmEWCyGYrGISCSCiYkJVY6jH3ppBqPHZDVKvV6X9ASjiQMul0uy9wSUFQuo1+uSxqhDoRAWFhYkez8lkcqdp1SXpUHk0rSiHiQFy8vL8Hq9iqwOJycnMT8/P/DkV+nSnVKphHg8jkKhgEgkgkAgoHnjCvSe3a+7izWK1JnFhUIBY2Njks+qtCAWMCgmkwlms3kkj1+q1YZSdamDHu9maLVYq9VQqVQkn+C2w2AwYGJiYuAJpFJJT+VyGfF4HPl8HtPT0yNjXCm9uooB3V2sSaQ2splMpmfN2H5QSixArpOUxiVHqVxE6niZ1+tFOp2W7P2kxGKxgOO4kU1SI4SICjtKQifUgyQmyi2vWalUEI/Hkc1mMTU1hWAwOFLGldJr20ndXaxRpDzRy+UyzGazLCeyUq6QYTo9dYJlWXg8Hs0amVZIvdKgWeJanWhoeRLQjXQ6DbfbrUrNaTAYxPLycl9NY+RsQFGtVpFIJJDJZDA1NYVQKKRactWwVCqVnu/PurtYo0hpZJeXlyVpPtEOJcQCpEx6asbtdqNQKIxMX2M53Hnj4+OarR3W+iSgHRzHoVgsSpps2A+DiLzLcZ1R45pOpzE5OYnJycmRNa6UfjyDupHVKIIgSHIichwHlmVlPamVEAuQ08gCo9UJSo7yCrfbrVkjC6wdn5QdjZRADTdxMxaLBWNjYz17AqScwNVqNczOzmJlZUU0rlrqIjUM/Yq36EZ2EyNVC8VOKNGUQoq60E5YrVYQQhSt+x0EucorGIaB3W6X5Xfsp29xO8bHx7G6uirNASnA6uoqxsbGZD1ne8Xr9aJQKPTUeESKySw1rktLSwgGgwiHw5vGuAL9N8XRY7KbmHq9jnq9LmsSA6CcWIDcJ+solPTI2YnH5/PJEvuUog8u7bc8Cs0peJ5HNpuVJdFwUHrtBjVMTJbjOCSTSdG4Tk1NaWKSITX9ZBUDurtYk0j1gywvL8Pv90vyXt2QUyxAqRPUYDBgbGxM025JOcsrWJaF2WyWXKVokBrZVoxKcwrqJtZSlyKDwQC/349UKtV2n0G9RdS4Li4uYmJiYtMaV8ogOt8jY2SfeOIJXHnllQiHw2AYBvv27RO3cRyH2267Ddu3b8fY2BjC4TCuu+66DUH/arWKPXv2wO/3Y2xsDLt378bs7OzQX0ZKpLgpEUK6Nq6WEjnFApRUBKGrOa0m2cjtNpfDkEn1+xkMBrAsq+nmFLlcDlarVXbv0SA4HA7U6/W2JXf9TuB4nsfc3JxoXKenpzX5vaWkXC737U4fKXdxsVjEGWecgbvvvnvDtlKphKNHj+Jzn/scjh49iocffhgvvvgidu/evW6/vXv34pFHHsFDDz2EQ4cOoVAo4IorrtBUZqkUmcXpdLovl8awyBmXVXKywDAM/H6/JvsaK6HvKUeXJSnVpLTcc7peryOdTivmPRqEyclJpFKpll6nXkMRPM9jfn4e8/Pz8Pl8W8K4Uvp1FQPquov7vure8Y534B3veEfLbW63G/v371/33F133YW3vOUtiMfjmJmZQTabxX333YcHHngAl156KQDgwQcfRCQSwWOPPYZdu3YN8DWkZ1gjSwhBPp9HNBqV8Kg6I6dYQKVSGUg3dVAcDgcymYzsq8Z+UardHZ1khMNhSd5PSk+E2WwGz/Oq9dftBO1NrCU3cTMMwyAUCmFubg7T09PrtnUb03q9jlQqBZ7nMTExIWu2vxYhhAzkZdzUMdlsNguGYcQ6taeffhocx+Hyyy8X9wmHw9i2bRsOHz7c8j2q1Spyudy6h9wMa2Tz+TxcLpeiF7ucYgFyNaLohBZLevL5vCLt7iwWC3iel8y7I1VMluLz+RTVwu2FQqEAo9E4EobHarXCarWuy9buVBZWr9exsLCAZDIJj8eDSCQyEt9Tasrl8sAetU1pZCuVCj71qU/hAx/4gNgzdGFhAWazecNyPxgMts0q/dKXvgS32y0+IpGInIcNYPi42yAuDSmQSyxACTdpM2azGUajUZGWkb3Saxs3KZBS1F3qVefY2BhKpZJm4uaCIGB5eRmBQEDtQ+kZn8+HbDa71rKyUkH2xZdgbzKygiBgcXERyWQSbrcbMzMzW9K4UgZVMRupmGyvcByH973vfRAEAffcc0/X/TvdxD/96U8jm82Kj0QiIfXh9nU83SiVSrDZbKr8sKMsFtCKQCCAVCqliZu50hMNmsimhe/eCi01z1hYWEAwGNS0m7gZhmEQDoXwyu2348VzzsXie96Dxbe/A6mvfQ11jkMqlcLs7CycTidmZmYUy4nQKoO6ioFN6C7mOA5XX301jh8/jv37969TvgiFQqjVashkMutek0qlEAwGW76fxWKBy+Va99Ayy8vL8Pl8qny2HGIBUkv+9QPLshgfH99wvqiB3B2vWuHxeDTx3VuhFSNbKpXAMMxIGqHMN74B/t8eAnm1ZItUKli595v4r9tvh8PhwMzMjKKSd1pm2HyITWNkqYF96aWX8Nhjj20wNmeddRZMJtO6BKn5+Xk8//zzOO+886Q+HMWpVqswGo2q9QWVY8amhnFpxO12I5/Pq64Ck8/nFU3+AgCn04lcLqfJ1Sw1bGq68wkhWFxcbDtB1zJCpYL0/d9uua2+7wewaiypTG0GdRUD6rqL+043LBQKePnll8W/jx8/jmPHjsHr9SIcDuO9730vjh49iv/4j/8Qg/XAWlsxs9kMt9uND3/4w7jlllvg8/ng9XrxyU9+Etu3bxezjdVmmPiVFuJCVCxAqmxStY0swzAIBAJYXFxUtQ9tpVKRvT1mMwzDwOVyiYl0WsPn8yGZTGJmZkaVz19cXBw5DVQKv7wsrmCbIeUy+OVlmJuyj7cqhBDU6/WB72kjVcLz1FNPYefOneLfN998MwDg+uuvx+c//3k8+uijAIAdO3ase92BAwdw0UUXAQDuvPNOGI1GXH311SiXy7jkkktw//33a0YVolarDeT353kehBDVS06oWIDb7Zbk/SqVimoqJhSbzYZ0Oq1KljPwmqtJjRmxx+NBPB4f2MhK0be4HQaDAUajUXb901ZUKhXU63XZWlzKBSEEhUIB6XIZsFiAFomKjM0Go4ZrfZVmWFfxSBnZiy66qOPB9vJFrFYr7rrrLtx11139frwiDHrDULKFYiccDgcWFhYkM7JSqRENSzAYxNzcnCqrJiWzipuhPYMHvdHI3a2L1vROTU3J9hnNEEKwsLCg2gp6EGijjFKpBKfTicjrX4/lP7sBK/d+c8O+4x/6ENgtnEXczOrq6tAhgZExsluBWq3W981MEATUajVNpNcrJRagNEajETabTZXYqBqf2Yjf78fc3NzARlZO74rZbIYgCLLI/7VjaWkJPp9vJNzElUoFy8vLIITA6/WuCzm4/vIvkc/nwT38CEi5DMZmw/i1H0Jp927NNWJRi2FdxcCIxWS3ArVarW/3aDqd1pTiBxULGPYmpEZ9bCf8fj9isRgcDoeix1Uul1XLGAfWJk4Gg2EgL4sSfadpcwol8hGq1SpqtZrquQ+dIISIJYcWiwWhUKjlb5BaXkb0058Ge+utSDz7LCKnnw7WagXP80gkEohEIptKom4Q5FS9UgLtTwNVoN8ZOSEExWJRkU5AvSKVWIBaMdB2MAwDn8+nuBKMFiYbg/YMlrJvcTuUquklhGhCiL0dPM9jYWEB8XgcADAzM9PWwObzedhsNhiNRrBWK9jJSdFFbDQaMT09jUQiAZ7nFf0OWmOYrGItoBvZNvRzQ81ms5LFP6VCKrEAtTOLW+F0OlEulxW7+QyaCCc1JpMJgiD0/b2VUlBSQtR9ZWUFHo9HEzkCjZRKJSQSCTEXIhqNYnx8vO19hBCClZWVjjkcJpNJNLSbMfzTC4QQzeSEDIpuZCVAi0ZWKrGAQWSllEBJcfdCoaBqPLaRiYmJvlfxcsdkKS6XS9a+4hzHoVQqaeZaI4QgnU4jFouhUCggHA5jenq6p6YYKysr8Hq9XSfzJpMJU1NTW9bQyqndrBS6kW2iX3dXoVCA3W5X3ZXYjFRiAUrqyPaD2WyGwWCQTT+3kWKxqJmuO1arFdVqta/GHEqp5TAMA7vdLovcIiEEc3NzkqkSDQPHcZifn0cikYDRaMTMzAwCgUDPq616vY5isdhzSZbZbMbk5CQSiYTqDVmURosLmH7RjWwT/RqVdDqtakJMJ6QSC9DaBIJCVXqUiANqKYuVitprEbmObXV1FU6nU7UJH61tjcfjSKVS8Hq9mJmZGUhpa5AOVTR5Kh6PbxlDuxlcxYBuZDfAcVzP2ZuVSgVms1lTN+BGhhUL0GIrv0ZYloXb7ZY1DqjFlTyNt2vx92FZFiaTSVIlKJ7nkcvlVFG1ouo+8XgclUoF09PTmJqaGjgZsFqtghDSNgTT6Te1Wq0IhUJbZkWrdtmcVGjTOqhIPyUSWmk+0Y5hxQJGoU5vfHwc2WxWtpuOVmNC9HtrEb/fL2n29/z8PMLhsKIelWq1imQyiWQyCavVipmZGfj9/qEn1J1WsQaDoet5bLVaEQgEkEgkNDnJkpLN4CoGdCO7gV6NLMdxYBhGc6ucRoZtJabFzOJmGIZBMBhEKpWS5f21WqMn9wp+GGgWtBSJOtlsFjabTZHJHiEEuVwOsVgM6XQagUAAkUhEsprsxpKdVrAs29OY2Ww2TExMbGpDSwjRXJhmUEb/G0hMr+UaS0tLml7FUqhYwCCMgpEF1m46PM+jVqtJ/t5KdjHqB4ZhMDY21jXJSM6+xZ2QYjVbr9eRyWRkz3mo1+tIpVKIx+Oo1+uYmZnB5OSkpIa9l5KdXlayFLvdDp/Ph9nZ2U1paLUqiDEIupFtopfZU71eB8/zmmrS0A4qFjAIWmtE0Qk5SnrU1NHthV6SjNSKKdtsNlQqlaEMAG06IdckoVwuY3Z2FnNzc3A4HIhGo/B4PLJ8Xi8lO/22Qx0bG4PH49mUhjabzepGdivTbUaqJYZtSqHVzOJmjEYjrFarpOUjWuvi1QzLsjCbzai0kUsDlKuRbcUwgvP5fB5ms1nySR4hBJlMBrFYDLlcDqFQCJFIRNYSrV5Ldnp1FzficDgwPj6OZDK5aQwtXc1reYLbD5vjWygIIQTlclkzdZPd2KxiAa2gjRqkutloNempkW5uWTWzowcVnBcEASsrK5Jq93IcJ7Y7ZFkWMzMzCAaDioxNryU7/biLG3E6nXC5XJibm9sUhnazZBVTdCPbQC/xq0wmo0opwTBQsYB+4Hlek7HITkjd11iL5TvNGI1GMAwDjuNableib3E7GIYZyJOysLCAUCgkiRelWCwikUgglUphfHwc0WgUbrdbMQ9Nt5KdRoaZELtcLjgcDszPz4+8oc3lcpvGVQzoRnYd3WpkafbhqM2yBhELGJWkp2ak6musVsLQIFA911aoPVHwer19uYyLxSJYlh3q3KMr4VgshlKphHA4jKmpKVXO534aTwziLm7E7XbDbrcr1m5UDjabqxjQjew6upXv0B62o3LzpQyymhhVIwu81glqGEqlkiZLd1phsVjA83zLG7SaMVngteYUneLGFEEQkEqlBhbnrtVqmJubw+zsLMxmM2ZmZjAxMaGaR6ZbyU4zg7qLGxkfH4fVah36/FeLzbaKBXQju45uRlZrmrG9MohYQLVaHVkja7FYwLLsUH2N8/m85uOxjfj9fqysrGx4Xqm+xZ3oVdSArvr6mcQSQpDP5xGPx8XmMDMzM6pPhnsp2Wlm2JUsxePxwGQyyVY7Lie6kd3kdDKyVI1m1FaxwGBiAVq4OQ9DIBBAKpUaOD7VT3tNLaCUnusg0JVcJxc+PfZeEwrr9TqWlpYQj8fBcRymp6cRDoc185v1qrLTCMuyknUu83q9MBgMA+kPqwUN0YziPbYTo3sXlYFOyT5ab6HYDanEAkYFg8EAp9M5UOtBLQi0D8IwJTNy0ykLmhCCxcVFhEKhru9TqVTE2la73Y5oNAqv16upCWG/KjsUqc852sRDyhaXcrJZ2ig2o50zUyO0OtFrtRoMBsPIZds20o9YwCgl/XTC4/FgdXW179VBuVzuSRNUawxaMqMEnST6UqlUx77AhBCsrq4iFothdXVVrG3Vasx8EJUduZiYmBATwbTOZivdoehGtgeWlpYkrdlTg37EAkap01MnGIYR3cb9MGrxWArDMHC5XMjn82ofSktaZRpXKhXwPN9yvHmex+LiIuLxOABgZmYGoVBI02VV/ZTsKEUgEADP85qVRwRe6662GSb3zehG9lXazf7r9ToEQdC8Gk03+hELqFQqI7mSa4XdbgfP823rSFsxypOMRpex1jwSDocD+XweQo0Hn65AqPFiTWwjpVIJiUQCCwsLcLlciEajGB8f19R3aYeWVrGNBINB1Go1zYYTNqurGAC0OyVUmHaN4Ec9FtsIFQvothKoVCqbym0TDAaxsLCASCTSdV86ERmFG3orGIaB1WpFqVSC0WjU1qqPANbflDF33xGAI4CRgf0tE2BnWLHdIS17CYfDIxee6bdkR2lCoRDm5+exurqK8fFxtQ9nHYVCoafrcxTRV7Kv0iqzWBCETbWq61UsQO0GBlJjMplgsVh6+u6jXLpEoeU85XIZpVKpr1W8nOT2x8D/cnnNwAIAT1A7nMLcI79FIpGA0WjEzMwMAoHAyBnYQUp21CAUCqFUKiGXy6l9KCKb2VUM6CtZkVZGdlTrYtvhcDiwsLCwad0ynZiYmEAsFoPdbu94MW+G5AuGYXD06FE899xzYjOKc845Bzt37lQtC5dwdRQOJVtvO5ZFZPepYEyjZVgbGaRkRw0YhsHk5CTm5ubAMIwmzvXN7CoGdCMrUqvV1qXcE0JQKBRk17JUkl56o2oxM1UKGIaB1+vtutool8uaX41048CBA3jmmWfEvzmOw8GDBwEAl1xySdvXUaFsQRDEB81J6PboSo4Dw7XZjxNQz3MwekfTyNKSnVE5bxiGQTgcRjKZFPtLq8lmdhUDupEVaW4+kMvlFG0krhRULKDdiqZb16tRxuVyIRaLwePxdHRHjvJvXqvVcOTIkZbbfvnLX+J1r3tdx1AAy7JtH0ajUfy/wWAQ/99LAwHC1TFnmgdpZWiNDOZWF2Dl7RgfHx+580/KZKdu16dUMAyDqakpzM7OgmEY1cqhNrurGNCNrEjzib26uoqZmRkVj0geaGegdhfVKPcs7oVQKITFxUWEw+EN26rV6sjc4Akh4DgO5XIZlUoFtVoNwJq7u10Mlud5ZLNZcZJhsVhgtVphsVhgMplkvdExJgMc508hfyCxYZvzgmm4ToyiUqkgnU6jVqvBZDLB7XbDZrNp+gYsdckO9TYp4dZnGAbT09NIJNZ+EzUMrRaTsKRGN7ItKBaLXWN3o4rD4cDq6mpHI7uZT3pamtNqMkEFILQGz/Mol8sol8vrunaZTCbYbDaxVy2VvDOZTC0NrclkQiQSEWtTS6USeJ5HoVBY1/Kw2QBLlQTnuiwKACj8IglSE8CYWTj+eAquy6JgGAY2m01MMqzVashms1heXgbLsnC73XA4HJq7JttN2AZFCpGAfmAYBpFIBIlEAgzDKK6TXSwWN1XeSyt0I9uClZUVTE9Pq30YstBNLGAzu4spwWAQs7OziEaj654vlUqqXvD1en2dMaXxcaPRCJvNBpfLBYvF0tHQ0CQnGoNt5JxzzlmXY8DzPPL5PIrFIgghMJvNcLlc4jlSrVaRzWbXxfGNRiOsVqtogPtZcTEsA/euE+C6OIJ6noPBaWqb7GQ2m8UGMPV6HblcTjQEDocDLpdL9QxkOUp2pBIJ6AdqaOPxOAKBgGLVFLSNrdYmTlKjG1msud7ozaJarcJkMmmqF6qUNIoFtDu5N/tJbzAY4HA4NmQ1KtWzmJaGUVdvo4amzWaDw+GAz+cb+BzcuXMnAODIkSPiypZmFzdiNBrh8Xjg8XgArE2wcrkcVlZWQAiBzWaD1+sVDTshBDzPo1KpoFgsIp1Or1t1mc1m0QCbzeb255fJ0FeSk8FgEI+Tqu7Mz89DEATY7Xa43W7Fm8XQkp3midqwKL2SpTQa2lAopEjIaCu4igHdyAJYu+lRN+LS0lJPjcpHGSoWsJljr93wer2IxWJwOp1gWRYcx0leG0wIWWdM6QqFukatVis8Ho/kEzqWZXHJJZfgbW97GwqFAhwOR09GyGw2ixmy9Niz2Syq1aqYHON0OsVH83et1WrrXkNhGEay+C9tHelyuUAIQblcxtLSEnieh9lsFvVU5Uaukp1eKgDkgmVZzMzMIB6PY3JyUvauZ6VSSdHqDTpRVHoRoRtZrLmjzGazGJfaTI0YWkHFAppvRputCUUnGIbBxMQElpaW4PP4sZBYgcMzWOIHNTDU1UvPI2pcqKtX6bE1mUziKrVfmmOkhBAUi0UsLy+D4ziwLAun0wmHwyG6/CwWS8sbsyAIovs5nU6vixcPE/+lMUQaR6xWq1hdXUUqlYLBYMD4+LgsuRVyluywLCsmsamBUoZWLVexbmRVQhAEmM3mTSEE0At2ux3pdHrDLHKzZxY3Y7PZ8YuHX8TxJ18AXxNgNLM445II3nLl68CyGy/Edhm9AERD4ff7R77PdStoLJTWVNbrdRQKBSwsLKBer8NoNMLlcmFsbGzDTYy6wVvF+ur1OiqVSsv4L+3U1Wv812KxiKU0NJN6ZWVFXP1Sr8WwyNmfWC13cSMsy4rJUHJp9GYymYEngIOiVhhMN7KAeJPgOG5kG8P3QzuxgHK5rHphupL8+t//Cy/94rWG6XxNwNP/XwwAcPafRHvO6N2KGAwGuN1uMabNcRzy+TwymQwIIbBYLHA6nV1LcAwGA8bGxjZkuw8b/zUajfD5fPD5fBAEAblcDslkEoQQjI2Nwe12D+RZkFtlR43Ep1YYDAYxRjs9PS35xLFUKinevKMfkRQp2fJGlnB1kCyH1eXMpuru1I1WYgHVanXLjAFfq+M3P91YswkAxx6LY/pNVjhcYz1l9OqsTT68Xq+YnV2tVpHL5UTBcLvdDqfT2fMklmEYmEwmmEymrvHfWq22TtiBrn6tVqvYQGN8fBzj4+Oi23txcRH1eh1Wq7WnBhh8rY5SroaV7CIiUfkqD9SMyTZDDW0ikZDU0NJkPDWuKd3IKggRCHL7YygcSsLJCagZs+AvmAa5zA6mhatws0HFApqzazdrVjUAsTa0VCphNVUEX2vtlqtzBIVsFYQRUCgUYDKZYDabYbFYYDabt0TZwbBYLBYx9EITlFZXV8UkKofDAafTOdBqstf478rKStv4bzAYhMFg6NoAQxAIfv3v/4Xf/DQBvibAYGKw49J625DCsGjBXdyI0WgUDW0kEpEkr0ANVzGgu4sVJ7c/tq77DMMT8W/3rhNUOirl2MxiATR2Sg0qXRkYDAbY7XZ4vV74PH48bk60NLRGM4s/OjkKo9kgrppoYlM2m13XuAGAaITpYyu7kVvRnKAkCAKKxSJSqZSYAEON7rCTvEHjv3TytLq6iqWlJdEd/tufpvD0j2KvvQdHxJDCOe96/VDH2gq1XJqdMBqNYmcoKQxtpVJRJfdFdxcrSCdFkMIvknBdHBlpRZBeaHZLaU3gu1doqUmpVEK5XBYvIpPJBLvdjkAg0PamcMYlEfGG2fy80bz2+3daNdHP53ke1WoVtVoNhUJhQ7cluoJqNMSjONZSQTOTqRu4Xq8jn89jbm4OgiCILuJWSVTD0Ev8t7EeOJmYwzOPtb5P/OanCZz9jhPE82SzYzKZMDU1hUQigZmZmYEbgVCPgRrnv25kFaSe51o3KgdAagJmX4yBHbdsKC3YbDfGxmbko5BZLAiCqJFaqVTE561WK+x2e981p2+58nUAILoCG7OLe6UxdtgKQgjq9bq4Gl5dXQXHcesudpZlReNLjfFmdts3Q8ttaGOC5iQqq9UKp9MJq9UqyzXYLv6bWy5D4FsbWb4moJSrweXfHFrTvWA2m0VDG4lEBjK0ajeg0I2sQqy1c2NbGlrGzGLqDVEILMTYTj6fX+ciZFl2UxjgRrEArRnZxvhprVYTlV5oRyS/3z/0mLMsg3Pe9Xqc/Y4TUMrVYHeZJV+ZMAwDo9EIo9HYti9soxHO5/Oo1Wrr4nIMw6xbBZvN5k1dz9yYREUIEa/BpaUl8Ryg7R/lZO18YNuGFOyuzd1+tBVmsxmTk5PiirbfyWClUkEgEJDp6Dqjx2QVpJMiiOOPp8CYDDAA6+JIjdTr9U1hgBvFAiqVyjo9XaXoJX4qt3vJaDaouiIxGAxt44jA2gqe4zhUq9W1pK3V1Q0ZqLSmtNEIa/W86weGYcRMYWDtfCmVSmKyEsuyYicqqSceRrOhbUjhhLNdyBXWFI02wzj3g8ViQSgUQjwe78vQUlexWujuYoXppAjSDWoARt0AN4oF0FphORkmfrqVoedNp7gwx3Go1Wpi6UxzcpbRaFznkh7V5Cza3pHGVVslUdFOVFK43NuFFN58xYnI53OIxWKK9frVElarFaFQSHQd9zLWamUVN6IbWQXpRxGkH0bJADeKBUiN1PFTnfY0upNbQePCNDmrVCq1jAs3J2eNwu/TnERFpfsak6hcLtfA7RU7hRTGx8fhdDqxuLgIYE2reBTGTCqsVisCgYDoOu42vmr3S9fdxSrRryLIMGjRAFssFpTzJSDHgXD1gSYaSsRPdQanMS7cTke4MS6cy+XWNXig79GcnKW21FwrjEbjuiQqGudOp9NiEtUgDUbahRQMBgPC4TDK5TLi8Tg8Hs+mLItrh81mw8TEhLiibTem1WpVdQlN3V28xVHDABOBwPBkDukjL4HhCeZM83Cc/6qIdofevWrGT3XkoZe4MDXCxWIRmUxmXVyYZug2J2epfS6YzWaxvSINV+RyObFcx263w+VyDR0rtNlsiEajSKfTiMVimJycHMqoqNHIflDsdjt8Ph9mZ2cxPT3d8ri14CoGdHexThvkMsC5/TFUDi2IfxNOEJPBXJdHW8ZPzWYzbDabHj/dYrAsuy4BqRk6AatWq6IhaxUXbnRJKz0ha6UsVCqV1ikL0aYYg6zSGYaBz+eD2+3GwsICTCYTAoFA39+Rdn3SoqegHWNjYyCEtDW01WpV9b7wurtYZyAGNcCMAODgbMv3zB2cRe4NDKwOux4/1emJXuLCPM+Lq+F2TTuak7PkPO9aJVE1KwvRphj9HAftkFQsFhGLxeD3+/sS3qAiAaNkZIG1agVCCJLJJKamptbcs1wdpeUCzAb1lal0d7GO5HQywNWlIpb4V1q+juEJ8qkc8uVi6+2vxlzpzJD28jUYDDAYDGBZVvx/44NlWTAMA5ZlxUfze+lsThobPrSLCzca4ebG/8D6ph30IaUhYllWFISnx9PYicpsNsPlcnVVFqKMjY3BbrdjaWkJmUwGoVCoJ7e01voX94PT6VwztLNJOH/Ho3AoudaPwMQie36tbShKCXQjq6Mo5nFrx4Ycp5x5WtskKJqtKggC6vU66vU6eJ6HIAjgeV78u1aridvpyU0IER+Nf4uf3WB02z3ofo3Gu9Fot/u71TYd7dCtaQeNC1erVRQKhZZNO6QUczAajfB4PGIskSaFraysgBCyrilGu89gGAaBQAAcx2F+fh42m61rIqCWlHgGweVyofLzeeQPp157siEUpWZveN3I6ihGLw052r721WxVqaEGVxAE0YDT/zf/3WzcG1/fyoA3G3KaWNKLMW9mGKOuG/fB6SUu3Cjm0K5pR6NLup/kLLPZLGqg0iSqRmUh2hSj1WrVZDJhZmYGudxabW0gEGg7mdCKpuygEK4O7snlltvU7A2vx2R1FGeYhhxyQI2bHHE4anA7Ge5Wf7eCGvdGI97JMLdzU/Vz0Tca6EEN/GY37kqKObRKoioWi+uSqGhTjEaXtsvlgsPhQCqVEl3IzS7vUVjJ0muA4zhwHCf+v16vAzkOTIfe8PU8p1jZZCO6u1hHceRqyKFF5Dbg/RjuRvd5L8fd7N5u/A6CIIiTB7o6bzyexv8PQjvD3auB15Jxb4wLt2IYMQeGWdPIpQlO9Xp9QxKVy+USk6hCoRCq1SpmZ2fhdDrXt2es1VCbm4NgtYJVsHlDR8PZBMuy4lgajUbYbLbXEtV4AXOm+ZahKGJkMJ9dhMswLom0Yb/oRlZHFZRsyLEZaYwPSw01kv2swPs14M1x7eYktUYj2ezSp6U7rYw6/f8gdDLiva7a+zXujU07pBJzCAaDMBqNG5SFLBYLnE4nIpGI6EIOTkwg/8//jPT93wapVJC2WuG94XpM7NkDZsBzSyrD2ddYdghFuS6YhjMaQT6fRzKZBCEEdrsdbrdb9r7GurtYR0dnA3Ia8EYj3Wyo6Y24eZ9eDTg1eEajsaMhb5WI1mjIm70EnWL2zQZ+EHo14lSQoXlbY71wOzEH6oLOZDKia9lutyP5D18F92//Ju5HKhWs3PtNAEBg717xeVUMZ590CkUxLAO32w232y3WKi8tLYHneZjNZrjdbllkDdVyFzNEjU8dklwuB7fbjWw2q4pyjI6OTnsajWQ/K/BeaWeoO63Ie7lhN6/S263MO23rdjttHheawIdqFcyf/TmYVwU71mG1wva9/w32VcPTbDjp/7WYVEe4el+hqGq1imw2i0qlAoPBIMawpfhexWIR1WoVXq936PfqxwbpK1kdHR1JaYx/S52F3mzomg1WuxV4r8fdvLpuNNy0WUbzPlJ8p2o8juOtDCwAVCowFgqA3S5+f9pkhkLH2mg0wmAwbPi/Wsa331CUxWIR9Wbr9TpyuRxmZ9ea5jgcDrhcrqG8OnpMVkdHR6cDSsS/W624aalY8/ZBEtharbjhdIKxWkEaFKvE19psCJ92WtskqMbVMc/zovu4XC6LdevtjlOrxhlYy7KmdcqEkHWJZFarFW63u69WjXpMVkdHR0dFGg241Ek4ncrHOI6DQAisf/qnKD/wwIbXGt/zbiSXlroed7PhtlgssNlsLQ17Yxb6KBhnhmFESUNao5zJZFCr1WA0GuF2u7vLGfIE9UwVxDWY2tig6EZWR0dHR2Z6aeDi/dRtWBqzI/3t74CUy2BsNnivv65jdnG3Bi40+7l5e690i3HTyQMtdWru8Nbq/YY1zs01yhzHIZvNYmVlBQzDwOVyrSsPIgJBbn8M+UNJgBMwZzreUW1MavTEJx0dHR0NIVQq4JeXYfT7Fa2TbUbKBi70/QCs24+uqFs1dml8UCNsMpk6GmdBEJDP55HL5cTyIMOTORQfT244HufOyMAtHvXEJx0dHZ0RhbVaYZ6eVvswNNPAha6OqfRmc3lXu/7nhBCUckWMHVpGq/WqUi0edSOro6Ojo6MoSiSw8TyP2nIJ2XrrPspKtXjUjayOjo6Ozqah0YCbA0bkOqiNGZzy69zqStw6Ojo6OpsSqjbWim5qY1Khr2R1dHR0dDYtaquN6UZWR0dHR2fTorbaWN/u4ieeeAJXXnklwuEwGIbBvn371m0nhODzn/88wuEwbDYbLrroIvz2t79dt0+1WsWePXvg9/sxNjaG3bt3i62zdHR0dHR0pGatxaNVcTnPvo1ssVjEGWecgbvvvrvl9r//+7/HHXfcgbvvvhtPPvkkQqEQLrvsMuTzeXGfvXv34pFHHsFDDz2EQ4cOoVAo4IorrtC8ULGOjo6Ojk4/DNWMgmEYPPLII3j3u98NYG0VGw6HsXfvXtx2220A1latwWAQX/nKV3DjjTcim81iYmICDzzwAK655hoAwNzcHCKRCH74wx9i165dXT9Xb0aho6Ojo6MW/dggSbOLjx8/joWFBVx++eXicxaLBRdeeCEOHz4MAHj66afBcdy6fcLhMLZt2ybu00y1WkUul1v30NHR0dHR0TqSGtmFhQUAQDAYXPd8MBgUty0sLMBsNsPj8bTdp5kvfelLosiv2+1GJBKR8rB1dHR0dHRkQZY62eYmz4SQro2fO+3z6U9/GtlsVnwkEgnJjlVHR0dHR0cuJDWyoVAIADasSFOplLi6DYVCqNVqyGQybfdpxmKxwOVyrXvo6Ojo6OhoHUmN7IknnohQKIT9+/eLz9VqNTz++OM477zzAABnnXUWTCbTun3m5+fx/PPPi/vo6Ojo6OhsBvpuRlEoFPDyyy+Lfx8/fhzHjh2D1+vFzMwM9u7di9tvvx0nnXQSTjrpJNx+++2w2+34wAc+AABwu9348Ic/jFtuuQU+nw9erxef/OQnsX37dlx66aXSfTMdHR0dHR2V6dvIPvXUU9i5c6f498033wwAuP7663H//ffjr//6r1Eul/Gxj30MmUwGb33rW/GTn/wETqdTfM2dd94Jo9GIq6++GuVyGZdccgnuv/9+WRQZdHR0dHR01EIXbdfR0dHR0ekD1epkdXR0dHR0dF5DN7I6Ojo6OjoyMZIqPNTDrXd+0tHR0dFRGmp7eom2jqSRpWIDeucnHR0dHR21yOfzcLvdHfcZycQnQRAwNzcHp9PZtZNULpdDJBJBIpHQk6T6QB+3wdHHbnD0sRsMfdwGZ5CxI4Qgn88jHA6DZTtHXUdyJcuyLKanp/t6jd4pajD0cRscfewGRx+7wdDHbXD6HbtuK1iKnviko6Ojo6MjE7qR1dHR0dHRkYlNb2QtFgv+9m//FhaLRe1DGSn0cRscfewGRx+7wdDHbXDkHruRTHzS0dHR0dEZBTb9SlZHR0dHR0ctdCOro6Ojo6MjE7qR1dHR0dHRkQndyOro6Ojo6MiEbmR1dHR0dHRkYiSN7BNPPIErr7wS4XAYDMNg375967YTQvD5z38e4XAYNpsNF110EX7729+u26darWLPnj3w+/0YGxvD7t27MTs7q+C3UIdOY8dxHG677TZs374dY2NjCIfDuO666zA3N7fuPbbi2HU75xq58cYbwTAMvva1r617fiuOG9Db2L3wwgvYvXs33G43nE4nzjnnHMTjcXG7Pnatx65QKOCmm27C9PQ0bDYbTj31VHzjG99Yt89WHLsvfelLePOb3wyn04lAIIB3v/vd+MMf/rBuH6XsxEga2WKxiDPOOAN33313y+1///d/jzvuuAN33303nnzySYRCIVx22WWisAAA7N27F4888ggeeughHDp0CIVCAVdccQXq9bpSX0MVOo1dqVTC0aNH8bnPfQ5Hjx7Fww8/jBdffBG7d+9et99WHLtu5xxl3759+NWvfoVwOLxh21YcN6D72P3nf/4nzj//fJxyyin4+c9/jt/85jf43Oc+B6vVKu6jj13rsfvEJz6BH/3oR3jwwQfxwgsv4BOf+AT27NmDH/zgB+I+W3HsHn/8cfzVX/0Vjhw5gv3794PneVx++eUoFoviPorZCTLiACCPPPKI+LcgCCQUCpEvgCj1CQAABaRJREFUf/nL4nOVSoW43W5y7733EkIIWV1dJSaTiTz00EPiPslkkrAsS370ox8pduxq0zx2rfj1r39NAJBYLEYI0ceOkPbjNjs7S6ampsjzzz9PotEoufPOO8Vt+rit0WrsrrnmGvKhD32o7Wv0sVuj1di98Y1vJF/4whfWPXfmmWeSz372s4QQfewoqVSKACCPP/44IURZOzGSK9lOHD9+HAsLC7j88svF5ywWCy688EIcPnwYAPD000+D47h1+4TDYWzbtk3cR2eNbDYLhmEwPj4OQB+7dgiCgGuvvRa33nor3vjGN27Yro9bawRBwP/9v/8Xb3jDG7Br1y4EAgG89a1vXecW1ceuPeeffz4effRRJJNJEEJw4MABvPjii9i1axcAfewo2WwWAOD1egEoayc2nZFdWFgAAASDwXXPB4NBcdvCwgLMZjM8Hk/bfXSASqWCT33qU/jABz4gqlPoY9ear3zlKzAajfj4xz/ecrs+bq1JpVIoFAr48pe/jLe//e34yU9+gve85z246qqr8PjjjwPQx64T/+t//S+cdtppmJ6ehtlsxtvf/nbcc889OP/88wHoYwesxV5vvvlmnH/++di2bRsAZe3ESErd9UKzziwhpKv2bC/7bBU4jsP73vc+CIKAe+65p+v+W3nsnn76aXz961/H0aNH+x6DrTxuwNpKFgDe9a534ROf+AQAYMeOHTh8+DDuvfdeXHjhhW1fu9XHDlgzskeOHMGjjz6KaDSKJ554Ah/72McwOTmJSy+9tO3rttLY3XTTTXj22Wdx6NChDduUsBObbiUbCoUAYMNMI5VKibOWUCiEWq2GTCbTdp+tDMdxuPrqq3H8+HHs379/ncaiPnYbOXjwIFKpFGZmZmA0GmE0GhGLxXDLLbfghBNOAKCPWzv8fj+MRiNOO+20dc+feuqpYnaxPnatKZfL+MxnPoM77rgDV155JU4//XTcdNNNuOaaa/AP//APAPSx27NnDx599FEcOHBgnQa5knZi0xnZE088EaFQCPv37xefq9VqePzxx3HeeecBAM466yyYTKZ1+8zPz+P5558X99mqUAP70ksv4bHHHoPP51u3XR+7jVx77bV49tlncezYMfERDodx66234sc//jEAfdzaYTab8eY3v3lDecWLL76IaDQKQB+7dnAcB47jwLLrb+MGg0H0EGzVsSOE4KabbsLDDz+Mn/3sZzjxxBPXbVfUTgyTsaUW+XyePPPMM+SZZ54hAMgdd9xBnnnmGTED9stf/jJxu93k4YcfJs899xx5//vfTyYnJ0kulxPf46Mf/SiZnp4mjz32GDl69Ci5+OKLyRlnnEF4nlfraylCp7HjOI7s3r2bTE9Pk2PHjpH5+XnxUa1WxffYimPX7Zxrpjm7mJCtOW6EdB+7hx9+mJhMJvJP//RP5KWXXiJ33XUXMRgM5ODBg+J76GPXeuwuvPBC8sY3vpEcOHCA/Nd//Rf51re+RaxWK7nnnnvE99iKY/ff//t/J263m/z85z9fdx8rlUriPkrZiZE0sgcOHCAANjyuv/56Qshaevbf/u3fklAoRCwWC3nb295GnnvuuXXvUS6XyU033US8Xi+x2WzkiiuuIPF4XIVvoyydxu748eMttwEgBw4cEN9jK45dt3OumVZGdiuOGyG9jd19991H/uiP/ohYrVZyxhlnkH379q17D33sWo/d/Pw8ueGGG0g4HCZWq5WcfPLJ5Ktf/SoRBEF8j604du3uY9/61rfEfZSyE7qerI6Ojo6Ojkxsupisjo6Ojo6OVtCNrI6Ojo6OjkzoRlZHR0dHR0cmdCOro6Ojo6MjE7qR1dHR0dHRkQndyOro6Ojo6MiEbmR1dHR0dHRkQjeyOjo6Ojo6MqEbWR0dHR0dHZnQjayOjo6Ojo5M6EZWR0dHR0dHJv5/CfnVFNV6x+gAAAAASUVORK5CYII=", + "text/plain": [ + "
" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "import matplotlib.pyplot as plt\n", + "from matplotlib.collections import LineCollection\n", + "\n", + "g = units[5]\n", + "pos, types = g.pos.numpy(), g.cell_type.numpy()\n", + "src, dst = g.edge_index.numpy()\n", + "keep = src < dst\n", + "fig, ax = plt.subplots(figsize=(5.5, 5.5))\n", + "ax.add_collection(LineCollection(np.stack([pos[src[keep]], pos[dst[keep]]], 1), lw=0.4, colors=\"0.8\"))\n", + "for code, name in enumerate(datamodule.input_spec.cell_types):\n", + " m = types == code\n", + " ax.scatter(pos[m, 0], pos[m, 1], s=18, label=name, zorder=2)\n", + "ax.set_aspect(\"equal\")\n", + "ax.set_title(f\"{g.unit_id}: {g.num_nodes} cells, {keep.sum()} edges\")\n", + "ax.legend(fontsize=7)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "b99713dd", + "metadata": {}, + "source": [ + "## A tile-level phenotype\n", + "\n", + "The dataset has no clinical outcome with enough samples to learn from, so the\n", + "vignettes predict a phenotype of each tile: **CD8 T-cell infiltration**, a tile\n", + "being positive when its fraction of `Tcell_CD8` cells is at or above the median\n", + "over tiles. It is derived from the data, which makes it easy to check what the\n", + "model and the niche analysis find. Tiles are grouped by image\n", + "(`patient_column=image`) so that splits and folds never separate the tiles of\n", + "one image.\n", + "\n", + "Labels are a CSV keyed by unit id; they are not part of the graphs, so they can\n", + "be changed without rebuilding them." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "954577e8", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-15T07:40:24.796454Z", + "iopub.status.busy": "2026-10-15T07:40:24.796372Z", + "iopub.status.idle": "2026-10-15T07:40:24.809077Z", + "shell.execute_reply": "2026-10-15T07:40:24.808859Z" + } + }, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
countsum
image
point161610
point23127
point8165
\n", + "
" + ], + "text/plain": [ + " count sum\n", + "image \n", + "point16 16 10\n", + "point23 12 7\n", + "point8 16 5" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "cd8 = datamodule.input_spec.cell_types.index(\"Tcell_CD8\")\n", + "rows = []\n", + "for i, entry in enumerate(units.entries):\n", + " fraction = float((units[i].cell_type == cd8).float().mean())\n", + " rows.append({\"unit_id\": entry.unit_id, \"image\": entry.region_id, \"cd8_fraction\": fraction})\n", + "labels = pd.DataFrame(rows)\n", + "labels[\"infiltrated\"] = (labels[\"cd8_fraction\"] >= labels[\"cd8_fraction\"].median()).astype(int)\n", + "labels.to_csv(DATA / \"labels.csv\", index=False)\n", + "labels.groupby(\"image\")[\"infiltrated\"].agg([\"count\", \"sum\"])" + ] + }, + { + "cell_type": "markdown", + "id": "c123d543", + "metadata": {}, + "source": [ + "## Folds\n", + "\n", + "Cross-validation folds are a JSON list of lists of ids. Region ids select every\n", + "tile of the region, so one image per fold keeps the tiles of an image together." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "20297ad8", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-15T07:40:24.810548Z", + "iopub.status.busy": "2026-10-15T07:40:24.810471Z", + "iopub.status.idle": "2026-10-15T07:40:24.812398Z", + "shell.execute_reply": "2026-10-15T07:40:24.812193Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "folds: [['point16'], ['point23'], ['point8']]\n" + ] + } + ], + "source": [ + "images = sorted(labels[\"image\"].unique())\n", + "(DATA / \"folds.json\").write_text(json.dumps([[image] for image in images]))\n", + "print(\"folds:\", [[image] for image in images])" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/vignettes/01_anndata_to_spatial_graphs.ipynb b/vignettes/01_anndata_to_spatial_graphs.ipynb deleted file mode 100644 index 54b6492..0000000 --- a/vignettes/01_anndata_to_spatial_graphs.ipynb +++ /dev/null @@ -1,589 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "7fb27b941602401d91542211134fc71a", - "metadata": {}, - "source": [ - "# 1 - From AnnData to spatial graphs\n", - "\n", - "This vignette takes a public spatial-proteomics dataset from `squidpy`,\n", - "converts it into the inputs `grass-mil` expects, and builds the cellular\n", - "graphs the models train on.\n", - "\n", - "**Dataset.** `squidpy.datasets.mibitof()` -- MIBI-TOF imaging of colorectal\n", - "carcinoma: 3 images, ~3,300 cells, 36 protein markers, 8 annotated cell types\n", - "(Hartmann et al.). It is a good fit because `grass-mil` is built for\n", - "cell-resolved spatial proteomics: every cell has coordinates, a type, and a\n", - "marker profile.\n", - "\n", - "**What you need to supply.** `grass-mil` reads a *manifest* -- one row per\n", - "image -- pointing at per-sample files. Everything else is config.\n" - ] - }, - { - "cell_type": "markdown", - "id": "acae54e37e7d407bbb7b55eff062a284", - "metadata": {}, - "source": [ - "## Setup\n", - "\n", - "The vignettes need the `io` and `interpretability` extras, plus `squidpy` to\n", - "fetch the data. `squidpy` is only used here to download a public dataset; it\n", - "is not a runtime dependency of `grass-mil`.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "id": "9a63283cbaf04dbcab1f6479b197f3a8", - "metadata": { - "execution": { - "iopub.execute_input": "2026-10-13T17:58:23.292933Z", - "iopub.status.busy": "2026-10-13T17:58:23.292845Z", - "iopub.status.idle": "2026-10-13T17:58:27.234464Z", - "shell.execute_reply": "2026-10-13T17:58:27.233614Z" - } - }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/local/home/lrabuzin/anaconda3/envs/spatial_leomed/lib/python3.11/site-packages/dask/dataframe/__init__.py:31: FutureWarning: The legacy Dask DataFrame implementation is deprecated and will be removed in a future version. Set the configuration option `dataframe.query-planning` to `True` or None to enable the new Dask Dataframe implementation and silence this warning.\n", - " warnings.warn(\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/local/home/lrabuzin/anaconda3/envs/spatial_leomed/lib/python3.11/site-packages/pyproj/__init__.py:91: UserWarning: Valid PROJ data directory not found. Either set the path using the environmental variable PROJ_DATA (PROJ 9.1+) | PROJ_LIB (PROJ<9.1) or with `pyproj.datadir.set_data_dir`.\n", - " warnings.warn(str(err))\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "working directory: vignette_work/\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/local/home/lrabuzin/anaconda3/envs/spatial_leomed/lib/python3.11/site-packages/anndata/utils.py:429: FutureWarning: Importing read_text from `anndata` is deprecated. Import anndata.io.read_text instead.\n", - " warnings.warn(msg, FutureWarning)\n" - ] - } - ], - "source": [ - "# !pip install 'grass-mil[interpretability,io]' squidpy\n", - "\n", - "import pathlib\n", - "\n", - "import numpy as np\n", - "import pandas as pd\n", - "import squidpy as sq\n", - "\n", - "WORK = pathlib.Path(\"vignette_work\").resolve()\n", - "RAW = WORK / \"data\" / \"raw\"\n", - "PROCESSED = WORK / \"data\" / \"processed\"\n", - "RAW.mkdir(parents=True, exist_ok=True)\n", - "print(\"working directory:\", WORK.name + \"/\")" - ] - }, - { - "cell_type": "markdown", - "id": "8dd0d8092fe74a7c96281538738b07e2", - "metadata": {}, - "source": [ - "## Load the public dataset\n", - "\n", - "`squidpy` caches the download, so re-running is cheap.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "id": "72eea5119410473aa328ad9291626812", - "metadata": { - "execution": { - "iopub.execute_input": "2026-10-13T17:58:27.235989Z", - "iopub.status.busy": "2026-10-13T17:58:27.235740Z", - "iopub.status.idle": "2026-10-13T17:58:27.262442Z", - "shell.execute_reply": "2026-10-13T17:58:27.262138Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "AnnData object with n_obs × n_vars = 3309 × 36\n", - " obs: 'row_num', 'point', 'cell_id', 'X1', 'center_rowcoord', 'center_colcoord', 'cell_size', 'category', 'donor', 'Cluster', 'batch', 'library_id'\n", - " var: 'mean-0', 'std-0', 'mean-1', 'std-1', 'mean-2', 'std-2'\n", - " uns: 'Cluster_colors', 'batch_colors', 'neighbors', 'spatial', 'umap'\n", - " obsm: 'X_scanorama', 'X_umap', 'spatial'\n", - " obsp: 'connectivities', 'distances'\n", - "\n", - "images : ['point16', 'point23', 'point8']\n", - "donors : ['21d7', '90de']\n", - "types : ['Endothelial', 'Epithelial', 'Fibroblast', 'Imm_other', 'Myeloid_CD11c', 'Myeloid_CD68', 'Tcell_CD4', 'Tcell_CD8']\n" - ] - } - ], - "source": [ - "# Downloaded once into the work directory (~20 MB); reused on later runs.\n", - "adata = sq.datasets.mibitof(path=WORK / \"mibitof.h5ad\")\n", - "print(adata)\n", - "print()\n", - "print(\"images :\", sorted(adata.obs[\"library_id\"].unique()))\n", - "print(\"donors :\", sorted(adata.obs[\"donor\"].unique()))\n", - "print(\"types :\", sorted(adata.obs[\"Cluster\"].unique()))" - ] - }, - { - "cell_type": "markdown", - "id": "8edb47106e1a46a883d545849b8ab81b", - "metadata": {}, - "source": [ - "## Reshape into the grass-mil input contract\n", - "\n", - "Two conventions matter:\n", - "\n", - "1. **One file per sample.** The manifest maps `sample_id` to a file, so we\n", - " split the combined AnnData by `library_id`.\n", - "2. **Name the cell-type column `cell_type`.** Several defaults key off this\n", - " name -- the sampler's `property_name`, the composition transform, and the\n", - " interpretability `cell_type_column`. You can override all of them, but\n", - " renaming once here is simpler.\n", - "\n", - "We also densify `X`: the h5ad loader reads the matrix directly.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "id": "10185d26023b46108eb7d9f57d49d2b3", - "metadata": { - "execution": { - "iopub.execute_input": "2026-10-13T17:58:27.264067Z", - "iopub.status.busy": "2026-10-13T17:58:27.263979Z", - "iopub.status.idle": "2026-10-13T17:58:27.392847Z", - "shell.execute_reply": "2026-10-13T17:58:27.392394Z" - } - }, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
sample_idinput_pathinput_typeregion_iddonordonor_label
0point16point16.h5adh5adpoint1690de0
1point23point23.h5adh5adpoint2321d71
2point8point8.h5adh5adpoint890de0
\n", - "
" - ], - "text/plain": [ - " sample_id input_path input_type region_id donor donor_label\n", - "0 point16 point16.h5ad h5ad point16 90de 0\n", - "1 point23 point23.h5ad h5ad point23 21d7 1\n", - "2 point8 point8.h5ad h5ad point8 90de 0" - ] - }, - "execution_count": 3, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "rows = []\n", - "for lib in sorted(adata.obs[\"library_id\"].unique()):\n", - " sub = adata[adata.obs[\"library_id\"] == lib].copy()\n", - " sub.X = np.asarray(sub.X.todense(), dtype=np.float32)\n", - " sub.obs[\"cell_type\"] = sub.obs[\"Cluster\"].astype(str)\n", - " sub.obs = sub.obs[[\"cell_type\", \"cell_size\", \"donor\"]]\n", - " path = RAW / f\"{lib}.h5ad\"\n", - " sub.write_h5ad(path)\n", - " rows.append(\n", - " {\n", - " \"sample_id\": lib,\n", - " \"input_path\": path.name, # resolved relative to the manifest\n", - " \"input_type\": \"h5ad\",\n", - " \"region_id\": lib,\n", - " \"donor\": str(sub.obs[\"donor\"].iloc[0]),\n", - " }\n", - " )\n", - "\n", - "manifest = pd.DataFrame(rows)\n", - "manifest[[\"sample_id\", \"input_path\", \"input_type\", \"region_id\"]].to_csv(\n", - " RAW / \"manifest.csv\", index=False\n", - ")\n", - "\n", - "# A bag-level label. See the caveat in vignette 2 before reading anything\n", - "# into supervised numbers on a three-image dataset.\n", - "manifest[\"donor_label\"] = (manifest[\"donor\"] == \"21d7\").astype(int)\n", - "manifest[[\"region_id\", \"donor_label\"]].to_csv(RAW / \"labels.csv\", index=False)\n", - "\n", - "manifest" - ] - }, - { - "cell_type": "markdown", - "id": "8763a12b2bbd4a93a75aff182afb95dc", - "metadata": {}, - "source": [ - "## Build the graphs\n", - "\n", - "`prepare_data()` runs once and writes `.pt` graph files plus an index. The\n", - "settings worth understanding:\n", - "\n", - "- **`coord_scale_um`** converts stored coordinates to micrometres. MIBI-TOF\n", - " FOVs here are ~400 um across ~1024 px, so ~0.39 um/px. Getting this wrong\n", - " silently rescales every distance downstream -- graph edges, ego-graph radii,\n", - " and the filtration thresholds in the interpretability report.\n", - "- **`sample_unit=tile`** cuts each image into square patches. A patch is one\n", - " graph unit, and later one MIL bag.\n", - "- **`tiling.tile_size_um=100`** gives ~16 patches per image at this scale.\n", - "- **`min_cells`** drops patches too sparse to form a useful graph.\n", - "\n", - "Run it through the CLI so the exact command is reusable outside a notebook.\n", - "\n", - "The graph is a Delaunay triangulation of the cell centroids; every edge carries\n", - "the bounded log edge length `clip((ln d - 2)/3, 0, 1)` that the GRASS-MIL encoder\n", - "conditions on (`log_distance`) and a `neighbor` flag for edges shorter than\n", - "`neighbor_cutoff_um`. `+trainer.fast_dev_run=true` only runs the precompute.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "7623eae2785240b9bd12b16a66d81610", - "metadata": { - "execution": { - "iopub.execute_input": "2026-10-13T17:58:27.394571Z", - "iopub.status.busy": "2026-10-13T17:58:27.394474Z", - "iopub.status.idle": "2026-10-13T17:58:31.656297Z", - "shell.execute_reply": "2026-10-13T17:58:31.655625Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "exit code: 0\n" - ] - } - ], - "source": [ - "import subprocess\n", - "import os\n", - "import sys\n", - "\n", - "cmd = [\n", - " sys.executable,\n", - " \"-m\",\n", - " \"grass_mil.train\",\n", - " \"task=finetune_mean\",\n", - " \"model=supervised_module\",\n", - " \"data.graph_builder.kwargs.edge_features=[log_distance,neighbor]\",\n", - " \"data.graph_builder.kwargs.neighbor_cutoff_um=20\",\n", - " \"data=spatial_omics\",\n", - " \"trainer=cpu\",\n", - " f\"data.raw_manifest_path={RAW}/manifest.csv\",\n", - " f\"data.processed_dir={PROCESSED}\",\n", - " \"data.h5ad.coord_source=obsm\",\n", - " \"data.h5ad.coord_key=spatial\",\n", - " \"data.h5ad.categorical_label_columns=[cell_type]\",\n", - " \"data.coord_scale_um=0.39\",\n", - " \"data.sample_unit=tile\",\n", - " \"data.tiling.tile_size_um=100\",\n", - " \"data.tiling.stride_um=100\",\n", - " \"data.min_cells=20\",\n", - " \"data.tiling.min_cells=20\",\n", - " f\"data.graph_labels.label_file={RAW}/labels.csv\",\n", - " \"data.graph_labels.id_column=region_id\",\n", - " \"data.graph_labels.tasks=[donor_label]\",\n", - " \"data.graph_labels.scope=region\",\n", - " \"data.split.train_val_test_split=[0.34,0.33,0.33]\",\n", - " \"data.split.split_by=sample\",\n", - " \"data.force_precompute=true\",\n", - " f\"paths.log_dir={WORK}/logs\",\n", - " \"+trainer.fast_dev_run=true\", # precompute only; no real training here\n", - "]\n", - "env = dict(os.environ, PROJECT_ROOT=str(WORK))\n", - "out = subprocess.run(cmd, capture_output=True, text=True, env=env)\n", - "print(\"exit code:\", out.returncode)\n", - "if out.returncode:\n", - " print(out.stdout[-3000:], out.stderr[-3000:])" - ] - }, - { - "cell_type": "markdown", - "id": "7cdc8c89c7104fffa095e18ddfef8986", - "metadata": {}, - "source": [ - "## Inspect what was built\n" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "id": "b118ea5561624da68c537baed56e602f", - "metadata": { - "execution": { - "iopub.execute_input": "2026-10-13T17:58:31.657557Z", - "iopub.status.busy": "2026-10-13T17:58:31.657469Z", - "iopub.status.idle": "2026-10-13T17:58:31.660598Z", - "shell.execute_reply": "2026-10-13T17:58:31.660146Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "graph units (patches): 44\n", - "per sample : {'point16': 16, 'point23': 12, 'point8': 16}\n", - "per split : {'test': 16, 'val': 12, 'train': 16}\n", - "\n", - "cell types : {'Imm_other': 0, 'Tcell_CD4': 1, 'Endothelial': 2, 'Tcell_CD8': 3, 'Fibroblast': 4, 'Myeloid_CD11c': 5, 'Myeloid_CD68': 6, 'Epithelial': 7}\n", - "markers : 36 ['ASCT2', 'ATP5A', 'CD11c', 'CD14', 'CD3', 'CD31'] ...\n" - ] - } - ], - "source": [ - "import json\n", - "import collections\n", - "\n", - "index = json.loads((PROCESSED / \"processed_index.json\").read_text())\n", - "entries = index[\"entries\"] if \"entries\" in index else index\n", - "meta = json.loads((PROCESSED / \"metadata.json\").read_text())\n", - "\n", - "print(\"graph units (patches):\", len(entries))\n", - "print(\"per sample :\", dict(collections.Counter(e[\"sample_id\"] for e in entries)))\n", - "print(\"per split :\", dict(collections.Counter(e[\"split\"] for e in entries)))\n", - "print()\n", - "print(\"cell types :\", meta[\"label_maps\"][\"cell_type\"])\n", - "print(\n", - " \"markers :\",\n", - " len(meta[\"molecular_feature_names\"]),\n", - " meta[\"molecular_feature_names\"][:6],\n", - " \"...\",\n", - ")" - ] - }, - { - "cell_type": "markdown", - "id": "938c804e27f84196a10c8828c723f798", - "metadata": {}, - "source": [ - "Note the split: whole images go to train/val/test. Splitting by `sample`\n", - "rather than by `patch` is deliberate -- patches from one image share a donor\n", - "and a batch, so a patch-level split would leak that identity across folds.\n" - ] - }, - { - "cell_type": "markdown", - "id": "504fb2a444614c0babb325280ed9130a", - "metadata": {}, - "source": [ - "## Look at one graph unit\n", - "\n", - "Each unit is a `torch_geometric.data.Data`: node features `x` (markers),\n", - "`edge_index` (Delaunay adjacency), `pos` (micrometre coordinates), and\n", - "`categorical_codes` (the cell-type index the encoder embeds).\n" - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "id": "59bbdb311c014d738909a11f9e486628", - "metadata": { - "execution": { - "iopub.execute_input": "2026-10-13T17:58:31.661948Z", - "iopub.status.busy": "2026-10-13T17:58:31.661868Z", - "iopub.status.idle": "2026-10-13T17:58:32.345987Z", - "shell.execute_reply": "2026-10-13T17:58:32.345463Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Data(\n", - " x=[59, 36],\n", - " edge_index=[2, 322],\n", - " pos=[59, 2],\n", - " edge_attr=[322, 2],\n", - " edge_attr_names=[2],\n", - " molecular_feature_names=[36],\n", - " categorical_codes=[59, 1],\n", - " categorical_labels=[1],\n", - " categorical_slices={ cell_type=0 },\n", - " graph_y=[1, 1],\n", - " graph_label_names=[1],\n", - " sample_id='point16',\n", - " region_id='point16',\n", - " patch_id='point16_point16_patch_0'\n", - ")\n", - "\n", - "nodes: 59 | edges: 322\n", - "x : (59, 36)\n", - "pos : extent [0.38999998569488525, 3.9000000953674316] to [98.66999816894531, 100.2300033569336] um\n" - ] - } - ], - "source": [ - "import torch\n", - "\n", - "sample_entry = entries[0]\n", - "g = torch.load(sample_entry[\"path\"], weights_only=False)\n", - "print(g)\n", - "print()\n", - "print(\"nodes:\", g.num_nodes, \"| edges:\", g.edge_index.shape[1])\n", - "print(\"x :\", tuple(g.x.shape))\n", - "print(\"pos : extent\", g.pos.min(0).values.tolist(), \"to\", g.pos.max(0).values.tolist(), \"um\")" - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "id": "b43b363d81ae4b689946ece5c682cd59", - "metadata": { - "execution": { - "iopub.execute_input": "2026-10-13T17:58:32.347416Z", - "iopub.status.busy": "2026-10-13T17:58:32.347329Z", - "iopub.status.idle": "2026-10-13T17:58:32.488605Z", - "shell.execute_reply": "2026-10-13T17:58:32.487617Z" - } - }, - "outputs": [ - { - "data": { - "image/png": "iVBORw0KGgoAAAANSUhEUgAAAk4AAAJLCAYAAAARlBlGAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOydd3wb5f3HP6e9rWXZlmU5gQbIggTCLjMklA1poAUCSSm/BggjjFIKpQTKaGhLKQEaVklICKMUCLSFNmwoe4SyCm1JbMtL09Zed8/vD3NXy5ZlbZ3k5/16+ZVIOt09d7p77nPfyRBCCCgUCoVCoVAokyKp9QAoFAqFQqFQ6gUqnCgUCoVCoVDyhAonCoVCoVAolDyhwolCoVAoFAolT6hwolAoFAqFQskTKpwoFAqFQqFQ8oQKJwqFQqFQKJQ8ocKJQqFQKBQKJU+ocKJQKBQKhULJEyqcKGXlr3/9K9asWVP096dNm4bjjz++fAMqkVdeeQUMw+CVV14p+LvRaBRr1qzJ+t1QKIQrr7wSixcvRnNzMxiGyXncUqkUbrvtNsydOxdqtRpGoxEHHXQQ3nzzzYLHVW4mG3su7r77bmzYsCHrZ7fffjuWLFmC6dOng2EYHH744TnXtXXrVhx22GEwGAzQarWYPXs27r333qLGVU62bNmC22+/vejvMwyDCy+8sGzjefTRRzFv3jyoVCrY7XasXr0a4XC4bOuvBBs2bADDMNi5c6fw3ooVKzBt2rSyrW+qsXPnTjAMg1//+te1HkrdQYUTpaz89a9/xfXXX1/rYZSNvffeG2+99Rb23nvvgr8bjUZx/fXXZxVOPp8P9957LxKJBE4++eSc62FZFqeccgpuuOEGnH766Xjuuefw8MMP4zvf+Q4ikUjB4yo3b731Fs4999yivptLOK1fvx5dXV048sgj0dzcnHM9v/zlL7FkyRLMmTMHjz/+OJ555hlccMEFSCaTRY2rnJQqnMrJww8/jNNPPx377rsvnnvuOVx33XXYsGEDlixZUuuhUSh1g6zWA6BQxIzBYMABBxxQ9vV2dnYiEAiAYRh4vV7cf//9Ey67bt06PPfcc/jHP/6RMZbjjjuu7OMqhkocHwD4/PPPIZGMPNvNmTNnwuU++OADXHPNNbjllltw5ZVXCu8vXLiwIuOqV1iWxY9//GMsXrwY9913HwDgiCOOgF6vx5lnnonnnnsOxxxzTI1HWZ9Eo1FoNJqKrDsWi0GtVldk3ZTioBanKcyaNWvAMAw++ugjLFmyBAaDAU1NTVi2bBk8Hk/Gso899hgWL16MtrY2qNVqzJw5E1dddVWGxWPFihW46667AIy4F/g/3hzOcRzWrVuHefPmCe6mAw44AM8888y4sT3//PPYe++9oVarsccee+APf/jDpPvDm55vvfVW3HTTTXA6nVCpVFiwYAFefPHFccu/8cYbWLhwIfR6PTQaDQ466CD85S9/yVgmm6tuxYoV0Ol0+M9//oNjjz0WOp0OHR0duPzyy5FIJISx8FaS66+/XjgWK1asyDg++fC73/0Ohx56aFkECu8Kfeqpp7DnnntCpVJhl112wR133DFu2e7ubixbtgw2mw1KpRIzZ87Eb37zG3Acl7HcWFcd7wZ5+eWXcf7558NqtcJisWDJkiXo6+vLGMtnn32GV199VTgeo10vvGiajDvvvBNKpRIXXXRRYQcjC/xv+9lnn2HhwoXQarVobm7GhRdeiGg0mrHsXXfdhUMPPRQ2mw1arRZz587FrbfeilQqJSxz+OGH4y9/+Qu6uroyrgmeRCKBG264ATNnzoRKpYLFYsERRxyR1QW7adMmzJw5ExqNBnvttRf+/Oc/F7Rvb7/9Nvr7+/GDH/wg4/1TTz0VOp0OTz31VEHry4ctW7bgwAMPhE6ng06nw7x58/DAAw9kLPPCCy9g4cKFMBgM0Gg0OPjgg7Ner5PBX//ZLJj5uJO3bduGk046CQ6HAyqVCt/61rewcuVKeL3ejOX4efPDDz/E0qVLYTKZsOuuu+Zc9xtvvIEDDzwQKpUK7e3tuPbaa3H//fePcxfy1+eTTz6J+fPnQ6VSCRb8fM43YOScmzNnDl5//XUccMABUKvVwjZZls06vttuuw3Tp0+HTqfDgQceiLfffjvn/kx1qHCi4JRTTsG3vvUtPPHEE1izZg2efvppHH300RkX5L///W8ce+yxeOCBB/D8889j9erVePzxx3HCCScIy1x77bVYunQpgBH3Df/X1tYGYOSmdMkll2DffffFY489hkcffRQnnnjiuDiDjz/+GJdffjkuvfRSbN26FXvuuSd++MMf4rXXXstrf+688048//zzuP3227F582ZIJBIcc8wxeOutt4RlXn31VRx55JEYHh7GAw88gEceeQR6vR4nnHACHnvssUm3kUqlcOKJJ2LhwoXYunUrzjnnHPz2t7/F2rVrAQBtbW14/vnnAQA//OEPhWNx7bXX5rUPPD09Pdi5cyfmzp2Lq6++Gi0tLZDJZJg9ezY2btxY0Lp4tm/fjtWrV+PSSy/FU089hYMOOgiXXHJJRqyDx+PBQQcdhL///e/4xS9+gWeeeQZHHXUUrrjiirzjbc4991zI5XJs2bIFt956K1555RUsW7ZM+Pypp57CLrvsgvnz5wvHp5ib92uvvYaZM2fiT3/6E3bffXdIpVI4HA5cddVVRbnqUqkUjj32WCxcuBBPP/00LrzwQtxzzz343ve+l7Hcf//7X5xxxhnYtGkT/vznP+OHP/whfvWrX2HlypXCMnfffTcOPvhgtLa2ZlwTAJBOp3HMMcfgF7/4hSBmN2zYgIMOOgjd3d0Z2/rLX/6CO++8EzfccAP+9Kc/wWw245RTTsHXX3+d9359+umnAIA999wz4325XI499thD+Lxc/PznP8eZZ54Ju92ODRs24KmnnsLy5cvR1dUlLLN582YsXrwYBoMBGzduxOOPPw6z2Yyjjz66KPFUCv/9739x4IEH4ve//z3+/ve/4+c//zneeecdfPvb3x4nTgBgyZIl+Na3voU//vGPWL9+/YTr/ec//4lFixYhGo1i48aNWL9+PT788EPcdNNNWZf/8MMP8eMf/xgXX3wxnn/+eXz3u98VxjfZ+cYzMDCA73//+zjzzDOxdetWLF26FDfeeCMuueSSccvedddd2LZtG26//XY8/PDDiEQiOPbYYzE8PJzvoZt6EMqU5brrriMAyKWXXprx/sMPP0wAkM2bN2f9HsdxJJVKkVdffZUAIB9//LHw2apVq0i20+q1114jAMg111yTc0ydnZ1EpVKRrq4u4b1YLEbMZjNZuXJlzu/u2LGDACB2u53EYjHh/WAwSMxmMznqqKOE9w444ABis9lIKBQS3kun02TOnDnE4XAQjuMIIYS8/PLLBAB5+eWXheWWL19OAJDHH388Y/vHHnss2X333YXXHo+HACDXXXddznHnWu6tt94iAIjBYCCzZs0ijz/+OPnb3/5Gli5dSgCQe++9N+e6x9LZ2UkYhiHbt2/PeH/RokXEYDCQSCRCCCHkqquuIgDIO++8k7Hc+eefTxiGIV9++aXw3tixP/jggwQAueCCCzK+e+uttxIApL+/X3hv9uzZ5LDDDpt03LmWUyqVRK/XE5PJRO68807y0ksvkWuuuYZIpVJyxhlnTLru0fC/7e9+97uM92+66SYCgLzxxhtZv8eyLEmlUuShhx4iUqmU+P1+4bPjjjuOdHZ2jvvOQw89RACQ++67L+eYAJCWlhYSDAaF9wYGBohEIiG33HJL3vvG78Po48+zePFisttuu+W9rsn4+uuviVQqJWeeeeaEy0QiEWI2m8kJJ5yQ8T7LsmSvvfYi++23n/Aef07t2LFDeG/58uUZx5W//h988MFx25roHB29vtHwc1xXVxcBQLZu3Sp8xs+bP//5zyfct9GceuqpRKvVEo/Hk7GPs2bNGjeGzs5OIpVKM66vbOQ63w477LBxYyaEkP/7v/8jEolEmFv54zV37lySTqeF5d59910CgDzyyCN57d9UhFqcKDjzzDMzXp922mmQyWR4+eWXhfe+/vprnHHGGWhtbYVUKoVcLsdhhx0GAPjiiy8m3cZzzz0HAFi1atWky86bNw9Op1N4rVKpsNtuu2U8qeZiyZIlUKlUwmvekvTaa6+BZVlEIhG88847WLp0KXQ6nbCcVCrFWWedBZfLhS+//DLnNhiGybC2ASNP8vmOMV94t1g8Hsdf//pXnHrqqVi8eDEef/xx7L333rjhhhsKXufs2bOx1157Zbx3xhlnIBgM4sMPPwQAvPTSS5g1axb222+/jOVWrFgBQgheeumlSbdz4oknZrzmLR2VOEahUAh33303Vq1ahSOOOAI33ngjLrroImzZsgX/+c9/Cl7n2GvijDPOAICMa+Kjjz7CiSeeCIvFIlwTZ599NliWxVdffTXpNp577jmoVCqcc845ky7LxyLxtLS0wGazFXUsJ3IRT+Y6ZlkW6XRa+Bvrsh3Ntm3bwLJszuv9zTffhN/vx/Lly8et9zvf+Q7ee++9qiY/uN1unHfeeejo6IBMJoNcLkdnZyeA7HMcbwmaDN66bbVahfckEglOO+20rMvvueee2G233ca9X8j5ptfrx11/Z5xxBjiOG2e5P+644yCVSjO2D5T/Om0kqHCioLW1NeO1TCaDxWKBz+cDAITDYRxyyCF45513cOONN+KVV17Be++9hyeffBLASPDiZHg8Hkil0nHbyobFYhn3nlKpzGs7wPj94d9LJpMIh8MIBAIghAguxNHY7XYAEPZ9IjQaTYY448cYj8fzGmO+8Mdijz32ECZxYOQmd/TRR8PlcsHtdhe0zomOD/C//fb5fCUdn9Fj51EqlQDyO18Kgd/O0UcfnfE+H+jMi8F84c//0Yw9Pt3d3TjkkEPQ29uL3/3ud3j99dfx3nvvCTF++V4Tdrs9r1iuUq+J0evI9tv5/X6Yzeac3991110hl8uFv1yinY+RdDgcEy4zODgIAFi6dGnGeuVyOdauXQtCCPx+/6T7VQ44jsPixYvx5JNP4sorr8SLL76Id999V4j1yXacs10f2fD5fGhpaRn3frb3JlpvoedbtnWPPYd5qnWdNhI0q46CgYEBtLe3C6/T6TR8Pp9wQb300kvo6+vDK6+8IliZAGBoaCjvbTQ3N4NlWQwMDOQ94RTLwMBA1vcUCgV0Oh1kMhkkEgn6+/vHLccHL49+Oqwlu+6664TZOoQQAPkHUfNMdHyA/02iFoulLo4PMPKEnG2fij0+Y89/YPzxefrppxGJRPDkk09mCNrt27fnvZ3m5ma88cYb4Diu4DEWw9y5cwEAn3zyCWbNmiW8n06n8a9//Qunn356zu8/++yzQvID8D8RnQ0+McLlcqGjoyPrMvw5tG7dugkTHyYSF9ngH2RGjxHIT+R/+umn+Pjjj7FhwwYsX75ceD+XtTLf5A6LxSKIxNFkO2cnWm+h51uu7WUT4ZTCoBYnCh5++OGM148//jjS6bRQcJC/kPknEZ577rln3Lomelrhn/5///vfl2XMuXjyySczLD+hUAjPPvssDjnkEEilUmi1Wuy///548sknM8bJcRw2b94Mh8OR1VReKOV4cpPJZDjppJPwxRdfZATRE0Lw/PPPY9dddy1YxHz22Wf4+OOPM97bsmUL9Hq9UK9q4cKF+Pzzz8dZax566CEwDIMjjjiiuB0aQ6FWk2zwLhPeHczz17/+FRKJBPvuu2/B6xx7TWzZsgUAcl4ThBAhzX80E+3jMcccg3g8PmEdq3Kz//77o62tbdz2nnjiCYTD4UlrOc2dOxcLFiwQ/nIJp8WLF0Mqlea83g8++GAYjUZ8/vnnGesd/adQKPLev5aWFqhUKvzzn//MeH/r1q2TfreQOa5QDjvsMLz00ksZ2Xkcx+GPf/xj3uso5HwDRua8sdnKW7ZsgUQiwaGHHlrI8ClZoBYnCp588knIZDIsWrQIn332Ga699lrstddegg/+oIMOgslkwnnnnYfrrrsOcrkcDz/88LibL/C/p9q1a9fimGOOgVQqxZ577olDDjkEZ511Fm688UYMDg7i+OOPh1KpxEcffQSNRlOWVHIeqVSKRYsW4bLLLgPHcVi7di2CwWBGYc5bbrkFixYtwhFHHIErrrgCCoUCd999Nz799FM88sgjeT9N5kKv16OzsxNbt27FwoULYTabYbVahZT75557DpFIBKFQCMBI3aInnngCAHDssccKlqZf/OIXeO655/Cd73wHa9asgcFgwP3334+PP/4Yjz/+eMHjstvtOPHEE7FmzRq0tbVh8+bN2LZtG9auXSts89JLL8VDDz2E4447DjfccAM6Ozvxl7/8BXfffTfOP//8sghLYOR8efTRR/HYY49hl112gUqlEs6h999/XxCLwWAQhBDh+Oy7777Ck/cPfvAD3HPPPbjgggvg9Xoxa9YsvPDCC7jrrrtwwQUXZDyh54NCocBvfvMbhMNh7LvvvnjzzTdx44034phjjsG3v/1tAMCiRYugUChw+umn48orr0Q8Hsfvf/97BAKBrPv45JNP4ve//z322WcfSCQSLFiwAKeffjoefPBBnHfeefjyyy9xxBFHgOM4vPPOO5g5cya+//3vF3tYsyKVSnHrrbfirLPOwsqVK3H66afj3//+N6688kosWrQI3/nOd8q2rWnTpuHqq6/GL37xC8RiMZx++uloamrC559/Dq/Xi+uvvx46nQ7r1q3D8uXL4ff7sXTpUthsNng8Hnz88cfweDwFPWgxDINly5bhD3/4A3bddVfstddeePfddwXRm4s99tgDu+66K6666ioQQmA2m/Hss89i27ZtpRwGAMA111yDZ599FgsXLsQ111wDtVqN9evXC/Fb+VgbCznfgBGr0vnnn4/u7m7stttu+Otf/4r77rsP559/fkb8KKVIahWVTqk9fHbIBx98QE444QSi0+mIXq8np59+OhkcHMxY9s033yQHHngg0Wg0pLm5mZx77rnkww8/HJfFkkgkyLnnnkuam5sJwzAZWSMsy5Lf/va3ZM6cOUShUJCmpiZy4IEHkmeffVb4fmdnJznuuOPGjfWwww6bNPuKzxJZu3Ytuf7664nD4SAKhYLMnz+f/O1vfxu3/Ouvv06OPPJIotVqiVqtJgcccEDGWAiZOKtOq9VOeDxH88ILL5D58+cTpVJJAJDly5dn7CuArH9js30++eQTctxxxxG9Xk9UKlXWseYDf3yfeOIJMnv2bKJQKMi0adPIbbfdNm7Zrq4ucsYZZxCLxULkcjnZfffdya9+9SvCsmzGcpggY+m9997LWC7bsdy5cydZvHgx0ev1BEBGlhSf4Zbtb2zmlM/nIytXriQtLS1ELpeT3XbbLetYJ4P/bf/5z3+Sww8/nKjVamI2m8n5559PwuFwxrLPPvss2WuvvYhKpSLt7e3kxz/+MXnuuefG7aPf7ydLly4lRqNRuCZ4YrEY+fnPf05mzJhBFAoFsVgs5MgjjyRvvvlmxvFdtWrVuLF2dnZmnE/5smXLFrLnnnsShUJBWltbycUXX5yRXVpOHnroIbLvvvsSlUpFdDodmT9//rjf7tVXXyXHHXccMZvNRC6Xk/b2dnLccceRP/7xj8Iy+WTVEULI8PAwOffcc0lLSwvRarXkhBNOIDt37swrq+7zzz8nixYtEjI0Tz31VNLd3T3uu/x1PjpLbjJef/11sv/++xOlUklaW1vJj3/8Y7J27VoCgAwNDQnLTTT/EZL/+XbYYYeR2bNnk1deeYUsWLCAKJVK0tbWRq6++mqSSqWE5fj58le/+tW4bY3dZ0omDCHfBAJQphxr1qzB9ddfD4/HI6qYlWLZuXMnpk+fjl/96le44ooraj0cUTJt2jTMmTOn4OKJU4UVK1YIrisKpZIsXrwYO3fuzCsDsxAOP/xweL3estflovwP6qqjUCgUCqWCXHbZZZg/fz46Ojrg9/vx8MMPY9u2beOqqFPqAyqcKJQGgGVZ5DIeMwyTUatlqsFxXM66Q8BIIH69kk6nc34ukUiqkrlHyQ7Lsvj5z3+OgYEBMAyDWbNmYdOmTRmV9Cn1A3XVUSgNwLRp03IWrDvssMMy+u1NNVasWDFpi5p6nQp5F3Uurrvuukl7tVEolPygwolCaQA++eSTcfVrRqPX67H77rtXcUTiYufOneOatY5lwYIFVRpNeUkmk+NS8Mdit9tzlg+gUCj5Q4UThUKhUCgUSp5QpzeFQqFQKBRKntRvNGQZ4TgOfX190Ov1ZSl8SKFQKBQKpX4ghCAUCuXVP5IKJ4z035qonxKFQqFQKJSpQU9PT87m1AAVTgBGAmeBkQNmMBhqPBoKhUKhUCjVJBgMoqOjQ9ADuaDCCf9roGgwGKhwolAoFAplipJPuA4NDqdQKBQKhULJE2pxolAoFAqlRFiWRSqVqvUwKHkgl8tL6qRAhROFQqFQKCUQDofhcrnqtvr8VINhGDgcDuh0uqK+T4UThUKhUChFwrIsXC4XNBoNmpubaUkbkUMIgcfjgcvlwowZM4qyPFHhRKFQKBRKkaRSKRBC0NzcDLVaXevhUPKgubkZO3fuRCqVKko40eBwCoVCoVBKhFqa6odSfysqnCgUCoVCoVDyhAonCoVCoVDqGKvVWtPtv/LKK3j33XeF1ytWrMCf//znGo6ostAYJwqFQqFQqkSK5fCnD1zo8kfRadbgu/s4IJfWtw3jlVdegdVqxX777VfyuliWLalUQDWo71+LQqFQKJQ6IcVyWHb/O7jqyU9w32tf46onP8Gy+99BiuXKsv5XXnkFRx11FL773e/iW9/6Fm6++WasX78ee++9N/bbbz94vV4AwOGHH44rrrgC3/72t7Hnnnviww8/xHHHHYdvfetbuPPOOydcfywWw1lnnYW5c+div/32w/bt29HT04P169fjl7/8JebNm4dPPvkEALBt2zYccMABmDFjBl599VUAQDqdxiWXXIL99tsP8+bNw9NPPw0A2LBhA04//XQcd9xxOP3008tyLCoJtThRKBQKhVIF/vSBC+/s8AMA0txIzad3dvjx5IcufG9fZ1m2sX37dnzxxRfQaDSYPn06rr76anz44Ye46qqrsGnTJlx66aUAAK1WizfeeAM33XQTvve97+G9994Dx3GYPXs2Lrzwwqzrvuuuu6DX6/HJJ5/g7bffxvLly/Hxxx/jvPPOg9Vqzfje0NAQ3n77bbz00ku44YYb8OKLL+L+++/HtGnT8Lvf/Q7BYBD7778/jjnmGADAu+++i48++qgu2p5R4UShUCgUShXo8kchkzCCaAIAmYRBly9atm0ceOCBaG5uBgA4HA5BmOy555545513hOVOPPFEAMDcuXOxYMECGI1GACNN7wOBAEwm07h1v/HGG7jyyisBAAcccABisRiGh4ezjoNf/z777IOdO3cCGLFCffbZZ9i4cSMAIBKJoLe3FwBw9NFH14VoAqhwolAoFAqlKnSaNRmiCRixPHVaNGXbhlKpFP4vkUiE1xKJBCzLjltu9DLZlssFIWTC1H5+nVKpVFgfIQT33nsvDj300IxlX3vtNWg05TsGlYbGOFEoFAqFUgW+u48D+083AxixNAHA/tPNWLK3o5bDyptvf/vb2LJlC4AR15pWq4XBYIBer0coFJr0+0cddRTWr18vCKnt27dXcrgVg1qcKBQKhUKpAnKpBJvP3R9PfuhCly+KTosGS/aun6y6VatW4f/+7/+w5557QqVS4cEHHwQAnHDCCVi6dCkee+wxbNq0acLvr1y5El9//TXmzZsHQgh22203PPnkk9UaftlgCO1KiGAwiKamJgwPD9eNj5VCoVAotScej2PHjh2YPn06VCpVrYdDyYNsv1khOqA+ZC6FQqFQKGIkFQe4NMCVp6QARfzUVDi99tprOOGEE2C328EwjFDTgYcQgjVr1sBut0OtVuPwww/HZ599lrFMIpHARRddBKvVCq1WixNPPBEul6uKe0GhUCiUKQfHAi/+AnhgERDsA7xfjfzbAE6cTz75BPPmzcv4W7p0aa2HJRpqKpwikQj22muvCQtu3Xrrrbjttttw55134r333kNraysWLVqUEYS2evVqPPXUU3j00UfxxhtvIBwO4/jjj887K4BCoVAolIJ5+Wbg9d8A6dg3b3BAeBAI9dd0WOVg7ty52L59e8bfE088UethiYaaBocfc8wxQo2JsRBCcPvtt+Oaa67BkiVLAAAbN25ES0sLtmzZgpUrV2J4eBgPPPAANm3ahKOOOgoAsHnzZnR0dOCFF17A0UcfXbV9oVAoFEpjQghBOp0W/thEBE1v3QUGWaxLEQ+gawUkNBKmURFtVt2OHTswMDCAxYsXC+8plUocdthhePPNN7Fy5Up88MEHSKVSGcvY7XbMmTMHb7755oTCKZFIIJFICK+DwWDldoRCoVAoooAQApZlwbJsphD65vVEngqGYSCTySCVSiGTySCLecEIlqaxG+EALgVIlNk/p9Q9ohVOAwMDAICWlpaM91taWtDV1SUso1AoxlU4bWlpEb6fjVtuuQXXX399mUdMoVAolGrBcVyG6BkrhCZKGBfEzzdCSKVSCa8lEsmEBR15IpEIvAkZdDJ1VvFEIAEkMuReC6WeEa1w4hl7EueqVJrvMj/96U9x2WWXCa+DwSA6OjpKGyiFQqFQCoa3Ao21/vB/3ATZarwViP+Ty+UZImiy+0ShJJNJDAwMQKVSoWOX3cAcuGokxmns/mitSKXSkEoJZDLR32IpRSDaX7W1tRXAiFWpra1NeN/tdgtWqNbWViSTyXF9ddxuNw466KAJ161UKjNKzFMoFAqlNAgh4DhuQjdYOp2e8LsZbjCZDGq1WnidjxWokrAsi8HBQRBCYLfb/yeGjrh65N9PnvlmSQmgs0Gib4Pim+8lEgnI5XJIKhjvxLIs9tlnHwAj90uZTAar1Qq9Xo/XX3990u+vWLECS5cuxfHHH49p06bh008/hU6ny7rsp59+iksuuQTd3d2QyWQ44ogj8Lvf/Q4PP/wwrrzyStjtdkQiEcycORM33ngj9txzz4zvL126FDt37sT7779f+o7XENEKp+nTp6O1tRXbtm3D/PnzAYwo/ldffRVr164FMNI8UC6XY9u2bTjttNMAAP39/fj0009x66231mzsFAqFUs+MDobOJoQmsgJJJJIMN5hCoch4XUsBVCiEEHi9XsRiMdhstvHFLSVSYOG1wH4Xgd2xE8Q6A9BoAQAM/icGU6kUAEAul4/sP5sCtm8BAjsA03Rg3hmAVF70OKVSqdC6ZM2aNbBarbjwwguLXt9ERKNRnHTSSVi/fj0WLVoEQgi2bNkixAufffbZ+PWvfw0AeOqpp3DUUUfhk08+EQwd27Ztg1QqLfu4akFNhVM4HMZ//vMf4fWOHTuwfft2mM1mOJ1OrF69GjfffDNmzJiBGTNm4Oabb4ZGo8EZZ5wBAGhqasIPf/hDXH755bBYLDCbzbjiiiswd+5cIcuOQqFQpiq8FWgiC1CuYOjRFiCZTJYhgippQREDw8PDCAQCsFgsaG5uzr2wXAVIZOCyhFQxDAOFQgGO45BMJiEhLGSPLAXT9SYgkY0UzvznY8DZW0sST2N58803cemllyIej8NqteLFF1/E4OAgVq5cCZfLBZVKhfvvvx977LFH3uvcsmULDj30UCxatEjYtzPPPDPrsqeccgqefvppPPLII1i9ejVSqRRuvvlm3HHHHfjBD35Qln2sJTUVTu+//z6OOOII4TUfd7R8+XJs2LABV155JWKxGC644AIEAgHsv//++Pvf/w69Xi9857e//S1kMhlOO+00xGIxLFy4EBs2bGgYZUuhUCgAsrrBRguhXMHQo0WQUqmEVqsV3q8nK1ClicVicLvd0Ol06OzszPvYSCSSCa1w/OdKpRLc+w8CXW8CICOZd8DI648fAfY+uwx7MOK6W7ZsGZ555hnMmTMHfr8fwEjNw2uvvRb77LMP3nvvPaxevRrPP/983uv9/PPPMW/evLyX33vvvfGvf/0LAHDbbbdh+fLlGffueqamwunwww+f8GIHRhTtmjVrsGbNmgmXUalUWLduHdatW1eBEVIoFEp5GR0Mnc0NNllK/Oh4ID4YWiqVNrwVqJKkUikMDg5CJpOho6Oj4GPJMEzOexmPZKgLRCL7n2gCRixP/h2FDjnnWHbZZRfMmTMHAGA2mwEAL730Er744oui15tPYtbY5QGgt7cXf//73/HCCy8IGfH1jmhjnCgUCkXsTOYGyyclfrQAEkMw9FSC4zi43W6kUim0trZCLi+fuywrpulguMwgecKlAdO0qpQv+OCDD4r2xsyaNQtvvvlm3stv374d8+fPx/bt2/H5559j+vTpSKfT8Hg8OPbYY/HXv/61qHGIASqcKBTKlIcPhp4oI2yyYGheCMnl8oyMMCqAxAkhBIFAAKFQCM3NzdBoNGVbb87ffN4ZIzFNo2OcOg9CcuZ3IUmlynbO7Ny5E59++qngqjObzTjssMNwzz334IILLgDHcfjss88wd+7cvNd55pln4pZbbsGLL76IhQsXghCCP/zhD/je9743btmtW7fi+eefx6233gqbzYb+/n5hXEuXLq1r0QRQ4UShUBqI0cHQE8UDZWNsMLRUKoVGo5kywdBTiXA4DK/XC5PJBKfTWTZxyzAMOI7LbdGRykcCwT9+ZMQ9Z54OZq/ToZTKwbIsksmkcP4Vi0QiwebNm3HOOecgkUjAZrNh27ZtWLduHc477zysX78e6XQaZ599dkHCSaPR4Omnn8Yll1yC888/HzKZDEceeSTOPnskNuuhhx7CCy+8gGg0ij322APbtm2DzWYrej/EDEPyccw2OMFgEE1NTRgeHobBYKj1cCgUCrL0BxsjhCaaukanxI+tD0SDoacuiUQCg4ODUKvVsFqtZTsP4vE4duzYgc7OTsjl8pLcffw5z3Fcxes/TWX432z69OlCmYlCdAC1OFEolIpStv5g32SEjS6MSKFMRjqdxuDgIBiGQXt7e8UyrvMNEJ9sHXK5HISQ8fWfKKKBCicKhZI3o4OhswmhSvUHo1AKhRACj8eDeDyOlpaWineLKOc5PLb+UzVLR/h8PixcuDDjPY1GU1BgeKNDhROFMgUZ2x9srBASS38wCqVQCCEYHh7G0NAQrFZr1eNsCk3bzwVf/ymdTpcl/ikfLBaLUImckh0qnCiUOqac/cH4ytB8TSAqgij1RjQahdvthsFgKKiAZbnIK0C8CPjrkr+mafxTbaHCiUIRCbQ/GIVSHMlkEoODg1AoFHA6nTUTFVKpFCzLVsQqROOfxAMVThRKmaH9wSiU6sCyLNxuN1iWrU4By0koR4B4PtuoVfwTZQQqnCiigIvHkfZ6IbNaIRnbhbyGlNIfbLQrjPYHo1DKByEEPp8PkUgENpsNarW61kMCUN4A8cmQSCRQKBRlq/9EyR8qnCg1hbAsPOvWwb9hI0g8DkalgnnFcjRfdBGYMk4CufqD5RsMTfuDUSi1JxQKwefzwWw2w2q11no4WSlngHguGIYBwzDYf//9QQgR+u1ZrVbo9Xq8/vrrk65jxYoVWLp0KY4//nhMmzYNn376KXQ6XdZlP/30U1xyySXo7u6GTCbDEUccgd/97nd4+OGHceWVV8JutyMSiWDmzJm48cYbseeeewIY6ZN3+eWXg+M4tLS04NFHHxV66NUjVDhRaopn3Tr41t8jvCbxuPDatnp11u/Q/mAUytQjHo9jcHAQWq22JoHf+TJZgHiKS+GZ/zyDnlAPOvQdOPFbJ0IuKd7FKJVKhSy46667DiaTCatWrYJMVt7bezQaxUknnYT169dj0aJFIIRgy5YtSCQSAICzzz4bv/71rwEATz31FI466ih88sknaGlpwSWXXIInnngCu+++O37yk5/gnnvuwU9/+tOyjq+aUOFEqRlcPA7/ho1ZP/Nt2Ij48ceDyVJ7hfYHo1CmDul0GgMDA5BKpXA4HKJ3R+UKEE9xKfzo7z/C+4PvQ8bIkCZp/PnrP+PexfeWJJ54eAv522+/jcsvvxzxeBzNzc148cUXMTg4iJUrV8LlckGlUuH+++/HHnvskfe6t2zZgkMPPRSLFi0StnXmmWdmXfaUU07B008/jUceeQSrV68GwzAIhUIARlrezJw5s+R9rSVUOFFqRtrrBYnHs38Yj4P4/WDa2jLeHl0/SC6XC/+nLjMKpbHgOA4ejwfJZBItLS1QKBS1HlJe5AoQf+Y/z+D9wfcBAGkyUirk/cH38ex/n8WSGUvKsn2WZbF8+XJs3boVM2fOhMfjAcuyWL16Na699lrss88+eO+997B69Wo8//zzea/3888/x7x58/Jefu+998a//vUvAMDvf/97fOc734FCocCuu+6KO+64o9DdEhVUOFFqhsxqBaNSZRdPKhVgMqGpqQl6vV6YjFiWRSqVQiqVQiKRQDgcRiqVyjpRSaXScQKLiiwKRdwQQjA0NITh4WE0NzdDq9XWekgFkcvi3RPqESxNPDKJDD2hnrJuf5dddhEa+La2tiKdTuOll17CF198UfR6C43bGj0n//a3v8W2bdswf/58XHnllbjlllvws5/9rOix1BoqnCg1Q/JNIPjoGCce2cknw9LWhmQyie7ubqhUKlgsFsHiNFkWzWiRxVfdjUQiEwaC8yJrrDVL7G4BCqWRiEQi8Hg8MBqNoo5jyodsQqND35EhmgAgzaXRoe+o2Dj4+k8A8NZbbxVd5X/WrFkFtV3Zvn075s+fD4/Hgy+++ALz588HAJx66qm47rrrCtq22KCP3pSa0nzRRZCffjoYXgipVDCccw52/dk1CAaD4DgOnZ2dMBgMGBwcRE9PD8Lh8KS1Unhfv1qthl6vh9lsRktLC9rb29HR0ZHx53A4YLPZoNfrIZPJkEqlMDw8jIGBAfT09Iz76+vrg9vtRiAQQCgUQjwezxmUTqFQcsM/IEUiETidThiNxroWTXyA+FhO/NaJWNCyAMCIpQkAFrQswAm7nlDW7e/cuROffvopAMDv9wMADjvsMDz44IOQSCSIx+P46KOPClrnmWeeiVdffRUvvvgigBFh+MADDyAcDo9bduvWrXj++edx+umnw2QywePxYMeOHQCAF198Ebvvvnspu1dzqMWJUlNiiQT0561E80+uRNrrhcRsRp/XC4lMBrvdjkAggJ6eHrS3t6O9vR0cx8Hv98Pn80Gj0cBsNpdsFRpdckA1SQ0pvrglb8lKpVKIxWITFraUSCQTWrLq+cZAoZQDlmUxODgIQgjsdnvZM8FqxUQB4nKJHPcuvhfP/vdZIavuhF1PKEtgOI9EIsHmzZtxzjnnIJFIwGazYdu2bVi3bh3OO+88rF+/Hul0GsuWLcOsWbPyfuDTaDR4+umncckll+D888+HTCbDkUceibPPPhsA8NBDD+GFF15ANBrFHnvsgW3btgl9Au+++26ccMIJkEqlaG9vx8aN2ZOC6gWG0MdkBINBNDU1YXh4GAaDodbDmVJ0d3fD4XBkxB319vbCZrMJ5uVYLIaBgQG0t7dnBIhGo1HhacpsNkOj0VR38HnAl0rg47L4/08ksrIFvlORRWk0CCHwer2IxWKw2WyTPrCImXg8jh07dmD69OnCfvBtUcQe0D5V27dk+80K0QGNIe8pdUkkEoFKpRoXrG21WuH1etH2TUadWq2G0+mEy+WC2WyGXq8HMPIEpNFowLIsfD4fPB4PtFotzGazaALA+UrhyixlFcbCW7J4gRWPxyds1MvHLYwWWMXGLlAo1WR4eBiBQAAWiwXNzc21Hk5FqJdrcGz7Fv7hrV7GXyuocKLUDK/Xi46O8UGRSqUSyWQyI7hSKpXC6XRicHAQ0WgUNpst4zObzQZCCCKRCHp7eyGRSGCxWOrqSVYikUCpVBYksngLViQSEV6PhRdZowUW/y+dICnVIhaLwe12Q6fT1X3gd75Uq4J4qfBzD8uy6O/vx7HHHpvxuUajKSgwvNGhwolSE8LhMNRq9YSWIaPRiKGhIZhMJuE9hmHQ2tqK4eFhIe5pdAwBwzDQ6XTQ6XRIp9Pw+XwYHByEXq+H0WgUjRWqHBQqsniBlU6nEY1Gc4qssQKLiixKKaRSKaEVSEdHR0Ndh7mYrIK4GJFKpWhra8N7770HjuNo+ZYJoMKJUhN8Pl9WaxOPwWBAd3d3hnDiaWpqgkqlQnd3N9ra2rJalWQyGVpaWkAIQSgUgsvlgkwmg8ViyUtsNBJ8M9B84i34mAdeaEWjUcFdODYccnRQ/Vi3IRVZFI7j4Ha7kUql0NraKsQsThVyVRAXM7yFeqrGP+UDFU6UqhMKhaDVanM+yTAMA5VKhVgslrVmk1KpRGdnJ1wuF5qamtDU1DThegwGAwwGA5LJJHw+H5LJpPAdOhlkwsc85CuyRge+x2IxQXRlyznJJrD4PoGUxoEQIpTqaG5uFmXSRjXIVUG8HqDxTxNDhROlqhBC4PP50NnZOemyFosFAwMDcDgcWT+XSCTo6OiAx+NBf38/Wltbc17UCoUCbW1tIIRgeHgY3d3dUCgUsFqtU+5puByMDlCfjNEiiw98z1X1PVs8FnUbiJ9wOAyv1wuTyQSn0zmlb7LV3HeZTIY5c+YIr++991788Y9/xK9+9SusWbMGVqsVF154YVHrPvLII3HnnXdi5syZSCaTQjmVFStWYOnSpTj++OMLWt+GDRtw7LHHCqUK6hEqnChVJRQKCS1UJkMmkwkVwCcydzMMA5vNhlAohO7ubrS3t09aC4ZhGBiNRhiNRiQSCXg8HqTTaRiNxrzHRimMYkQWL7Roax3xk0gkMDg4CLVaPWUCv/OlGgHiRqMR27dvz3hvv/32y/v7+bgUpVIpJBLJhK77fNmwYQMWLFhAhROFkg+EEPj9/rysTTwWiwU+n2/Si0yv10OpVKKnpwetra2TtmThUSqVsNvt4DgOQ0ND49q7UKrPaJFFW+uIm3Q6jcHBQTAMMy5Zg5I9QJykUhh6+mmkunsgd3bAePLJYMps8X7llVdw55134oknngAAvPfeezj88MPR19eHW2+9FSeffDI2bNiAv/3tbwgGg9BqtfjlL3+Jc845B36/H9OmTcOGDRtgNpsBAPfddx/eeecdpFIpPPLII5gxY4YQA0UIwZo1a/Dcc88hEongpJNOws033wwA+PGPf4xnnnkGSqUSp512GmbPno33338fS5cuRVNTE955552y7ne1oHcGStUIBoMwGAwFPX1pNBp4PJ68ntoUCgU6OzvR19cnVBXPF4lEArPZDLPZjHg8jsHBQXAcB5PJBK1WS5+gRcroAPXJ4EUWb8lKJpNChuFEIitbXJZEIpny5wMhBB6PB/F4HC0tLVMu4SJfxgaIk1QK3ef8ENH33gNkMiCdRvCZZ+H8wwMliaehoSHMmzcPAHDAAQfg+9//fsbnn3/+Of7xj39gcHAQBx98MI4++mgAwLvvvouPPvoIBoMBxx9/PC644AKcdtppWLt2LdasWYM77rhjZNyE4O2338aTTz6J1atX469//SskEgmkUimSySRWrVqFNWvWgBCCY489Fh9//DEcDgcee+wx7Ny5ExKJBMPDw2hqasKCBQtw5513ZrgW6w0qnChVgQ8YLcTaxKPX6xEKhfKq6i6RSOBwOOD1etHb2wu73V7wTU6lUlWsvQuldhTbWof/4zMMp3JrHT4+cGhoCFarta7dLdVgbID40NNPj4gmAPimHEj0vfcwvHUrjEuXFr2dsa66V155JePzU045BQqFAh0dHdhjjz3w5ZdfAgCOPvpoYV5977338OyzzwIAzjrrLBx33HHC93khdsopp+CCCy4Q3ufLorz00kv49a9/jWQyiYGBAXz++eeYPXs2mpqacM455+Dkk08uOBZKzFDhRKkK/NNGMTcRo9EIl8tVUDscq9WKSCSCrq4utLe3FxX8LZFIYLVaYbVaEY1G0d/fD0C87V0o5YNhGKHqez5FVMe21onH43XbWieejsMb88KqtkIl+9++R6NRuN1uGAwGGseUJ2OPUaq7R7A0CchkSHb3VHUc/OvR89joZcZa+HP91vF4HJdffjneffdd2Gw2XHDBBYjFYpDJZHj//fexbds2bN68GZs3bcITj20p1y7VFCqcKBWHEIKhoaGirE3A/240yWSyoN5PWq0WCoUCvb29aG5uhlarLWr7QH20d6HUjlJb60wksqrZWoflWNy1/S5s+nwT4mwcKqkKZ806C/836//g9XihUCjgdDrp+V4EvBCROzsyRRMApNNQOCeuaVcOnn76aVx++eVwu9348ssvsdtuu+Gjjz7KWGbBggX405/+hKVLl2LLli045JBDhM8ee+wxHHjggdi6dSv22WefjO/F43EwDIPm5maEQiH87W9/w/777w+/349UKoXjjzsO87/VhsOOOQVwfw69nCA08DUwezZQp+KbCidKxQkEAjAajSVN9Hz/OrvdXtD35HI5Ojs70d/fj1gsBqvVWvQYgMZp70KpHaW01gmHw5P2Lyy2tc5d2+/CfZ/cJ7yOs3Hc98l9CIVCuPKgK2nJjiIZHSBuPPlkBJ95NiPGSbPvvmg66aSKjmH+/PlYvHgx+vr68Lvf/S5r0sUdd9yBH/zgB7jhhhvQ2dmJjRs3Cp+p1WoccMABQnD4aIxGI5YtW4Y5c+Zgl112wQEHHACZTIZYLIaTTz4ZyVgE4NJYe/XFAIAVp52AFSsvhtG0Bu+8/2FF97tSMKSeK3SViUK6IlMKgxCCrq6uspj2u7u74XA4in7iDQQCCIfDaG9vL+tTM9/eJR6PN2R7F0p9MLa1zmir1ljGttZhJSyO3no04mx83LJqmRqvfe+1DLcd5X/E43Hs2LED06dPz/rwxHEcWJYVhCdJpTC8dSuS3T1QODvQdNJJZc+qEwuEY4GBT8FgfPIFGAnQMheowVyZ7TcrRAdQixOlovj9fpjN5rK4FUwmk9BVvdjvq1QqIe6pELdfLmh7F4oYKKS1Di+yEokEotEodvh3ZBVNABBLx/DJ15+gVd2asS2+rs/ov7Hv8a8ZhpmyMVEMw4wIiDQHSORg5PKSAsFrCW9nIYRk/D/bvwAANglFNtEEAIQDuBQgqb85kgonSsXghUSxsU1j0el06O7uLlo4ASMmZ6fTCZfLBbPZDL1eX5axAbS9C0W88AHr8XgciURCuLlJJBJwHAeD1AClRIkElxj3XbVMjbm7zBUsTnzG4eg/lmWFf/l4rbHL5EtDiTJCwIT6IQ+7AZARK4u2GdC3VTW+J5fIySp4JoE/3qOPPW9lH/0Z4aQgkExscZLUp6WNCidKxfD5fLBYLGWb1BiGgUajQSQSKSnQWyqVwul0YnBwELFYDM3NzWWfeGl7F0q14QsSxuNxxGIxJJNJ4TO+BINer4fFYkEsFkMgEADHcWhqaoLD4cDZ4bMzYpx4ls1cluGmG51xWIl9qLQoyyXCyi7KQv1AeBDCGggHhAdH/m/IL15zMitPsYKH/zeb4CkHI/F5aci1VjAR9/gFtM01cdOVAyqcKBWB4zhEIpGSg7HHYjab0dfXV5JwAkYmh9bWVgwPD6Onp6diVY9pexdKuSGEIJFICBYkvoM9AKHautFohEKhyDi/UqkUfD4fEokEtFot2traMs75VfNWAQA2f7EZsXQMaqkay2YtE96vBrUWZaPfy1eU8cH6fCNcQYCAQBr2INsVTsIepFUWkKyfZmcyK49Y5hK+ZRIhZOQcVNgBhgEJe0YsT6OtbnUKFU6UisBbm8oNX+smnU6XpSVKU1MTVCoVuru7YbfbhZikiWrZlAJt70IpBI7jBIEUi8UyyhUolcqMc2eimybHcRgeHkYwGIRcLofZbJ4w+1MqkeLivS/Gj/b8EXqHeiFPyuG0Oyuyb7WgUqKMDzTmA+15SDqR3UUFgAEHKUPAyBTC2BoBPgt0XBsjgx0ppRng0lAoNXVraeKhszWl7HAch2g0iubm5oqs32q1wufzoaWlpSzrUyqV6OzshMvlgs6gw+avN4+rZbNq3ipIJeWZcGl7F8poWJYVrEfxeFywcjAMA6VSCbVaDb1en7e4JoQgGo0iEAiAEIKmpiY4nc68zyuVTIVdrbuiq6ur6H2aiox2e3EcB5YwkIABgyyuNEYCiVRRt3WMxjLOypRtvxgJIFXUvWgCqHCiVACv11t2F91oVCoVBgcHy9p1XCKRoKOjA2v/sRYP//dh4X2+lg0AXLz3xWXZ1mhoe5epA1/skv8bHaCtUqmgUqlgNBqL/t0nc8UVikqlQjwep/XJCoCvIM8wDOQKJRid7X8xTaOp4/iesfD7PJWaZTfGL0cRDRzHIRaLlRyDNBlNTU0IBoNlXWeCTeBPO/+U9bPNX2xGPJ09Xbsc8O1dOjs7odVq0d/fD5fLhWg0WrFtUsoPH6AdDAbhdrvR09Mj/LndbiSTSWi1WtjtdnR0dKCjowPt7e2wWCzQarUF33g4jkMgEEBXVxc8Hg+MRiM6OzthtVpLvomZTCb4/f6S1jEV4GOnkskkOI4TSkIwDDMSx6NrAeFvtYwE0LWUPb6HYZiMHnL9/f2QSqVYs2ZNUeubNm0awuHwhJ+///77uOKKKzL2mT/fJntofvLJJ7HXXnth9uzZmDNnDu6++24AwIoVK7DLLrtgr732wu67746VK1fC5/MJ3zvllFNgMpmwdEwph1WrVsFms2HBggVF7WsxUIsTpax4PJ6KuehG09TUhO7ubjQ1NZVtnd6YN2ctm892foZdLLtAo9FMbI4uA7S9i/ghhCCZTArxR2MDtFUqFQwGA5RKZdnPk1JdcfmiUCiQSqXKatltJPjG5YFAAACE1jgZMIwQ36OQMoBEDpYAX/6jH8OeGJqa1dj9wFZIpaVd12azGW+//TZYloVUKsUTTzyB2bNnl7TOXMyfPx9z586FTCYraE7avn07rrrqKvztb3/D9OnTEYvF8Nhjjwmf33HHHTj++OORTqfxs5/9DCeeeCL+8Y9/AAAuvvhinHPOORkVzQHgjDPOwDnnnIOVK1eWZ+fygM7ClLLBsiwSiURVGuDy8R/xePmsQFa1FSppdreEWqbGHh17QCaTYWhoCD09Peju7kZvby8CgUBGbZxywbd3cTqdUKlU6O3tRW9vb1n3mZIbQgji8TgCgQD6+/sF65HL5cLQ0BCAkZuWw+EQLEitra0wGo1QqVRlFRypVAoDAwPo7u5GLBZDW1sbOjo6YDAYKiZstFottXqOgWVZuN1udHd3QyaTwW63Z2TTZYWRADIlWAI8c/t2vLz5X9j+Qjde3vwvPHP7drBs/iUVsq6eYXDIIYfg1VdfBQA89dRTWLJkCYCRiti77767EDv3n//8BwceeCAA4N1338UhhxyCvffeG9/97nezWpnWrl2LOXPmYO7cudi8eTOSySRefvllnHnmmZBIJBgcHMQRRxyBBQsW4Gc/+1nW8fFz4+23346f/vSnmD59OoCRunorVqwYt7xMJsNNN92E/v5+bN++HQBwxBFHZK27d/DBB2dNRHrzzTex//77Y6+99sLChQtzHb6CoRYnStmolrWJx2KxwO12o729vSzrU8lGAsEnqmWjVWoBJTIuXpZlEY1GMTQ0hERipHigVCoVrEblsEwxDAOdTgedTie0dxkcHKTtXcoIx3EZ8Ud8Bhsv0FUqFaxWa0Wa6042rnyz4iqB0WjE4OBgxV3v9UAqlRLKiVitVthsNgAo6EHmy7cG0PfvIQAAx46Iib5/D+HLtwcw6+DC+nCO5bTTTsOmTZuwxx57CDXjvF4vDAYD9t13X7z00ks46qij8NBDD+Hss89GMpnEFVdcgWeeeQYmkwm/+tWvcOedd+Kqq64S1vn+++/j8ccfx/vvv49gMIgDDjgAhx9+eEaiwvXXX48TTjgBl112meB2GwsvnP71r3/h8ssvz2t/pFIp9tprL3z55ZeYN29eQccikUhg2bJleOaZZzBnzpyyu5ypcKKUBZZlkUwmszaPrBRyuVyotVIu8XDe3PMQDAbxTO8zI7VsZGosmzlxLRupVAq9Xl81MUXbu5QGn8EWi8UyArQZhhECtJuammoa5FotV1w+yGQysCw7pd11fP01AGhubi7pOhv2xCCRMoJoAgCJlMGwJ1byOA866CBcdNFFePTRR7F06dIMQfeDH/wAGzduxMKFC/H444/jrbfewpdffol//vOfOOKIIwAAyWQShx9+eMY633jjDSxZsgQMw8BkMmHhwoX44IMPMkIk3nzzTcHStGzZMvz85z/POj6GYQq2yhdrxf/yyy+xyy67YM6cOQBGrMLlhAonSllwu93CE1g1MZvN8Pv9ZcviCwVDWL1gNa749hVF13GaTEwlk0kQQkoSU7S9y8TwqdHZWoxIpVKoVCohc1FM1rpyZ8WVC51Oh3A4XNb2RPVAJBKBz+eDXC5Ha2tr0bXWRt/8m5rVGaIJGLE8NTWX/sDJMAwOPfRQ/PKXv8QXX3yBRx55RPjsyCOPxMUXX4y//OUvmDt3LkwmE3p6erD33nvjpZdemnDcLMsKDYrHVhgfvdxk8w2/zKxZs7B9+3bss88+k+4Px3H45JNPcP3110+6bLURz6xBqVv4yrm1SFvmW7CUK74oHA5Dq9VCJVPBoXeUrfglL6ZaWlrQ0dEBp9MpTMZDQ0NwuVxFx0zx7V2czpFihd3d3ejv788IWG5E+ADt4eFhDA4OZsQfeb1epNNp6HQ6tLe3C/FHdrsdZrMZGo1GFKKpkllx5cJoNArxXI0O3yKpq6sL0WgUDocDbW1tJRWoHS0sdj+wFfYZRgAjliYAsM8wYvcDWif6ekGsWrUKa9euHRfzwzAMlixZgh/96EdYvnw5AGCPPfZAV1eXEEMUiUTwn//8R/hOMpnEwQcfjGeeeQapVAqBQAAvv/wy9t1334x1H3zwwUKA95YtW7KOiz8Gl112GdauXYudO3cCGHFzrl+/ftzyfHC43W7HXnvtVfBx2GOPPbBz5058+umnAEBddRTxUStrEzAyIWi1WkQiEeh0upLWlUgkKpotN5ZyW6Yatb1LsS1GxIqYXHH5MLqooxjEZiXgM+SCwSAMBkNZf4/RwkkqleDE1fPw5dsD/8uqO6D0rDqeGTNmYMaMGVk/+/73v497770X3/nOdwCMPHA9+uijuOCCC4Sg8LVr1wpN2RUKBQ444ACceuqp2GeffcAwDK6//nq0tbXhyy+/FNZ73XXX4fvf/z42bdqE7373uzmPwT777INf/OIXOPHEE4XaT6PLKFx88cW4+uqrEY/HceSRR2Lr1q3CZ0cffTQ+/PBDRCIROBwOPPXUU9h3331x7rnn4i9/+Qt8Ph8cDgfWrVuHU045BZs3b8Y555yDRCIBm82Gbdu2lXZwR8GQcqcC1SHBYBBNTU0YHh6GwWCo9XDqilQqhcHBQTgcjpqNgeM4uFwuweJSLP39/bBYLFAoFGUaWXngxVQ0GhUat0okkrzEFN/eJRQKib69C99ihI8/ytZiRKVSQS6Xi1Zk5GKsK85kMonGqjQZw8PDIITAaDTWeihlhS/5EYvFYDKZinrA4FuuTJ8+PavVnS+IWevf+g9/+AM+//xz/PrXv876+YTtUspAKpUSGicnk8maz7HZfrNCdIA4Z1BK3VBLaxMP3808lUqNr6OSJ7zbp9YXdDbysUwB2cWUGNu75NNixGAwiFbgFUKts+LKhcFgQE9PT8MIp4ky5CqBGALrf/SjH+G1117LGs+UV7uUEhHDMSgn9T8zUWoGXxxPDGKDT71tayuuIm+9Bb8WI6aUSmVB7V3SSRbRYBIagwIyReFPoHyAdiwWy4jZKleLEbFSb664fOBFOF9gsV4pZ4ZcvohBNNx7771Z36+Ulelvf/sbfvKTnwivCSE49NBDsW7dOuF1rY9JKVDhRCkat9tdtka7paJUKoWYoGIuyKGhobLVg6oVhYqp1tZWpNNp9Pf3A4AQNM1xBO8++zU+frEH6SQHmUKCvRZ2YL8TdoFEMj6jJpVKZWSwjR6PSqWCTqeDxWJp2PgYHrFmxZUL3o1R7tTualCuDLliEZtI4K9bABWxMh199NE4+uijhdejrfl8WQKxHZNCoMKJUhT8TbhY11gl4LN/TCZTQd/jY2ka8caer5hiGAZutxscx2HnOxF88apHWD6d5PDBc11gWRZzFrZUvcWImGkUV1w+6HQ69PT01I1wIoQgGAxiaGgIGo0GDoejotd4rnBhMV0TtW7KW0w9p3JT6vapcKIUhZisTTwGgwHd3d0FC6dixFY9k0tMBYdC+PIfXVm/98nLvdjj0GaYzea6DdAuB43oissHPsCZv+mKlUpmyGWDvxb4zgljt8U3wq01vJWJYRjIZDKkUqmqlSwZfQz4gqq1OocIIfB4PGAYpugHf/Ge/RTRkkgkIJFIRGVtAv5X/TkWixVUwTwcDtfNU3S54XuxRSIRxGIxRAJJcOnsT2NsisDTH0A0FhbilJRKJZRKZdVbkdSCRnfF5YPJZEIgEKhqa6V8GZsh19nZWZVzUiqVwuFwwOVyCfWJRiMGocl3WJBKpTW5TkcfA0KIMJZawTAMHA5H0WOgwolSMG63u+gg7EpjsVgwMDCQd3mEatduqjXpdBqRSASRSERIk1apVNBqtbBYLGBbOLyqGIltGotMIcG3du+ETCEVGjrH43GEQiGk02lhuUYSVVPJFZcParVaCK4WC9XMkJsInU6HGTNmZLXg9Pf312y+TKfT8Hg8UCqVMJlMNbsORx8DQggGBwfR2lqeop/FIJfLSxJuVDhRCiIej0Mmk9X8CWoiZDKZ0CognwvD7/dn7azdCBBCRqxIkYjQl00mk0Gr1cJms2X9DWUKKfZa2IEPnhvvrttrYYeQXTe6KOdY6l1UTVVXXD4wDAOFQiGK0h21yJDLhVQqHTfncBwnxABWm6GhIQwPD6O9vb2mvxXHcVAoFBnHgE8cqVfEefejiBaPxyNaaxOPxWKBz+eb9MlTzLWbiiGVSgnWJJZlwTAM1Go1dDodrFZr3jf+/U7YBalUGp+92g82lZlVlw+iFVWpGBAeBHQtgHy8K5e64vKDd9fVKsax1hlyhVALN10qlUJ/fz90Op0oBH86nW6460i8ZxxFdMRiMcjlclFPVMBI/zqPxzNpymu91W4aDcdxGdYkYMT8rNVq0draWtJEJZEwmLPQhrkL2yBh5UXXccpGTUQVxwIv3wy8fTeQigJyDXDABcARV4MDQ11xBaJSqTA4OFjVbY7OkFOr1RXPkCsXpRTlLRRCCPx+PyKRCNra2kQTgyqGGK9y01h7Q6koHo+nbmod6fV6hEKhnKXz66l2UzKZFKxJfM8wvsJ2tkyeUkkkErBYLFWdfCsmql6+GXj9NwC+CXpPRUFe/w1CoSCG97mYuuKKQKVSIR6PV1xkVjtDrtyk0+mqXEOJRAL9/f0wGo0lt54qNyzLZhVO9VzLiQonSl5Eo1Eolcq6MbkajUa4XK4JhZOYazdxHIdoNIpwOIxkMimkzWq1Wtjt9qqMma8mLBaKFlVcEva37gKDzExBBgT6zzbDcNwNWd12lNyYTCZ4vV7Y7faKrL9WGXLlJpVKZT1nywWfWp9MJtHR0SHK+TmbeJTJZFUTlZVAPDMjRdR4vd6aNvItFIlEAplMNmEMk1hqN/FxVuFwGLFYTLAm8a1QalkvqV5uVDlFlfdrMOlY1u8xqSj6//MxJObpkMvlUCgUkMvlkMvlohTUYkKhUAgtl8p5noghQ66cVNJVF4/HMTAwAIvFIurjlE6nx1km+fOHCidKwxKJRKBSqeruZsL3r8v2VFyr2k0sywouN74YnUKhgE6ng8lkqrtjLHYkhlZwMhUk6fj4D+UatO66J1IYEdh8cD0vCEYjk8kyxJVCoYBEIqkbcVkJtFototEotFptyesSW4ZcuahEfA+fzs+yLJxOp+jnjGyuOrlcjmQyWVFrXCWhwokyKV6vFx0dHbUeRsEoFAqk02nBisNTrdpNY4tLEkIglUqh1WphtVpF+7TFcVxDCIJUKoXePjfa910JyVt3ABnuOgY44AIwCg0UQM7MSr68BS+uwuEwUqmU4O7l4YvCjrZcNXKFdaPRiMHBwZKEUz1lyBVLOX//aDQKt9uN5ubmsgjWapBNPCoUCoRCoRqNqHQa7yyllJVwOAyNRiP6p5qJ4FOnR9dqqlTtpsmKS9bLDTSZTNb9E38oFILf74fD4YDMeR0gkwtZdUSuAfNNVl0+8C0qJrupcxwniKtEIoFQKJS1ICIvqMZar+oNmUwmtM8o5Nyu1wy5WsJxHAYGBiCRSOrCyjSabOcHb3GqV6hwouTE5/PVpbWJR6fTobu7WxBK5ardVExxyXohHo/XrXDig2U5jsvMwFp4LXDoFUgGXPAlZGjrmF72bfPZfbkyzQghSKfTgsDiGyaP7WUmlUrHiatatcvIhU6ny7usR71nyNWKUCgEn8+HlpaWglpJiRmpVCqK/n3FUr+zO6XihEIhaLXaunq6GQvDMNBoNIhEItAqJIgMfg29uvCg8HIVl6wHEokEjEZjrYdRMCzLore3F01NTWhqahq/gFwNhW0Gkl1dNUuF5jMkJ3PTjnYNRqNRDA8PZ2QLjl1XrQLbjUYjent7IVfL4Y15YVVboZJlCsdGyZCrNizLor+/HwqFgh4zkUGFEyUrhBD4fD50dnbWeiglYzY2Ifzna4DPN0OXioLI1cABq0ZcNZLx6buVLC5ZD9RjNfV4PI7+/n7Y7fZJrWUGg2HSGl+1RiqVQq1W57Qw8N3uaxnYTkCw8b8bsfXVrYizcaikKpw16yysmrcKHMs1VIZcoZQSKzg8PIxAIIC2tra6tf42MlQ4UbISCoWg1+sb4ilH+tpaGD6+V3jNpGLfFEQEsPDaqheXrAfqaZ8DgQDC4TA6OzvzsrY0NTXlrPFVL/AZmbUMbL9r+114rOsx4XWcjeO+T+5DKBTC8l2XN1SGXKEUk1GXTqfR19cHjUZDrUwihgonyjj40v2NYG1CKga8fTfGTz8E3Ft3oW+X70Om1le1uCSlPBBCBFdGIXF4EokEUqm0IVtBjKWSge0sw2LT55uyrm+raysuP/hyKGVTUzQBhdVw4uO/QqEQ7Ha7aDNuC2Ws9XM0Eokk72bsYqOxZw1KUfDBmw3xtBMeHOlPlgVJOgaHUQ6YWqs8KPGSa6ITE6lUCr29vUWnZZvNZvj9/innPpqIYgLbu4a6EGez1McCEEvHsMO9A9NM0yCXy0UZ2F5p8q2MnUwm0d/fD4PB0BgPq6PgOG5CYSSXy5FKpahwotQ//JNPvV/AqVQKgUAA8VACjhwFEKGrTYd3sVIP8U2RSAQej2ek1ECRFiO1Wg23213X/bKqzdjAdpVOBZVUlVU8qaVqWNVWRCIRJJPJSV2DCoWisMbNdcBk7Vb4ONJYLIb29vaGtH7msury1cPrsal24/1SlJIYHh5GU1NTXU5g6XQagUAA0WgUcrkcJpNpxKJw4IWZTV6/+V9w5pmQpwg0jWEVLwuJREK0MSmEEHi9XqRSqbLEf2g0mrJVvp5qEEIQ8ASwpHMJtny9Zdzny2YtQ7OpecLvj3YN8j0Gx7oGxZA1WAq5XHV8U16z2Qyr1VrlkVWPdDqd0+LEJ9/UG1Q4UQQIIRgaGqoraxPLshgaGkI4HIZMJoPJZBpfHuCIq0EIAXn7bkjSMeCbAoj6w66C2+uDz+eDzWYTrWCoJolEAjqdrtbDGAfHcejt7YVer0dz88Q35EIwm83o7++nwqlAeDep1WrFlQdfCa1Wi02fb0KcjUMtU2PZzGVYNW9VznXk6xrMN2twtOWKdw3WmmzWFkII3G430um0aJvylpNs7VZ46rl6OBVOFIFAIACj0Sh6axPLshgeHkY4HIZEIoHRaITZbJ543BIpkof8BMOzVsCm5kbcc3I1JABaW1uRTqcxODgIAGhpaWlIk3m+JBKJilRVL4VEIoG+vj60tbWV1awvlUpBCBnXkocyMcFgEIFAIMNNevHeF+OsGWehx9+D3dp3G1fHqVjyzRpMp9OCwAqFQkgmk1kLimYTV5Wa67h4HFx/P0hzM5hvztlYLIbBwUFYrVZRPpxUgnQ6PeEDqUwmy5qAUA9M3TsEJQO+DYJYrU0cxyEYDCIYDIJhGBiNRnR0dOQ98aVSKcjVesA0vvilTCZDe3u7EKQpl8ths9mm5M2UECKq/R4eHsbw8DCcTmdFns6zteShjIfPYJTJZFkrfuvVeljl44tfVprR7ryJ4ol4cTzaclWpuCvCsvCsWwf/ho0g8Ti+UqlgWn422FNPBeqwXUqppNPpCS26Yn9AzwUVThQAI/3bclptagAv5oLBIICRwoWFiKXR5BP0zKe1x2Ix9PT01F2PuUaCEIKBgQFIpdKif/N80Gq18Pl8VDjlIJlMoq+vL2cGo1QqHSdExALDMHkVFC1H3JVn3Tr41t8jLE/icfjvuReGNIv2H19RmR0UMblinOoZKpwoIIQgFAqJwtpECEE4HMbQ0BAIITAYDGhvby/5KS2VSuVtHler1ejs7EQoFEJXVxeMRmPdBswXglhKEaTTafT29sJisVTcpcEwDJRKJeLxeF1m91SaoaEhBIPBSeNxGuHaKDXuiiQSiDy4Iev3Qlu2gLvoQkim2DlWr3WaJoMKJ4rwxF2ryY8QgkgkgqGhIXAcB51OB7vdXtYLrpBidDx6vR46nQ5DQ0Po6upq+NgEMRSEjEajcLvdaG9vr1oRQIvFAo/HA7vdXpXt1QMcx6G/vx9KpbKiFr96I1fcVdLlwn8TiazfI7EY0l4vFA5HpYcoOiY7d+qxJAgVTlMcjuMQiUSqnhJLCEE0GsXQ0BBYloVWq0VbW1vFnk6KvTgZhoHJZILRaITX6xUy8BqlS/loal2KwOv1IpFIVL3VhFwuRzqdrssJvBLwqfItLS0NeZ5XCpnVCkalAsmWYq9SISyVwUTPsQz4Iphirx03lqkTpUbJSrXjO2KxGPr6+tDd3Y1YLIaWlhY4nU5YLBZRm3QZhkFzczMcDgeGhobQ09ODZDJZ62GVlVoJJ47j4HK5IJVK0d7eXpMbi8FgwPDwcNW3KzYCgQDcbjc6OjoKFk0Mw4zLZptKSFQqmM4+O+tnluXLIVGr0NXVhUAgIBq3eK3hhVO9QS1OUxiO4xCNRstWF2ci4vE4AoEAkskk1Go1mpub67YXk1QqRVtbG1KpFNxuNxiGQUtLi6hFX74kEgk0NTVVdZvJZBK9vb1obW2tqXWjqakJPT09MBqNNRtDLeE4Dn19fVCr1QX1/RsNfxOcqvXQCCGIn3IyjIRgePNmkFgMjFoN8/Kz0XzRRWCkUjQ1NWF4eBhdXV1oamqqi/IvlYSvHl5vUOE0hfF6vRVz0SWTSfj9fiQSCahUKlgslpqZYyvxdCeXy9He3o5EIoHe3l4olUo0NzfXdapxrmJ1lYCvCVSpUgOFwDfDLSYWrt6Jx+Po7+8vuU4WfxOcqsLJ5/PBYDTCdPllUJ15BkhgCMbp0zICwvlSKlRAjSCXyxGLxWo9jIIR9SyfTqfxs5/9DNOnT4darcYuu+yCG264IcMcTAjBmjVrYLfboVarcfjhh+Ozzz6r4ajrA47jEIvFylo1mbfCdHV1wefzwWg0orOzEy0tLTX1YVcys0OpVMLpdEKn06Gnpwc+n4+a4SeBLzUQi8VEIZp4zGYzfD5frYdRVfx+P7xeLzo7O0vOKpTL5Q3nvs6XeDyOeDwO0zd14ji5HKpO54RZdLyA4uP5GtGFl0/MILU4VYC1a9di/fr12LhxI2bPno33338fP/jBD9DU1IRLLrkEAHDrrbfitttuw4YNG7DbbrvhxhtvxKJFi/Dll19Cr9fXeA/Ei8fjKYuLbsL+cCKiGsGHWq0WWq0WwWAQXV1dMJlMMBgMU/IpMhcsy8LlcsFsNovu+lSpVEgkElMiSJxvYaPVauEoU6aXXC5HJBIpy7rqCT4DcXQ5l3wzVBvZApXPA6tEIqnLuDhRC6e33noLJ510Eo477jgAwLRp0/DII4/g/fffBzCiaG+//XZcc801WLJkCQBg48aNaGlpwZYtW7By5cqajV3MsCyLRCKBlpaWor8/aX84EVFN94vBYIBer0cgEEBXV1fOooFiohr1VviWE3a7XbRZNDqdDpFIpKHLTsRiMQwMDJS9hU29BvqWSn9/P1pbWzPc9IUWfmxEASWG8iaVQtSuum9/+9t48cUX8dVXXwEAPv74Y7zxxhs49thjAQA7duzAwMAAFi9eLHxHqVTisMMOw5tvvjnhehOJhFCRenRl6qlCMdYmlmXh9/vR3d0t1HdxOp1ob2+HRqMR9cWdTCarGrfCMAzMZjOcTicikQi6u7tF3wW80hl1fr8fPp8PTqdTtKIJ+F8LlkaEEAKv1wu/318W19xYJBJJQ7ma8mF4eBgKhWJcYkOx/Q8byYXXyMJJ1Hv1k5/8BMPDw9hjjz2Ekv433XQTTj/9dADAwMAAAIyznLS0tKCrq2vC9d5yyy24/vrrKzdwEcOyrJDdNhmj+8NJJBI0NTXVZTG8WgX8SiQS2Gw2sCyLwcFBcByHlpYWUQYfV0o4EULQ19cHlUpVNpdQJZFIJGAYpuEqHrMsi97eXhgMhqrXbGtUUqkUhoaG4HQ6s35eyjzZCBaoRm23AohcOD322GPYvHkztmzZgtmzZ2P79u1YvXo17HY7li9fLiw39kSaLEbhpz/9KS677DLhNd9SYCrgdrtzxiCVsz+cWKj1k49UKoXdbkcqlcLg4CCkUilsNpuoJpVEIlF291QqlUJvby9sNtuEDVjFiMlkgt/vr3iZjmrBV2Nva2ubshlv5YYQgt7eXjgcjorOjfUsoFiWzcu6zBtFxDQfToaohdOPf/xjXHXVVfj+978PAJg7dy66urpwyy23YPny5WhtbQUAwV/P43a7c8bvKJXKKTeBcPE4EgMDSKXTUI06VkDl+sOJCTFMMnK5HA6HA/F4HL29vVCpVGhubhbF2FKpVFnFZSgUgt/vh8PhqDtzvUajgdfrrfUwSoZ3zaVSqapWY58KwfWDg4OwWq1VO7frUUDl+8DKZ2PWU5V6Ud8Zo9HouJu3VCoVovCnT5+O1tZWbNu2Tfg8mUzi1VdfxUEHHVTVsYoVwrJw3347vjrgQOz8zjGILj0V7ttvB5dOIxwOw+VyoaenB6lUCna7HU6nE0ajsaFEk9hQqVRwOp3QaDSiimEoxwRMCIHb7UYkEoHT6aw70QSMHAe1Wl2X9WV4WJZFT08PFAoF7HZ71W6uvPWgkYlEIiCE1CSBoJ5ioAoRTvWWVCDqWe2EE07ATTfdBKfTidmzZ+Ojjz7CbbfdhnPOOQfAyEm0evVq3HzzzZgxYwZmzJiBm2++GRqNBmeccUaNRy8OPOvWwbf+HuE1icfhW38PgsNBmC5cVdH+cJTc6HQ6aLVa4QnSYrGILkW/EPg4mqampqpXIC83ZrMZg4ODaG9vr/VQCiYSiQhNi6sdiM/fBOtRMOcDy7LweDwZpQdqQT1YoPINkFcoFHX3kCLqs3vdunW49tprccEFF8DtdsNut2PlypX4+c9/Lixz5ZVXIhaL4YILLkAgEMD++++Pv//973V9AyoXXDwO/4aNWT9LP/00TD+5EpIGF01idxuMngB9Pp8QW1PNmCCO40o+Rnz1abvd3hBucJlMBpZli86OqgWEEHg8HrAsW/VGyTx8QcN6crvkCx/XNJkFrxzXU77Ug4CaDLlcXnd9IhkiRhtflQkGg8KJZzAYaj2cspF0ufDfoxZN+PmuL2yDog4ynUohlUrB5/MJ8XBih+M4uN1upFIp2Gy2qoiQeDyOUChUdDB0IBBAOBxuuLi4YDAIlmWFatBiJp1Oo7e3Vyi8Wiui0Sii0WhDZu55vV5IpdJJz4dazjmEEAwPD2NoaKjmAqqnpyevpCtCCFwuV80TtArRAY0zy1HGIbNawUxUq0WlAqZAQ9Nq13AqFYlEgtbWVrS1tcHr9aK3txfpdLqi24zH40UJNL7UAMuy6OjoaCjRBAB6vb4uarzxsYp2u73mD371GK+SD2NbquSillm89RQDxVMvlrHRNNZMR8lAolLBvGJ51s+MZy3DYCCAgYGBhg7mrEa7lUogk8nQ3t6O5uZm9Pf3Y2BgoGKtCYqp4ZRKpQS3QCNaF4CRCV2hUIi2/xrf8y8cDqOzs1MUDwgymaziQr/a8C1V7HZ7XsuLoX5RPQqoekLUMU6U0mm+6CKEQ2Ek/vQnIB4Ho1ZDevJJaLnkEkhkMsRiMbhcLqjValit1oazGqRSqbJXSK4mCoUCHR0diMVi6OnpgVarhcViKetTWjKZLEhc8sHH9VhqoFAsFgt8Pl9GuRMxwNfIEltCQT1aDyYjW0uVXNS6btxoGiEGSoyI49elVAxGKoVs+dlQLjsTFokEipYWhJNJeP1+2Gw2qNVqdHZ2IhwOo7u7GwaDASaTqWEuqnpz1U0E/zuFQiF0dXUJk2G5fqd81lOrukC1hLc4iSnJgK+R1d7e3hDntpiZqKVKLliWFd3DWrUFVKFJFQzDiOoam4zGMi9QxsGfjBKVCrL2dkhUKhgMBqTTaUSjUWE5nU6Hzs5OSKVSdHV1YXh4uCHMuhzH1dxsXk70ej06OztBCEFXVxfC4XBVtstxHFwuF+RyeVXrAokBg8GAUChU62GAEIL+/n5Eo1E4nU4qmioM31KlUFe0GFx1E1EtF16hx6DeYuOocGpwotGo0IR3dIxMa2ur0D+Nh2EYNDU1obOzEyzLVvXGTMkfhmFgMpnQ2dmJWCyGrq6uouug5DNhJhIJdHV1obm5GcYpkFAwlqamJgwNDdV0DHxMmV6vR0tLi+iFa70/dPGlB9rb2ws+1mJy1U1EpQUUy7IFHQO+eni9QIVTgxMOh6HT6cZ1LpdIJGhra0NfX9+47zAMA7PZXJYbM6VyMAyD5uZmOBwODA0Noaenp+DJZ7L4puHhYQwODsLpdIrO/VAtJBIJpFJpzYKeg8Eg+vr64HA4alKtulDqzXqQjcHBQVgslqIEECGkbmJFKyWgChWP9XbO1MevSymaRCIBhUIh+JBHo1KpoFKpJnya5m/MHR0dGB4eRk9PDxKJRBVGTSkEqVSKtrY2tLa2wuPxoK+vL++b/EQZdbxbKJFIoKOjQ7Suh2phNpvh9/uruk2+3EM8Hq+r9jV8Ecx6hW+pIqag+0pTbgFVqHCqt3OGCqcGhj/pGYYZ56rjsVgsCAaDOS0Vo2sL+Xw+uFyuujjJ66nqczmQy+Vob2+HxWJBX1/fOFdsNrIJp3Q6je7ubuj1ethsNtG7haoB37uuWi6oZDIpBPDW229Qb26X0fAtVeqlYG65KZeAKibGqZ7OmalzV5mCjL4pjnXV8TAMA7vdjr6+vkkvDplMBrvdDpvNhsHBQaH4oVhJpVJTMoBWqVTC6XRCr9ejp6cHPp9vwt92rHCKRqNCMcV6cAtVE41Gk5FQUSmGhoYwMDCAjo4OaLXaim+v3NSb24Un35YqU4FSBVShMU4T3Z/EChVODQwf3wQgq6uORyaTwWq1YnBwMK/1KhQKOBwOmM1m9Pb25mXZqAVTVTjxaDQaoTBitkzJeDqO/mg/ktzIk57X60UgEBBNMUWxYTabEQgEKrZ+juPQ29uLVCpV1+7RehVOPp8Per2+LgvmVopiBVQ9BMiXQuPuGQWxWAwWiwUAJnTV8eh0OoRCIUQikbyfclUqFZxOJ6LRaMWKM5ZCoYUdGxWDwQC9Xo9AIICuri6YrWZs+GoDNn2+CXE2DtW7KpzccTJWzl2J9vb2Wg9XtEilUhBCKuICTiQS6OvrQ0tLS1UbPFcCqVQqygepXPAtVRwl9u6sJ6tJIRRSB4p/ILOyVqhkjZlQQi1ODczogmL5mEL54OJC3W+8ZUOpVIqqtH+9tlupBHympNPpxN3b78Z9n9yHOBsHAMTZOB7d+Si27NxS41GKH5PJVHarUyAQgNvthtPprHvRVI8U2lJlsnU1clxlLgsUy7G448M7cMijh+Cct87BIY8egjs+vAMsl//9RAz3jXxo3F94ijPWVJrLVTd6mYlKFOQDX5yRv6Bq3SB1qrvqspHkkvjTzj9l/WzT55swHGmMwqeVQqvVlq22GV9UlOO4unbN1TuFtlTJRaGxPfVKNgH1m7d/M+6B7L5P7sNd2+/Ka51SqVTUMbOjafxfeIoyOr4JyE84ASOBxVqtFn6/H2azueDtjjbp+v1+dHV1wWq11iTItZ5K+FcKjuMQi8UQDoeRTCYxEBsQJraxxNk4drh3wCKzZLyvUCigVCqFv0Z+op4MhmGgVCoRj8dLqmsVj8eFG3Yh7Two5aWYliq5aPTYnrHw871Sq8Qf3/hj1mU2f7EZP9rzR5O67fiSBPVw/MQ/QkpRRCKRjJTayWKcRmM2m4WYpWw1fvKBYRhYLBaYTCZ4vV74fD7YbLYpW0SxWqRSKUQiEUQiEXAcB4ZhoFarYTQaoVAo0Mw2Q/WuKqt4UsvU2KNjj4wJjhCCZDKJZDKJSCQCv9+fcR5JpdIMUSWTyRperFosFng8nqJdO36/H9FoFJ2dnQ0rQiUSCViWFbUVjW+p4nQ6y7ZOMbdbqSS+uG/CB7JYOgZvzAuHPnf8GF+SoB4eJKhwalDGTlqFpnva7Xb09PSU3MxVIpHAZrOBZVm43W6k02m0tLTQ2KMyQAhBLBZDJBJBPD4yaclkMmi1WrS2tmadwBORBE7uOBmP7nx03GfLZi4b91TIW1iUSuW4goCEELAsi0QigUQigVAolFF4UyKRZFir+EKs9Y5cLkc6nS7YoslnzWm12pKDkMUOn1knVhHBlx5wOBxlPSdZlp2S4QFWtRUq6cQPZFb15P3+FAoFIpFIJYZXdqhwakCyBSjm66rjkUqlaG5uxsDAANra2koeE1/dOpVKwe12AwBaWloqZpZtxDiddDotWJP4WAC1Wg29Xg+r1TrpDSAQCCAej+Mn3/4J9Ho9Nn22CXEuDrVMjVOcp+Dsb51d0HgYhoFMJhPE2lg4jhNE1dDQ0LgCd/XsAjQYDBgeHs67d18sFhOupalgdeXdLmLd11JaquQinU7XhcWk3KhkKpy+2+l48IsHx32W7YEsG/VUxoIKpwYkEomMy84pxFXHwwfChkKhsrUf4KtbJxIJ9Pf3Qy6Xo7m5uexPpmJ3E0wGIQTxeByRSEToEyiVSqHVamGz2Qqe8P1+P5LJpCCCL977YhxrPhYqiwpWtRVK6UhGpF6rL9txk0gkUKvVWW8kvAswkUhkdQHKZLIMYSU2F2BTUxN6enomFU6EEPh8PsTj8YZ2zY1FLpcLVlCxUcmWKlPVVReNRnGy7WTIZDJs/mIzYukYVFIVzpp1FlbNW5XXOmrZD7JQqHBqQCKRyLjA7mIrs9psNnR1dUGtVpf16UypVKKjowOxWAwulwtqtRpWq7VsN5Z6K0XAsqxgTeInD5VKVZbaWD6fD+l0elwbCaVUmRF3wFeQ7+joKHpb+TLaBTiWyVyAo79bKxcgb23LlbnJsix6e3uh1+sb3jU3FrlcjlAoVOthjINvqdLZ2Vmx9ddDcHM5iUaj8Hg8mN45HRdPuxg/2vNH8Ma8iHljmDF9Rt7rEdOD0WRMrV94ipCt8GOhrrrR3xt9Qy33ya1Wq9HZ2YlwOIzu7m4YDAaYTKaStyPmUgS8tSUcDgstPCQSCbRaLaxWa1nH7fF4AIy4RceOYSwKhQIajQZDQ0N5u6AqQb24AM1mM3w+X9a+ZtFoFIODg7Db7UUnWNQzYnS7VKOlylTL5OVFk9PpFPZbJVPBoXegN9jbsFmGjbdHU5yJxFExrjoehUIBg8EAn88Hq3XyIL9i0Ol00Gq1CAaD6OrqgslkgsFgKHoSSqVSook14DhOsCbxNxO+7IPJZKrYjd3tdkMikWT9zSaa0MxmM7q7u6HVakUrPPN1AYbD4awuQN5KVaoLUKVSIZFIIJaKwRf3CS5Pr9eLVCqFadOmTamb6GjEuN+0pUp5ySaaRmMwGBAKhWAymWowuspChVODEY/Hs95QSm2iaDQa4XK5Sq5fkwuGYdDU1ASDwSC0B7FarUU1m00mkzAYDBUYZX7b5q1J/BOoVquF2Wyu2qQ9MDAAhUIxYS2uZDKZ1RLCMAza29vR29s74YQoZqrpAmQ5Fo+6HsUTrz8x0rpGqsJJjpNwwbwL0NzcXJH9oxRHuVqqUEaYTDQBIzGyvb29BQkn/gFf7LGAVDg1GGMLX46m1Ewzu92Orq6uige58u1BxtaAKsSCVC0T8dgCk8CIm0Kn08FoNFZ9AiCEYGBgACqVKueElUgkJhRxMpkMRqOxohbGWlBuF+Bd2+/C5v9sFj6Ps3E81vUYDAYDLt774orvDyU/+JYqlYprmmrkI5qA4h7WeRev2N3bVDg1GPF4POvNrhyWA4lEgtbWVvT391elGSzDMGhubgbHcXC73fB6vbDZbHlfVJWwlkxWYLKWFhpCCPr7+6HRaCaNUUomkzmzivissUQiIfpJrFwU4gLs9/Tjoc8eyrqefCslTwXEEPNTzpYqU518RROPUqksaA7hy1iIfc6hwqnBqPREpVarIZfLEQwGq+YK4wVbOp2G2+0Gx3FoaWmpeAxOMQUmawUf+GowGPL6XfJpbVCuIqiNwFgXYCKUQIJLZF0230rJjY5MJkM6na5prFy5W6rkohFrx42mUNEEjMQ5BYPBvF3XfPVwsUOFUwORTCarMkk1Nzeju7tbEFHVQiaTwW63I5lMYnBwEBKJBC0tLeMEDBePg+vvB9fcDEkB8VilFpisFYQQuFwuGI3GgmrTTLY/UqkUVqsVbrd7XFbeVCdXpWSVRIUmWVMNRiUueLdLrYRTJVqq5ILjOFE9TJWTYkQTMJJAwWf25kO9VA+nwqmBiEQiRQVSF0qtA4gVCgUcDgfi8Th6e3uhVCrR3NwMhhB41q2Df8NGkHgcX6lUMK9YjuaLLgIzZkIrd4HJWkEIQU9PD8xmc0V+e51Oh2AwiGg0Oq6o6lRGJRsp7nffJ/eN++zMmWdiyDuEiDwCm802ZV1EvNulFlSqpUouGjX1vljRBIzcK/hSOPl8V4xlLLLReL/yFCYSiRTdeLRQZDIZTCYTPB4PbDZbVbY5FpVKBafTiWg0ip6eHuDhLYg+9L+4ExKPw7f+HgCA5aKLKlZgslbwoslqtRYkagp1KbS2tlYlKaDe4Csi861rVBIVzpo9UilZKpEiFosJzbLr9RwrBblcjnA4XJNtV6qlSi4asWp4KaKJR6vV5v1QX2y9wWpDhVMDUe00ToPBIKTd19IaodFo0NHSgq8efzzr574NG5E88UTozOayF5isFRzHoaenp+BsQ6DwJ2M+xmxgYKBqwrwekEqkOH/u+TjRdiJkTTIk/AlMc0yDVDJy8+SLu5arNlm9USvrQSVbquSi0SxO5RBNAKDX6+HxeKriDakW9PGxQahVb7bW1lYhYLuWpL1ekIl6Y8XjsCkUMBqNDSOauru70dLSUlTQ60Q1nHKhVqshlUprZkEQK8lkEnq1Hg69A23NbfD7/eOWMRgM6OzsRDqdRnd3t1AtvtGRSqVCrGC14FuqZKvmXo1tN4pwKpdoAkYEdL30oMsXKpwahEgkkrU2TaXhrRF9fX1V3/ZoOIMBmEAMMGo1ZA1Sj4hlWXR3d6Otra3oQqS5ajjlwmazwev1Vv1mKGZGi1CNRiMUPR0LwzCwWCzo6OhAKBRCT09PXWQPlUK1LWvVaKmSi0axOJVTNPEU2sBX7O46KpwahHA4XBPhBIzECqlUKgwNDVV923xxO28wCNPys7MuY15+dkHZdWJltGgqpc5JMRYnYORG2NbWVnORLCbG9oXUarU5LUp8JmhbWxs8Hg/6+vqoEC0TtW6p0ggxTpUQTcCIuy7fps8ymUz01wQVTg1CreulWCwWBIPBqj1FE0Lg9/vR3d0No9EIh8OBlksugeW8lWB495VKBcMPz0HzRRdVZUyVhHfztLe3l1wcLp8aThOhVCqhVqtrIpLFyNgSICaTKau7biwymQzt7e2wWCzo7e0Vhbu7nuFbqtSyL1qtwiXKRaVEEzCSnZuvm78eajnVv12RIorqvAzDwG63w+VyVbxgYiQSgcfjgdFozNgWI5XCtno1rOedh7TXC4nZDJfbDdR5JlgqlYLL5YLD4SibOC7l97FYLKJvBFxNRh9LqVQKhmHyvokqlUo4nU5EIhF0d3ejqakJRqOx5tdzuahG7zExtVSp19+tkqIJKKz9Si3LWORLfd9RKACAWCxWlcq4kyGTyWC1WjE4OFiR9SeTSfT09CAcDsPpdE54g5GoVFA4HJBpNHlbAMQKL5o6OjpEI1J4kdzX1yf6WIRaYDabCz7ntFqt8BDQ1dXVMEH41cisoy1VSqPSoomHb78yGfVgcaJnWgOQq7FvtdHpdCCElLX6K/9E6Xa70dbWhpaWlrwnyaamJoTDYdH7zLORTCYF0VSuoNNyCR25XI6mpqa6FqWVIleQeC4YhhGsqLFYDF1dXUKrn3ql0taDarZUaUSqJZqA/7VfmQxqcaJUhXg8LqqmiK2trfB4PCWLlWxxTMUIiJaWlopZwSpFIpFAX18fnE5nWTN1ypn5YzQaEY1GRf90WClyBQNPFiSeC765tcPhQCAQgMvlEv2NZCIqaT3gW6pka2pOmZxqiiZgJImI79KQC4lEIvoHXSqc6hz+qVZMvvVyZF9FIhF0dXVBIpGgs7OzpCdKlUoFQkheZmIxEI/H0d/fD6fTWfZg02Iz6iZiKrvsxmbUjaYcLmKpVIq2tjbYbDYMDAxgYGCg7gLIK+Wq40sPtLe3i2ruqxeqLZqAzPYrky0ndqhwqnPKfSMsF0qlElqttuCbR75xTIVSL1anWCyGwcFBOJ3OisRsFFvDaSKkUiksFgvcbnfZ1lkv5Lr2RgeJl4pCoUBHRweamprQ09MDr9dbN0K1UsKpFi1VciGGBJ18qYVo4uHbr9Q7VDjVOYXGN1VzwjWbzYhEInlZekqJY8oHmUwGtVqddy2RWhCNRuF2u9HR0VGxQNdKCG29Xo90Op2XGb6RyGVxAooLEs8F38JFoVCgq6sLw8PDohdQhWRT5UutWqrkol5KEdRSNAEjc0U+cU5ihwqnOqeQPnG1aKBot9vR398PQgi4eBxJlwvcqIDXcsUx5YPVaoXP5xPlzSYSicDr9VbM0sRTSg2nXLS1tdWlK6kUxtZwGkuxQeKTMVVbuAC1bamSi3qoGl5r0QTk335FIpGIei6hwqnOKcREXImnv8mQSqWwms3YcdNN+OqAA/HfoxbhqwMOhPv22xH+pvlpOeKY8oFve+H1eiu6nUIJhULw+/3o6OioyoRWqTot9eIOLReEkElFbilB4rmYii1cat1SJRdirxouBtHEk0/7lVo1iM4XKpzqmEKfcmphcQKA6B/+gMTmh4UmvCQeh2/9PfCsW1fWOKZ80Ov1iMViomk6GQwGMTw8DIfDUfMJrVQ0Gg0YhmmYGkTloNJ1xOqlhUs55p1at1TJhZgtTmISTUB+7VfEXsuJCqc6ptDGvnwV32rCxePwb9iY9bPEE38CanBxtLa2isIyMjw8jFAoVLXMoGqI5paWFtoIeBTlDBLPhZhbuEil0pL3XwwtVXLBsqwohZPYRBOQX/sVsddyosKpjik0MLwWrrq01ytYmsZCYjGka+A2UygUkEgkNQ1mDgQCiEajVXU7VOOpmGEYtLa2or+/v6LbqTUsy+Ydi1buIPFc8C1ctFoturu7EQgEah7TV6rbhU8csdvtZRxVeRGjxUmMognI7z5ELU6UilFoJkctXHUyqxWMSpX9Q5UKcaWyJhM7H49Ti237/X4kEgm0tbVVdUKrVukKlUoFpVKJ4eHhim+rVkyWUTeaSgWJ50JMLVxKtR7UQ0sVscU4iVU08UzWfiXfIPJaId4zkZITjuMKviBq4aoLRCKQf3dJ1s/My89GSiJBd3c3fD5fVccmkUhgMBiqfnP3+XxIpVI1yQoqdw2nXFitVgwNDYl68iuFQkVopYLEcyGWFi6lWA/qpaVKpRsZF4LYRRMwefuVWsXj5os4fmlKwUSj0YLim4DquuoIIULl8F2uvhqW81aC+WbyY9RqWM5bCdvFF8NiscDpdEKhUMDlcqG/v79qN1uTyYShoaGqCTaPxwOO49DS0lKV7Y2lmsVSRzcCbkQKsTgBlQ8Sz0WtW7gU66qrt5YqYhAp9SCagPzbr4gVcTllKXkTDodhNpsL+k61VDzLsnC5XLBYLEIMlm31aljPOw9prxcyqxWSUe47hmGg1+uh1+uRSCQwMDAAQgisVmtFnzQZhoHNZoPH46m4mHG73ZBIJDW9CVSqhtNEyOVy6PV6+P3+gs9VsVOocBodJF4rlw7fwiWZTGJgYAByuRw2m63ilhKZTFbwwxBfeqARsk2rRb2IJiCz/YrYx5oNanGqUwqduIHquOoSiQS6u7vR1tY2LnBdolJB4XBkiKaxKJVKOBwO2O12BL+p81TJCskajQbJZLKiT+ADAwOQyWSieHKu9iRlMpkQDodFHehZDMW4ZqoZJJ6LardwKeacE1tLFbFTT6KJJ5/2K2J111HhVIcUezJV2lUXDAYxODgotIUoBalUipaWFjidTnAch+7ubsHVVW5aW1sxMDBQ9vUSQtDf3w+lUtlwFpdCmMqNgEdTiyDxXIi1hYsYW6qImXoUTcDk7VeKsVRWCyqc6pBEIgFVDqvNRFTSVefxeBCNRsveZ41hGJhMJnR2dkKj0cDlcqGvr6+sFgy5XA6FQlHW5pO8aFKr1aKoPVPLG6JMJoPZbIbH46nZGMRCLYLEJ0NMLVzE2lJFrNSraAImz5wTcy0nKpzqkELrN/FUQjgRQuByuSCTydDa2lrRi1er1cLpdMJqtcLr9aK7u7tsYqe5uRkej6csx4ePz9DpdDAajaUPrgzUus6MwWBAKpWqSVZXuSkmo5WnlkHiuahGC5fJri0xt1TJRSnnQynUs2jiydV+Rcy1nKhwqkNisVhRQdPljnFKp9Po6uqC2WyuqlVFoVDAbrfD4XAgGo2iq6ur5EJ/EokEJpMJgUCgpLHxQrKpqQkGg6GkdZWTambUTQTfCFgM7qBSKCa+kKdalcSLpVItXPLJrBNzS5Vc1KJqeCOIJiB3+xVqcaKUlWIzEcoZ4xSLxeByudDe3g6NRlOWdRaKRCJBc3MznE4nJN/UgxocHCx6om9qakIoFCr6+4QQ9PT0wGQyiS4+o5o1nCaC/70qEU9WTUoRToB4gsRzUe4WLpPdBMXeUiUX1bbmNopoAnK3X6EWJ0rZSKVSkMvlRX23XK66oaEheL1eOJ3OosdSThiGQVNTEzo7O2EwGNDX1weXy5WzMu1E8BXFC4UXTVartSg3aqURg8UJgFB7rJzxZNWm1GMptiDxXJSrhUuum2A9tFTJRTWrhjeSaAJyP8xLpVLR9FscCxVOdUax8U1A6a46QggGBgaQTCbhcDhEUyl3NGq1Gh0dHWhpaYHf70d3dzdCoVDek71KpQIhpCDRxWf9NTc318z6NhnVruGUi9bW1oplSFaDcljvxBgknotSW7jkctXVQ0uVXFTLVddooolnsvYrYqQ+z9QpTCQSKfrmXIqrjuM49PT0QKPRwGazif7ClcvlaGtrg8PhQDKZRFdXV95tXQqxOvGiqaWlRfRtIcTym9V7I2CO40q2MIg1SDwXpbRwmUg41UtLlVxUw1XXqKIJmLz9ihihwqnOKGXSLtZVxwsPm80mqoDnfJBIJLBYLEK9mnzaushkMqjV6gmDFnlYlhWKfRZTHmIqo1KpIJfL627CLBdiDxLPRTEtXLK5XeqtpcpEVNpV18iiCajP9itUONURpTaSLMZVF4lE0NfXB6fTWdfigG/r4nQ6YTabMTAwgJ6engkvWKvVCp/PN6HQHC2axBA7lAuxxtI0NzfD7/eLtshdpamHIPFc8C1cbDYbBgYGMDAwkPf8wpceaG9vr3sxUEmLU6OLJiCz/cpYJBKJKB8uqHCqIyKRSMGNfUdTqKvO5/NhaGgInZ2dNeuvVQnyaevC17Xxer3g4nEkXS5w37gl+EKB7e3tohdNQO1rOE1EPTYCLmdvrXoKEs9Fvi1cEmkOPf4o4im2oVqqEEIqEp81FUQTz0TtV4ptEF1p6v+snUKEw+GSzdr5TNKjW4W0t7eXtD0xw7d1IYRgaGgI3d3d0Gg0sFgskEgk0Gk0GLztt/A99RRIPA5GpYLxrLMQPelEdHR2iiKjMB/EklGXDYVCAZ1Oh0AgUBep6Mlksqy/Ox8kXsoDkVjgW7jwDyMmkwkGgwEcAX677Svc//p/EU8TqOQSfH+eDdeesnethyxappJoAkbqOXk8nnGJT3wZC7F5O6hwqiNKKUUA5BcczLIsXC4XLBaLKNPqKwHf1sVkMiESiQiV0JlHHkXqkUeE5Ug8jsB998HEAPLLLqvhiAtDDDWccmE2m9Hd3Q2dTid6MVpqDaexmEwm9PX1NYRw4jEYDNDr9UJW65ZPQrjnjS7h83iKw4b3BqDTfYUrjt69hiMVJ1NNNAETt1+Ry+Wi7DZAXXV1QjldBBORSCSEuJ2pIprGwrd1Met0GHrooazLDG3aLLjt6gExW5x46qURcLmPZT0HieeCd3U3t9qx8e2erMv84R87EE811n6XylQUTTzZ2q+ItXo4FU51QrFtVvIlGAxicHBQyD6bChBCkE6nEY/HEQ6HMTw8DJ/PB7fbDe+//w1MUFuExGJIe71VHm3xiKmG00TIZDIYjUZ4RX5cK2G9q/cg8Vz4IinE09kDxqNJFp5QfdXvqSRTWTQB2duvyGQyUQoncc+mFIFwOFyxUgAejwcsy6Kjo6MuL1hCCFiWRTqdRjqdFv4/+t9sMAwDqVQKqVQKmUwGmUwGlUo1IjLUaoRUSiA+fmJn1GrI6iyFuh5+16amJrhcLsTjcdHFNPCwLFv2RAmNRiM0mK6H36kQmvVKqOVSxLJYltRyCUzq+k46KZeFdKqLJmCk/Upvb29GrKNYjwUVTnVCPB5Hc3NzyeshiQSSLhdkVisYpRK9vb3QarVlWXfJY5tAAPH/z+XOkMlkWQWQTCaDRCLJ+wJkWRYejwfJZBKGM5ch+MAD45YxLz8bEpHe2OudtrY29PT0CFWqxUglxtVIQeKjUcmlOGNBKx54q3fcZ2fv34EhnwcBQmAymaDVakX7m09EqSViACqaeMrZS7XSUOFUR5RyURGWhWfdOkQe3ID/JhJgVCpITz4Z9ssvg7bMDWnFIIAKJZVKCZa35ubmkdYrl10KuVwG/8aHQGIxMGo15KecAt2551ZkDJWgXiYiHqlUCqvVisHBQbS2ttZ6OFWjEYPEgZG+lufs2wK1So0//GMHokkWGoUU5xw8HZcu2g1SyUhtuUAgAJ/PB5VKBbPZLPokAZ5S261Q0ZQJ335F7DGZVDjVAeXI5PGsWwff+nuE1yQeR/rRRxExNkG7enXW79SjACqUZDIJt9sNYKQg4+gLlpFKYVu9GtbzzkPPP/+Jjj33BBQKdHV1walW10VtK7HWcMqFTqdDKBRCNBoVVe+/SorQ0UHi9XBe5UMwGEQ0GoW9rQ1X2Ntw4ZHfwj//vRNzdnVCo/yfMOKr+1ssFsTjcbjdbrAsC6PRCL1eL5q5IhulXF9UNI2Hb78y1gMiNjd2fc2oU5RSGvsCABePw79hY9bPfBs2In788WAmUPi82OFFkFgFUKHE43F4PB6hllOuJ1yJSgVJW5vgnmtvb4fL5aqLCa8eMuqy0draip07d6Kzs1M0zV9LLQcyGXyQuBjc5qUSCoUQCoVgt9uFa0Qll2LXFiNIOgkosx9HlUqF9vZ2cByH4eFhdHd3Q6FQwGKxiDJppdh2K1Q0ZUelUsHj8WS8xxfBFNPvT4VTHRCNRmE0Gov+ftrrBZkofT4eR4tSCWVHR9Hrryei0Si8Xq/QBLiYp0WFQgGTyVQX7iSx13CaiNGNgMVShLXcNZzG0ihB4pFIBENDQ3A4HOP2Q6PRIBwOT+qSlEgkQm21RCIBr9eLdDoNg8GApqYm0RwflmULFtNUNE3M6PYr/LERo3ASx6McJSelBiDKrFYwEwQzM2o15A3whDsZ4XAYXV1dCIVCaG9vL1o08RgMBjAMI/omtfVqcQJGKlHL5fJJmy1Xi2ocSz5IvF6JRqPw+XxZRRMwYlEotKChUqmE3W4Xsn57enrQ19eHxATlQqpJoa46KpomZ2z7FTHWcqLCSeSUI+ZBolKh6axlWT9r5AwxQojQ/iEWi6GjowMtLS1liyGx2WwIBAKimMAnoh5qOOWiubkZPp9PFAUiK21xAkaCxOu1phPv/s5V1mSiZq75wDAMmpqa4HQ60dzcjEAggK6uLvj9/oKbl5eLQlx1VDTlh16vz3gglcvlSCaTNRzReKhwEjmlxjfxpL77XRjOOQf4RiQxajUs561E80UXlbxuscH3nuvq6kI6nRYm2nLHyjAMA4fDgb6+vppN3PlQz5O0mBoBV0OE1msl8UQigYGBgbxqwZUinnjkcjlaW1vhdDohl8vhcrnQ29uLWCxW0noLJd+sOiqa8mds+xUxWpzq91F0ihCJRGCz2UpaRywWg0Klgua8ldAuPxuaVAoyq7XhLE2EEPj9foRCIRiNxqrUApJKpWhra0Nvb++E7glKaSgUCmg0GgwNDZUU61cOqvH71luQeDKZRH9/P5xOZ14PJ7y7rhydEBiGgV6vh16vRzqdFir/a7VamEymimco5hOPRkVT4fDtV/gkJLE9mFKLk8gpRzq52+1Gc3PziKtBr4fC4Wgo0cRxHDweD7q7uyGXy9HZ2Qmj0Vi1SUqlUkGn04muXUi91XDKhdlsxvDwsOiePCuBRqNBNBqti98vlUqht7cXHR0deVt01Wp1RSxDMpkMLS0tcDqdUKlU6OvrQ09PDyKRSM2OJRVNxZGt/YqYoMJJxJQjuyYcDkOj0UAqlVYlRqOasCyLwcFBuFwuaDQadHZ2CkHb1cZkMiGVSmUENdaaeqzhNBEMw6C9vb1mjYCrvc16CBJPp9NwuVzo6OgoyLLDC8NKwTAMdDodOjo6YLfbEY1G0d3dDbfbPa6JbCWhoql4dDodwuFwrYcxIY0xqzYopRYAJITA6/XC6XQCqEyfrVqQTqfh8XiQSqXQ3Nxc0ebHhdDW1oauri4oFApRVD6u54y6bPCNgH0+H6xV7hVYbREq9kriLMuip6cHDoej4ONSzdYaUqlUcHlGo1EMDAyAEAKj0QidTlcxQUNFU2mIvf0KFU4iJhwOlxTTEQwGodfrM0zo9XwRp1IpuN1uEELGVfkWA3ywuMvlEkWvtXqt4ZSLpqYm9PT0VL0tQ7WttWKuJM6Lpvb29pIeEKpdr0qj0UCj0QgtXrq7uyvS4oWKpvIwuv2KVCoV1bVAXXUippQbHyEEgUAAZrO5zKOqPolEAi6XS4jVcjgcohNNPDKZDM3NzaLIAms0ixOP3W5Hf39/VZ9Ia3Es+SBxMcFxHHp6etDW1laSkOSLGtYCvsVLZ2cnmpqa4Ha70d3djWAwWNA5lW1ZKprKB99+BRBfSQIqnEQKf1EWe/EFAoGqBkhXglgshp6eHvh8PrS2tqK9vb0uLCharRZKpRKBQKCm46j3Gk4TwTcC5nsMVoNaxAeKLUicF02tra0li8hKxznlC9/ixeFwgGVZdHd3o7+/P6+bNMdxGRYQKprKi0qlEpIIaim0s0GFk0hJJBJQFZn5xhd+bGpqEt5jWVY0Pb8mIxKJoLu7G8PDw7Db7bDb7XUnAKxWKyKRSNXryoylUSdwnU4HlmWrdvNNJpM1iVsTS5A4IQQulws2m63oeWk0YhFOPHyLl87OTpjNZni9XnR1dWFoaGhC4RqOJTAYYRFPsVQ0VYDR7VfEVsupvu5GU4h8+jlNhNfrhcViybiAxdbrZyyEEITDYfj9fqjVajgcjroRehNht9vR1dUFp9MpGt98I9Ha2oqurq6qNQKuVbZmrYPEedFksVjKloghk8lEW+STb/HCP4D29PRAJpPBYrFAqVSC5Qh+u+0r3P/G14inOKjkn+LUuRZc990FVDSVGb79ilKpxPDwcK2HIyD6O1Nvby+WLVsGi8UCjUaDefPm4YMPPhA+J4RgzZo1sNvtUKvVOPzww/HZZ5/VcMTlIRaLFZVRx3EcotEo9Hp9xvtiLUVACMHw8DC6urqQTCbhdDphs9nqXjQBI0+x7e3tcLlcVXe3iMW9U0kkEglaW1sxMDBQ66FUjFpXEieEoLe3FyaTSbQZfpViohYvNz/7Me58+T+Ip0aKMsZTHDZ96MHtL/y7xiNuPPj2KzKZrKqlJCZD1HenQCCAgw8+GHK5HM899xw+//xz/OY3v8nINLv11ltx22234c4778R7772H1tZWLFq0SNTFs/Kh2IwTj8eTteJwrVwNE8FX+e7q6gLHcejs7BxnJWsEFAoFTCYTBgcHq7rdRqrhlAu1Wg2pVCrqmi+lUqsgcUII+vr6YDAYytL2aSx8deh6gG/xYmtrx8Pv92dd5oF/7EA8JU4rWr3Ct18R231B1DPr2rVr0dHRgQcffFB4b9q0acL/CSG4/fbbcc0112DJkiUAgI0bN6KlpQVbtmzBypUrqz3kslDsTY9lWSQSCbS0tIz7TCyuOo7j4Pf7EQ6HhZgCsV0U5cZgMCAWiyEYDMJgMFRlm42aUZcNm82Grq4uQUSVm0IauVYCjUYDj8dT9fT9gYEB6HS6ip2zGo0GsVhsnHVczAwOxwRL01hiSRYff7UTdsPIPCuXy6FQKIQ/uVze8HNdJRCjwBa1xemZZ57BggULcOqpp8Jms2H+/Pm47777hM937NiBgYEBLF68WHhPqVTisMMOw5tvvjnhehOJBILBYMafmCi2se/g4OCEfe1qPflzHAe3242enh4olUpMmzYNTU1NU2YisdlsCAQCSCQSVdleI9ZwmgiGYdDW1laxEhBicHNXO0h8cHAQKpUqI8Gk3KjValEFiOciHo/D5XIhHfZDJc9+29QopNhrt2no6OiAw+EQYqJSqRT8fj9cLhd6enqEv97eXni9XgSDQcTjcdH1YxMLYmy/ImqL09dff43f//73uOyyy3D11Vfj3XffxcUXXwylUomzzz5biG0Ya2FpaWlBV1fXhOu95ZZbcP3111d07KUQiUTQ2tpa0HdSqRQ4jsuZ8VILkcKyLDweD5LJJKxWa8kNi+sVvjhmd3d3VYKZk8lkXT3Jl4pSqYRara5II2AxWO+qGSTu8Xggk8lgMpkquh2FQiGq2jxjIYRgaGgIwWAQKpUKra2tkMlkOPfbcdz58n/GLX/OwdOhko88nDIMA7lcDrlcPuFvxrIskskkkskkQqEQkslkhnhiGCbDYqVQKISYt6mETqdDb2+vkGEnhv0XtXDiOA4LFizAzTffDACYP38+PvvsM/z+97/H2WefLSw39kBOdnB/+tOf4rLLLhNeB4NBdHR0lHn0xVNMhdRc1qZakEql4PF4wLIsmpuby5LCXO9IpVK0tbWht7cXDoejohNAo9ZwyoXFYkF3dze0Wm1Z4/mSyWRFLS/5UK1K4nyjaovFUrFt8IjhBpiNVCoFr9eLVColBIePHuuli3YDAPzhHzsQTbJQyRice8iuwvv5IpVKoVarJ8xU5DgOqVQKyWQSsVgMQ0ND45IEZDIZlEplQ7sD+fYrfEmCWlt/AZELp7a2NsyaNSvjvZkzZ+JPf/oTAAhWmYGBAbS1tQnLuN3urHE+PEqlsuZPkBPBcVzB1ohEIgGJRCKKEyqZTArxGGJsi1JrVCoVdDodvF5v1iD+ctJoE+hkMAwDu92Ovr6+stbTEUtiBR8kXqnzxu/3g+O4qj6AMQxT1JxXbvhyKIFAQCiwOtHcJZUwuOLo3XHhkd+CJ5RAOuxHe6sNUkl5rzeJRJLzXkUIQTqdFqxW4XB4XK0jfh2jrVa1PtbFwB8DMbjNAZELp4MPPhhffvllxntfffUVOjs7AQDTp09Ha2srtm3bhvnz5wMYObCvvvoq1q5dW/XxloNIJFJwGQK3250hHMdSjeKX8XgcHo8HUqkUNptNFDcascK7XSKRyJRL8a40crkcTU1N8Pv9ZbOaEEJEcbOpZJD40NAQkslkwSECpaJWqxGLxWp2HbAsC7/fj2g0Cp1OV1D9OJVcig6zBgmtFD6fL+ccXAnK5Q6Uy+UZ4kqM7kCDwQCv1ysK0QSIXDhdeumlOOigg3DzzTfjtNNOw7vvvot7770X9957L4CRH3316tW4+eabMWPGDMyYMQM333wzNBoNzjjjjBqPvjgikUhB/eVisRjkcnlOt0wlzZvRaBRerxdyuRxtbW1Tzj1ULG1tbejq6hLM6+VkKtRwyoXRaERPTw/0er1oJtpywQeJl1NoBINBRKPRqt/4gRExWEqx32KJxWLw+XwghMBsNpdkxVMqlUgmk6KJvxlNo7gDVSoV0um0aKqHi/out+++++Kpp57CT3/6U9xwww2YPn06br/9dpx55pnCMldeeSVisRguuOACBAIB7L///vj73/9et4GxhZoi3W43HA5HWdeZD+FwGD6fT+j1RCtjFwYfLO5yucpekmGq1HDKhd1uR09PT8OVuyh3kHgoFEIwGER7e3tNjpNKpRLiqirNRMHe5cBoNGJ4eLjsiQmVpl7cgQzDQCqVIhgNIhlKwqq2QiWrXdwsQ6b64ykg9HUbHh6uWp2dbPCtDfINVA+Hw4jFYpM+LXm9Xmg0mqIqkY8dXygUQiAQgEajgcViEYULoxr09PRUJIEgEokIPfnKuc5EIlGQ5bIRCYVCiEajOeMdJ4NlWQwMDKC9vb2MIysNl8uFtra2kh9WIpEI/H5/xRMVJqNS1xbP2GBvg8FQ9v0lhAgZs1ON0e7ARCIhZHjzlMMdyHIsfvP2b/D4fx9HgktAJVXhrFlnYdW8VZBKyvPQXogOmNqPpSIjHo/nnX1GCIHX64XT6Zx02VJddXxblKGhIRgMBtrIsoxotVrEYjEEAoGypX9PpRpOueDbNcRisaJ7rIklGHU05QgSj0aj8Pl86OjoEMW1XG43VyHB3uWAYRgolUokEokplxBTDXfgXdvvwqZ/bxJex9k47vtkpKbjxXtfXIG9yg0VTiKikMKXwWAQer0+L4tPscUv+bYooVAIRqOx4dweYsFqtcLlckGlUpWliepUq+GUCz6WrNjaWWIUTqUGifOJHGJ5AFKpVIjH42U590sJ9i4Vi8UCt9stKuukGCjVHZhgE3jos4eyfnfzF5vxoz1/VHW3HRVOIiIej8NqtU66HCEEgUCgILNwIRMkx3Hw+XyIRqNTpi1KrbHb7ejq6oLT6SzZBTMVazhNhEQiQUtLCwYHB4sKfharCC02SDyRSGBgYEA0ogn4X2ZdKcKpnMHexSKXy8GyrCjKK9QTk2UHdg11IcFl77gQS8fgjXnh0OeO8y03dHYVEfk+QQYCARiNxrJPfCzLwuv1IpFIwGKx1GTymapIJBK0t7fD5XKV5aYmlpuiGNBoNAgGg0W1MhKjxQkoLkg8mUyiv78fTqdTVDd2jUaD4eHhgr9XyWDvYjGZTAgEAlUpINrI8PG0Q0NDYMFCJVUhzsbHLaeWqWFVT25sKDfiuXqmOPkW2SOECEFs5SKdTqO/v1/ohO50Oml9oRqgUChgMpkwODhY66E0HC0tLfB6veNiKyZDrNaD0ZXE8yGVSqG3txcdHR2i2x+JRFJQn7ZUKoX+/n709PRAIpHA6XSipaWl5qIJGGkPEg6Haz2MuoUX993d3WBZFg6HA7s4d8FZs87Kuvyymctqkl1X+zONAmAkwyWfp2Gv1wuLxZK3RSEST2IglEJzihX6KPGkUim43W5a5VtEGAwGxGIxBIPBojI8aZJsdhiGQWtrK/r7+yct31Ev5Bsknk6nhWxdsZYNYdJxEP8OMPpWQD7eZVftYO9iYRimLK7HqcRoy6FCoYDFYhln5V01bxUAYNNnmxDn4lDL1Fg2c5nwfrWhwkkkRCKRSVPSOY5DNBrNy4XGcgS/3fYV7n/ja8RTHNTyr/DDb0/HpYt2Qzo10haFYRg0NzeL0hUxlbHZbOju7i6qNRCt4TQxKpUKSqUSw8PDNe89Vw7yCRJnWRY9PT1wOBziPC84Fnj5ZtjfuhNMOg7INcABFwBHXA1IpDUN9i4Ws9mMgYGBhhHolSIej8Pn84Fl2aw9AUcjlUhx/tzzcVLLSZAapDWv4yTCK2lqko9LwOPx5B139NttX2V08I6lWNz58n8QDAZx/rcdoogHoGSHL47J14Up5EaRTCZF+SQuFqxWq9AIeLLzn+M40ceK5QoSZ1kW3d3daG9vF28LpJdvBl7/DST4xlKaigKv/wapdBqDs/+vpsHexSKTyUAIqXhD5nqE4zgEAgGEw2GoVKqC2nOlUinoVDpY9LWPHxO3dJ8i5HOBsSyLRCKRVxHLeIrFA2/syPrZH//pg8lqo6JJ5EilUrS1taG3t7cg9xut4ZSb0Y2AJzuuYg0MH43JZILf7x/3Psdx6Onpgd1uF+8+pGLA23cDGPs7EEjfuwetliZ0dHTUZbwl70aljBCNRuFyudDb2wulUinEpRUi6MXSbBugFidRkE+z18HBwby7lntCCcRS2YNGYykWnlACHebSqohTKo9KpYJOp4PX6837iVus6fNiQi6XQ6/XIxAI5KyuXg/CSQgSj4chjXkBXQs4qRI9PT1obW2tufWRt7yk0+lxfySwE22paNbvSdIxSOI+QF2f57JWq4XX6xVl/7pqkU6n4ff7EYvFoNFoSq52n0qlCs6KrRQFC6edO3fi9ddfx86dO4V4m/nz5+PAAw/Mu+o1JZNwOJxTFPEl7PM9vs16JdRyaVbxpJIxkKYiYFklNSPXAXzaeT7iGqA1nPLFZDKhu7sbOp1uQnGUTCbFb+3gWLR8di8k790LpGMgcjVCs5bBtnhN2efjXCIonU5PaMGTyWQZf0qlEjKZDFKTfiSmKYt44mRqBJJyGOvY3aXVavNO+mkURgfxSyQSmM3mvB/4JyOVStWfxWnLli2444478O6778Jms6G9vR1qtRp+vx///e9/oVKpcOaZZ+InP/nJlOzXUwqTBfQWYm0CAKVMgqV7mrHpA8+4z354yC4waNUYGBgAy7JQq9UwGo2iOSEp4+GrX/OtCCZjqj7hFordbs/ZZDmRSJStDU7FePlmyN+6A7y7i0nFYPj4fjCGJmDhtRN+rZwiSCqVQiaTFRG0LR8JBH/9N8h01zFgDlwFpbYJfX19AICmpibo9fq6OrfNZjN6e3unhHBKpVLw+XxIJBLQ6/UVCeIXU8xYXsJp7733hkQiwYoVK/D444+P64+WSCTw1ltv4dFHH8WCBQtw991349RTT63IgBuNyUy5iUQCEokkb5cBIQQ9PT348XdmocnQiwfe2IFYioVaLsH35tlw2aLdIZUw0Ol0IIQgFovB4/EgnU5DqVTCaDTW3LxPyYQPFs91k6cUjkwmg9lshsfjyfpgwnGcaCbqrEwQI8SAgLx1FwZ3PwtpyKokgorkiKsBANxbd0GSjglZdcwRV+P/23vvMGfO+t77O6M26qvedrXPY2xTYopjcwi9BfsECOHwHkLAFROC4wI2IdSEcPKCzUmDk4AJGAPGFAMBEk7g5LVz4oBNMxgbjEkw4MeSdtV7HUkzc79/rGe8RdpVmRmNtPfnup7LXkkr3TuSZr73r3x/LtYEl8sFSZLQaDSQyWRgMpng9/uXotWfZVkwDLOyna6yp2Cj0YDZbEYgEND82mGUcx9DJqg8/frXv46XvOQlEz1huVzGqVOn8NSnPnXuxenFNFOR1abb7aLb7Y4dtZLJZBCLxSb64smiKRAIKCkGfijixw8+jCefeQLNWgU2m21sKzbP86jX60oR3traGjiOM8yHdZFoPcF9EjqdDhqNxljbCkKI4tdDmZzt7W0EAoEDqS0jvOf7kSQJPM+D53kMCr9A9AvnjX3s4Iofwhx8jOHb9wFg69QvkPCax/o4yeyum+E4Dn6/39DRctmTLRKJLHopqtHv91GpVCAIAjweD7xery7XCK2/j9PogIlk8KSiCdhp951k3hplh3a7PfZN6vV6sFgsM4smAOAsJsQ9VnAWE2yhEDKZjOJnsx95bAGwU99Rr9dRKpVgMpng9XrhdDqpiFogTqcTvV4PtVptZAppVXe2WhOLxRTrB6N8vncLpF6vpzhrMwyjDIP2nvy1sTVCxOKA1bcOLIFoAgBi5sD4j74oms1mJTooDyuWi4bX1tYMFyG02+2KybBRPluzIEkS6vU6Wq0WbDYbgsGg4RsntGTms2yxWESxWDxglf+kJz1p7kUdJ3ieH9sxVSwWJzJRGyea9sMwDBKJxET+QFarVTlBCYKARqOBarUKlmXh8XiWrt5gVQgGg9ja2lIunruhHk6zwbIsQqEQ8vm8MghYrwvdRALJ6x0vCEbUCBEwaD7hAgjNLvx+40eMZ3W75zgO8XgchBB0Oh3D1kO53W60Wi3dsxlqIA9PliQJPp9vYcOhjTYRYWrhdM899+CSSy7Bf/zHfyh/DMMwyolm2llQx5ndx28/7XYbDofjyB3UJKLJZDIp0QiTyYRoNIpsNjuxs62cvw4EAhBFEc1mE5lMBgzDwO12w+PxLEU6YFWIx+NIpVJIJpN7Ph/Uw2l2nE4nms2m0r2otmfM3AJpHI/UCOF7NwDDLiSzHXjaH8L7wj9BvdlCOp1GLBYz9Oei3+/PJfgZZqdm06j1UGtra9ja2loa4SSKImq1GjqdDux2uyHMkgVBMFRKduqj8drXvhZnnnkmbrrpJkQiEcOo+mVkXISAEIJyuXygCH/U4yaJNNlsNvT7feXDb7fb4XA4UKlUpp7ibTKZ4PP54PP5IEkSWq2WYtLocrlmO/lTpoJlWSQSCWxtbe3ZAVIPp/mIRqNIpVKw2+3oN5tgy2VIXi/YKdr6NRNI42BNO91zz3kL0C5gaFlDtdVDjDVhbW0NLpcLuVwOHMchGAwa8nyt5lw3lmWV85NcD1UsFhdaD8WyLMxms6F9weSoXa1WA8Mw8Pl8U81E1RojWREAMwinU6dO4Stf+QpOP/10LdZzrGi32yNbVZvNJtxu96FRnElFE/CocNr9OL/fj62trblOWizLwuv1wuv1Kl88anOgD1arFT6fD4VCAdFoFPxQRKrSgdcfBD3is8EwDCKhEE69730YfPkrAM+jwnHwX3oJQldfDWaX2NFdIB2FxQ74TsAGYFhtKq3bZrMZGxsbaDabSKVSiMVihkvn9no91bx+dmOkeqhAIIBKpaKkgo2CIAioVCrgeR5OpxPxeNyQG18juYYDMwinF77whfjxj39MhZMKdLvdA87FhBDUarVDvbCmEU3AjnBqtVoHbh+X8pmF3eFyanOgDx6PB+1OF+/72k/wmR9kH7Gd+IUyzNnEGmO3uEy0brwRg89+TvmZ8Dwqf/9R8D0eposvWrxAmoBgMIhyubynk8vj8cDpdCKfzyuCwijRBFEUNU8FLboeymazYTAYGKJInBCCVquFer2upDON3vVnJNdwYAbh9PGPfxyXXHIJfvrTn+Kss846oAJf9rKXqba4VWfUl6hWq2FtbW3sl2ta0QTs7LwEQThw+7iUz7wwDAOHw6HM1eN5HrVajdocaMBnf1LHjd/JKD/Lw5wB4C3nP3ZRy5obuf6PEAJCCCRJUv5//7/D7pvkfuU1+310PvmpkevpfvGLeMwbr4bF6E7iABwOB0ql0oHB4SaTCYlEAu12G6lUCpFIZOH1P3qzyHqotbU1NBoNrK2tafYahzEYDFCpVDAYDODxeDQxqdSKpU/Vfec738Fdd92F//N//s+B+2hx+OSMah2XDcXGRZtmEU3A4aZhcsqnWCxqtuugNgfawA9FfOKuh0fe94lvn8JVLzgdnMW0R4TMIzKO+l21O18Yhhn5TzYWHPVPnt02ze/J9DMZPNTvj1wL6fVAajVgCYQTsJOKr9VqI2sYXS4XHA4H8vk86vU6otHowr5/kiQt7LX1rofyeDxIp9O6CidZHDabTVitVgQCAcPWWR2GkVzDgRmE0xvf+EZcdNFF+NM//VPDh/eMzKjZY+VyeWxB3qyiaRI8Hg+63S5arZbmxcXU5kA9Dhvm3B3sGJ/GPTsnSb1FyLLRarVQ7nYBjgN4/sD9jN0O8xL507lcLqRSKfj9/pHvC8uyiMfj6HQ6ePjhhxEOhxcyl0/NwvB50KMeimEYpd5U67IFnudRqVQgiiK8Xu/CbATUxEjrn1o4VSoVXHvttVQ0zUm73VaiMMDOzkAemrwfNUQTy7KHqvZIJIJUKgWO43QLiVKbg9mQBSffaIEzM+CFg5Eeh9WEJ595ApzFOLs0IzIcDpHP52Gz2bB5xhn41ctfDuHWWw88zvqKV6DUaCBssxnqBD4OhmEmSg05nU6cOHEChUIBjUYD0WhU1+9cr9dTUvpGYX89VC6XAyFElXqoQCCAYrGIRCKh4op3kCQJtVoN7XYbHMchHA4bKr21SkwtnF7xilfgjjvuwGMe8xgt1nNs2C9iSqWSZqIJeLSzbtxJatHz0KjNweHIKU6e52Ey7bSan+734/ef3VdqmnZz2TNPUtF0CIQQlEol9Pt9RKNRWCwW5HI5xN58LfIMA+Ef/xGk1wNjt8P23/8fnHj729Hp9QzbmTYKr9eLVCp1ZGqIYRhEo1H0ej2k02kEg0HdCnF5nj/QIGMUtKiHslgsEEXxQP3ZrBBC0O12UavVQAiBz+cbG2WkqMfUwunMM8/EO97xDtx111144hOfeEDRvvGNb1RtcavK/i+NKIro9/sHonhqpueOEk7ATgQoFAohl8uNnYemB9TmYIdRswNDodCek+K1LzoTwE5NU3fwyDDns8PK7ZSDdDodlEolBAIBJT3TarXAMAzsTifsv/86xN72VgjlMszBILLlMgaCAJfLBbvdjlwup4ydMPIFSr7wj7M92Y/dbsfm5iaKxSLq9TpisZjmmxVCyFJEldWsh/L5fGPrzyZl98w+h8Ohy3u1KIzmGg5MOOR3NydPnhz/ZAyDhx56aO5F6Y3eQ37b7TYGg4Gy08pms/D7/XuGjKpd0zQYDFCtVvekB8dRKpWUC7WRkG0O6vW67jYHegx8nefv44ciSq0+Qm4bKsX8yKG1xx1BEJDP52GxWPa04wuCoERa6/U6zGbznlo/URQPdJ42Gg1FXBi52FaSJGQymUPtTUbR7/eRy+Xg9/s1PScacZDyNPA8j2q1OlU9FCFEGXs1DYQQtNtt1Go1sCwLv99vuDSnFgyHQ1QqlYmuXfOg+pDf3Zw6dWrmhVF2aLfbimgaDoeQJElT0QTshIiHw+FEjw0Gg8owYCNdfFfR5kCOqDUaDSWiFgqFpo6ocRYTNvw7xyUajSoXy2U6FlpBCEGlUkG320U0Gt0jdAghyGaziMfjYBgG7Xb7QP2JyWSCx+NBrVZTvrdyN2g2m4XL5YLP5zPksWZZFhzHTV2EbbPZsLm5iXK5jEwmg1gsprrX0nA4XPgoj3mZpR6KYRjY7faJ3xNZOPT7fbjd7qWyEVADo1kRAHMM+aXMzm7r/UKhsMc1V6vuuWlO6nK9UyqVOnIY8CI5zOZgbW0NDofDkBczuYar2WwqNVzRaFS1ULtcL1atVudKB6wC3W4XxWIRfr8fwRFdcZVKBR6PB1arVbFUGPV5X1tbQzqdhsfjUS72ZrMZyWQStVoNmUwG8XjckEIgGAwim81OHdlhGAahUAiDwQDb29vwer2qRqGN0lGnBtPWQ/n9fuTzeQQjMSVSvLsmkRCCRqOBRqMBi8VyICNxnDCaazgwg3C67LLLDr3/E5/4xMyLOQ7szoz2+32wLKuIKC0tB6aFZVnEYjFsb29jfX3dkAJkN6NsDiqVimFsDuSuQbmWxu12I5FIaCZKvV4v0uk0vF6vIS/mWiOKIvL5PEwmE5LJ5MjjLI9MkYddH3YhZxgGsVgM+Xz+wHBsn88Hl8ulibhQA5PJBJPJNHMbvNVqVQSiPDRYjQtZr9cz3LFSg0nqoRjWhI99N4cv/eR+9IYS7BYTXvesk7jiOZuo16oQBAEej2clbATmxWiu4cAMwqlWq+35eTgc4qc//Snq9Tpe8IIXqLawVaXf7ys7h2KxqMwu0kM0MQwzVTcHx3FwuVyoVCojd+tGxSg2B7KA63Q6ioDb2NjQ7UQYi8WQy+WWuoZkWgghqFarit3HOKEgSRJyudyeOpNms3loh5fVaoXVah1ZbG2xWJBMJlGpVJTok5GKdUOhEAqFwgHRNykMw8Dv98PtdiOXy8HpdM7dvWXkobdqMc4f6tP31vDpe4rK42TH/3anjXe+9Ikrf1ymYSVSdV/96lcP3CZJEq644gqcdtppqixqlZFPur1eDxaLBWazWbdIkzwvaZqQr8/nw/b2Nrrd7lIWIqphcyDxPKRcDlIoBPaIYzfKNmBR7cEWiwUcx6HZbOrS9LBoer0eCoUC1tbWjiy8zefziEQiewT0JBfyUCiEVCoFh8NxQHwzDINgMIh+v698n7U2lJ0U+cIzamLBtM+TTCZRr9eV6NM8F/njFE2RSwsa7S4++8OfjXzMF+8r4e2/bRzBbQSM5hoOqFTjxLIsrr32Wjzvec/DW9/6VjWecmXp9XoIBAJIp9NYX1/XNT1ns9nA8/zUuXJ5GPD6+vpSp32mtTkgoojS3/0dqp+6GYTn8SDHwX/pJQhdfTWYXV/kSWwDFkUwGEQqlYLL5TJsrdq8iKKIQqEAhmHGpuV202w2YTab92wEJj05MwyDcDiMYrE4tstHLqwulUpoNpuIxWKGOPahUAilUkmJcs/D2toaXC4XcrkcOI6b2prBiC3maiHby8j/djflMAyDMg/wgjTyd7uDne5YudGDsoMRzqW7Ue0q+Ktf/WrkIFnKXuQLtt1uB8uyutY02Ww21Ov1qX+PYRgkEglsb2+vTM59dzGnbANQKpX22AA0PvIRVP7+o8rvEJ7f+ZkArjf8wR7bAJ/PZ0hTRIZhEIlEkM/nF+rNpQWEENTrdTSbTUQikYk2BHLNyf6I1DRROYfDgVqtdugmRBZYPM8jnU4jFAotvG7RZrNhOByqtoM3m83Y2NhAs9mc2hhUj7EjWiIIwh5xtPvax7IsbDYbbDYbnE7nTk3TrnNmcCjCbvnpyHFJnIVFr1ZAUdhJhS7zRnWVmfpdefOb37znZ0IIcrkcvv71r+OSSy5RbWGriNx+Wy6XsbGxoXshuNVqxWAwmOl35c6OQqGguZ+G3oyyOajm86h/8lMjH1+5+Wbglf8doUjEcLn3UdjtdtTr9ZXqYuJ5HoVCYaoCWkIItre3kUgkDjx+lA3BYUxq+cBxHDY3N1EoFNBsNhc6UBfYiUCWy2VVR2Z5PB44nU7k83mlpueov9Hon0VCiCKOeJ5Hv9+HJD0aJTKZTLDZbOA4Dh6PR5nrOAncI4Xgoxz/f/9Zp+HMx5xU0s6iKMLj8dDxUwZjauF077337vmZZVmEQiH89V//9ZEdd8cZfijil1tlOM0EHrcb29vbunfPzXvCdrvd6Ha7K18zw3EcAiYT6v3+6AfwPNaApRBNMpFIRDHdW+aIoSRJKBQKIIRgfX19qshJuVwe6Tp/mA3BOOT6tUksH+SRJt1uFw8//DCi0ejCRIPD4UCpVFJt5IeMyWRCIpFAu91GKpVCJBI59G/s9Xp7bFgWASEEg8FgT+RodwrRYrHAZrPB4XDA5/OpWmdz7YvORI/v4XM/zKE3lOCwmnDZM08qjv92ux2JRAKEEDSbTWxvb4NhGPh8PsParGiBUVO6UwunO+64Q4t1rCyiRPCB2x/ETXedQm8owmZm8LtPCuKPf+sJCw/dz0I4HEY6nQbHcSvd+WEOBsFwHAjPH7yT41CVJAQM2O0xDpZlEQgEUCqVFn7BmpVGo4FarXbkRXkUvV4Pg8Fg5DzIWaMfsrfTpJYPDocDm5ubyOfzaDabE0VmtMDv98898mMcLpcLDocD+Xwe9Xp9bIRNFEVd0lCSJB0QR7uxWq2w2WxwuVwIBAK6RXVMLIPLzg3hzec/AdWucMDHSYZhGKUuUxRF1Go1lMtl2Gw2+P3+lT4HAzspUSOeY2kCVWM+cPuDe0KyfYHglh+V4PVu4y3nP1b39TAMA0LIzCds2RxTjl6saviY5Tj4L7kElY9+9MB9gUsvgSccVlqLnU6n6jtSLXC73UoR+zKdcPv9PvL5PFwu10wRM0mSkM/nx3baHWVDcBjRaHQqyweWZRGPx5XIzCIGBrtcLqRSKc26PeW/sdPp4OGHH0Y4HN6zSZymS3USjirGluuNvF4vrFarYaI1kiTByVnh5Cb7LppMJgSDQaVzs1KpYDAYTDzqZRkxohUBMKFw+q//9b/i3e9+N57xjGcc+rhWq4UbbrgBLpcLV155pSoLXGb4oYib7ho9ouYT3z6Fq15wuu4T7OU6p3lO1iaTCZFIBLlcbqq6kGVDetXvwj0YoH3rrSC9HsBxCOzqqpNHLXS7XeRyOUiSBLfbDa/Xa1hBKZuaTjsnaxFIkoRisQhRFJFIJGaOUORyOUSj0bHvyTxC0mq1guM4tFqtqawHFjkwmGGYneaHRkNTA0qn04kTJ06gUCig0WggEgqh8uEPH9mlOop5irFXEZvNhlgsNtWol2XEiK7hwITC6ZWvfCV+93d/F263Gy972ctw7rnnIh6Pg+M41Go1/OxnP8Ndd92Fb3zjG3jpS1+Kv/zLv9R63UtBqdUf2TkBLK7t1GazqdLR4nA40Ov1UK1WZ96tG5lGowHGZML6294K6U1vxKBYRGk4RPgxj9nzOIZh4HQ64XQ6QQhRfKIAGPIkJg+wrdVq8Pl8i17OWJrNJqrVKsLh8Fz+YY1GA1ardWwqTo0OM9nywel0Tl0ntb6+jkajoYon0jR4vV6kUil4vV5NP59yfVev18ND77sOwq23KvcpXaoAQm96k2bF2EZEFEVVNlfjRr2YzeaVGNNiRNdwAGDIhNVXg8EA//AP/4AvfOELuPPOO5W2doZh8IQnPAHnn38+Xv/61+Oxj9U//TQv00xFngZ+KOLsP799pHhyWE340Z++SPeIE8/zaLVaI2s9ZiGTySAUCi39F3Q3PM+jWCwecPmeZpK7fBJrtVrKCAajFHXK09mnLa7Wg8FggHw+D4fDgUAgMNfxGg6HyGazh3bd1Wo1RUzOQ7fbRaPRmNkjSRAE3QcGy7Uyeph0SjyPB3/j6WNrBp3/8CVYXS4lcmSz2Qz32VSTTqcDnuc1myU5HA5RrVbB8zzsdvvSWhtks1mEw2Fd1j6NDph4NVarFa95zWvwmte8BsDOTk42czRiKM0IHNZ2etkzT+oumoCd93F/geQ8JBIJpFIpJJPJlTjRiaKojOIYdfGatD5s/7yqer2OcrkMq9UKn8+3UKEpRwFyudzMIzjUhhCCYrGI4XCoyrBcQgiy2exI64HdTGtDMA6Hw6E4xs/y3i5iYLDf70cmk9FFOAnl8mjRBAA8jyjHwbpiNieH0ev1NJ3EYLFYFMuJZbY2MKJrOADMfPS8Xi+i0SgVTUdw7YvOxFXPPx12y86hdlhNuPS/RPH7vzG/e+8ssCyraounXAgqjzFZZggh2NraGjt81263o9frTf28ZrMZwWAQm5ubCAQCaDQaSKVSKBQKewpZ9cRms8FisaDdbi/k9XfTarUUd3O13OlLpRJ8Pt+hzzWLDcFhRKNR5PP5ub4HPp9PqUObxax2GliWBcdxM32mp0XuUh0FY7fDvESzMNVgVoE9C7K1gRxB397extbWFjqdzlKcs40Qpd/P8sXulgwTy+At5z8Wr31aDA88tIWnnnUmOIsJmUwGds621O65MnLHyjK3ugM788sCgcDYOhOn04l2uz3XTtFqtSo7wd1DPxfRmRcOh/Hwww/D6XQu5OQ0HA6Rz+eVESVqraHb7SrT5Q9DbRNGOco471BsPQcGB4NBZLNZzQdBsxyHtYsuQu3GGw/cZ33FK8CswHlwGtQU7JNy3K0N1GQ54nUrgM3MYn2Ng4V9dIRJNpuFKI4uHtcatXcaXq8XgiAYIoIxC3Kty2GFiBzHgR+XbpgBjuMQj8eRTCaVDqt0Oo1arbanMFYr5LEghUJB89faDSEEpVIJ+Xwe0WhUVT8j2SBzEnd7uaZBTbxeL7rd7tyRRHlgcDgcRiaTQavVUmmFezGZTDCZTKqm78cx+H9eAd/rXw88Emlh7HYELn8Dwm96I1Kp1MKir8cR2dpgc3NTEfupVAqVSmVh16RlgkacdMRsNqPf78NsNsP0SDv71taW7vPfLBYLhsOh6juMWCyGVCqlpIGWhV6vh06nc2S9j+yBpTaL7MxzOp1z1eZMS6fTQalUQiAQUK1BYTfZbHbiobpa+VnFYjHkcjkkk8m5n0uPgcGhUAiFQkHTerd2uw3O4UD4j94M2wWvAep1eE+cUHycOIcDW1tbCAaDhuyiUhNBEAxVt2NUawMjpxFpxEknCCGwWCx7ZsXJIdJ8Pq/rWjiO02SHuXsYsJE/9LsRBAGFQmHiAbgmk0nTHRnDMPB4PNjY2EAikYAoishkMprWJKhRm3MUgiBga2sL7XYbm5ubmhQk1+t1cBw3kQDUsujUYrHA4XCg2Wyq8nxyZDAQCCCdTqPT6ajyvDLyJkerIe2EEJTLZUUomxwOsLHYHvNLs9mMzc1NtFotFIvFpTl/zIJR5/TJ1gbr6+t7zj3ZbFbVSPukGNU1HJhBOF166aX41re+pcVaVhpZOO0XLG63G2azGbVaTbe1yF5OWmCxWBAMBnUXg7NwVDH4KBwOB7rdrsYr20GumUkmk4oXTjqdRi6XU/VEZjKZ4PP5UK1WVXtOGfmiKbcVRyIRTXawg8EAzWZz4vZurectBgIBVKtVVVOu8sDgVqulGK6qRSgUQqlUUu35diNHGOX3nWXZkWtnGEZxUs9kMiubMjKqcNrN7nNPKBRSGlqKxaJmAns/RnUNB2YQTq1WC+eddx7OOOMMXHfddUpKgXI0ZrN5ZB4/FAqh0+no0t0CaCucgB1XZJPJhEajodlrqEE2m0UoFJrqy+l0OlXf8U+C1p15Xq8XnU5H1ZNit9tFKpWC1WpFMpnUrPhUth6Ix+MTi7J2u63prEiGYRCJRFTfQMhWErKBpVrnDJvNhuFwqLpYGQ6H4Hl+T4SRYZhDRZ/X61WGUut1TtQTNQyI9US2NpAjxYVCAel0GvV6XdNaTKO6hgMzCKcvf/nL2N7exlVXXYUvfelLOHHiBH7rt34L//AP/0CL+w7hKP+feDyOQqGgi5oft+NTE3mXokfR6SxUKhXY7fapL577062LQO7M29zcVLoZU6kUyuXyXBc+uTZnXkRRxPb2NprNJpLJpKaRHQAoFosIBAIT2xiobUMwDrvdDoZhNLn4ywODa7UaCoWCKqmtYDCIcrmswuoeJZ/PHzAFncQSRa7tqlQqmkRCF40RW+wnQU9rg5WKOAE7Yeg3velNuPfee3H33Xfj9NNPx0UXXYR4PI5rr70Wv/jFL9Re50pw2JeFZVkkEglsbW2tRH5/d+egHh1i0yC79s4yKoZhGEOd9NTszLNYLOA4bubaHEIIKpWKUuR72Hw4teh2uxBFcaqaKT1TJZFIRDVhsx/ZQ83pdCKVSs29SXE4HOB5XrXva7vdHtkoMunGjWVZrK+vgxCC7e1tw51HZmEVzu3Ao9YGGxsbiMViShlBPp9XbWO5csJJJpfL4bbbbsNtt90Gk8mEF7/4xXjggQfwhCc8AR/4wAfUWuNKsDviNO7LY7FYEA6Hkc1mdVuTlphMJkSjUd3+nkkYDocolUoTF4OPYlzKdZHInXnr6+vY2NiAyWTC9vY2MpkMms3mxO91MBicqTan1+shlUrBZDJhc3NTl1SEKIooFotTjznRwoZgHCzLIhAIqB7J2Y3L5cLGxgZKpRJKpdJc32u/369KveX+gvDdHJWq208gEIDP50M6nV54tHdejCwGZkUrawOjuoYDMwin4XCIL3/5y3jpS1+Kzc1NfOlLX8K1116LXC6Hm2++GbfddhtuueUW/Pmf/7kW611a5JOZ2Ww+9MPkcDhgt9s1PdHK69AjLWi32+FwOFCpVDR/raOQJAnb29tYX1+fK2q0qDqnSZmnM2/a2hxRFJHNZlGv15FMJrG2tqbSX3E0svXAtO+lVjYE43C73eB5XlOxLQ8MtlqtSKVSMwsMl8uFVqs196Zqf0H4bmaZXuBwOLC+vo5cLqdat+IiWIbC8HmQrQ2SySRsNhtyudzUm7fdGCm6v5upfZxisRgkScKrX/1q3H333XjKU55y4DHnn3++rifQZYFhGGVW3GH1GH6/H9lsFu12WzNPE7lAXI/dj9/vx9bW1kJPGnK4PxKJzD3Sw+l0Ip/PL8VnfNzMPIvFMnZ6ut1uR71eR7fbHeuSTghBvV5Hs9lEJBLRffZerVaDw+GYOrK1qF1sLBZTBg5ridfrhdPpnHlgMMMwWFtbQ6PRmPnzLReEj5siMGuNpTzPr1gsotvtatahqSU8z8Pn8y16GZojWxu4XC5l4Hkmk4HZbB573lkmpo44feADH0A2m8WHP/zhkaIJ2Jm3dOrUqXnXtlLIattms020G4zFYqhUKpqFprXychpHPB5HPp9fWItxuVyG2+1WRbhp7eWkFbs784LB4J7OvP2fs0gkgmKxCLHXw2BrC9Iu+wOe55FOpwEAyWRS95PgYDBAq9WaqUZNaxuCcZjNZjidTl06TWWBwTAMMpnM1JFlr9eLer0+c9RpVEH4bqZN1e3/3UgkAqfTiXQ6rVtrvFoYuVNMK2ayNhj2YGptAUNjdlVOvfW+6KKLtFjHyiPXOE06WJVhGKyvryOdTmNzc1P1Ilubzab5ENHd7C5+19spvdVqQRAEVZ2qZRfxZdvxyuyfmVcul/fMzGMBmL/4RTz4uc8DPA+G4+C75GKIr3wl8EjR7iIiN7L1gNzVMy3tdhuJREKDlR2N3+9Xhhnrcex8Ph9cLhe2t7fh9XonjiDJ0YJ2uz21Uem4gvD9zz8vbrdb8XuKRCJzzY/Um2U9Z6iBbG0A7KQtC4UCRFGEx+OBx+MBCwLccR3I9z6M+LAHWBzAb1wBPP+dAGuceifqHK4jsnCatNbBZDIhFotp0mm3iKiJ1WqFz+dDsVjU7TUHgwGq1epEs8umQe+InZbs78zL5/P41Xvfh9YnPgk8EmkiPI/qRz8G3HqrpkNnj6JQKCAYDM70+nrZEIxD9mDS0xxWHhgsCMJUppJ+v39qG4DDCsK1wGq14sSJE6jVaoaooTyKVemoU4tR1gbNr70d5M6/BiNHmoZd4M6/Bu64brGL3QcVTjohf2mm3W1wHIe1tTVdxYaWeDweZR6b1qhVDD4Kh8Nh6ALxWZA782KBAISvfnXkY5qf/dyetJ2eyIXts9b9GaEwl+M4mEwm3dzngdkGBrMsC47jpvKgOqwgXCtk2xOGYbC1tWVoy4LBYLBUxpd6oVgbRINwP3ALGOwXmAT43g2GSttR4aQjs55Q5JoMoztxT0okEkG1WtW0y0gep6JVdMRut+t68dMToVwGGSOOSK8HQeOOz1GIoohSqTRX5FBPG4LDCIfDC5nHJptK9nq9iXyRpjHEHOUQrid+vx+BQEAVPyutMIJwNyqCIKCa+c9HI037GXaBdkHfRR0CFU46sb8eZtqTZjgcRqPRUH1G2SKKK/UYBlwsFrG2tqbZDm9R6R49MAeDwJiCb8Zu37lfZ2a1HtiN3jYE42BZFsFgULPZcIexe2BwKpU6NGpqMplgMpkmEiJHFYTrgd1uRzKZRKFQMOQms9frLX03mZqIoohqtYp0Oo1CoQCbfwPEMkZYWhyAK6LvAg9hdc/+BmO3QDjKy2kUcrF4LpdTrTZJ65l1h2E2mxEKhVQZ8bEf+aSpdffUPN1BRkWSJGyXSnD+3qtG3u+/5OI9U+31oFqtwul0ziWCjWam53K5MBgMFmboyHEcTpw4ceTA4EmG/05SEK4XJpMJGxsb6Pf7yOVyhqorEgTBEMdokUiShHq9jnQ6jWw2C4vFovjNOdeC6D35sgOJOoDZKRAfJ6oWABVOOiLvlmUvp2lReyzLIoUTsOOHZLFYVO3u43kejUZjrIeMmjgcjpVK1/X7faRSKYRCIWz88R8jcPkbwMipBY5D4PI3IHT11bqvqdPpzGQ9sJtF2RAcRjQa1WTjMCmTDAyWL/TjItN6F4RPghxVc7vdSKVSS2dZsGrINa2yCS8AZcKB2+1Wros8z6PyxD9A80mvB7E80iVpcQDP/qOdrjoDwRAjSfIFIdc+NBoNzU6uhUJBSR21220Mh8OZjdDa7TaazeZcY0OAnbqEcrm80BA7IQSZTAbhcHjuMLYoiprZN4yC53k0m01dRJrWNBoNNBoNJBKJPZEZiechlMso9PuInziha9SGEIJUKqWMkJmHTCaDRCJhuBRrtVoFy7ILN1OVJAn5fB4mkwnhcHhPSrTf76NarY48TxSLRdjt9qlrmzKZDDY2NuZe91EMh0Nsb28jFApNPdBbTSRJQjabxfr6+sLWoCeEEHS7XdTrdYiiCJfLBa/XO/Z7LEmSsmnrdrsI+9w7NU2uiG6Rpml0gLHOIivOvBEnGZfLBZvNNvfUcL3GrhzG7hTkPGkvuRhcz4vjoiN2akAIQS6XQ7/fHylQWI6DdX0dnmBwIv8xNcnn8wiFQnOLpkXbEByGz+dDo9FYuKHqYQODbTYbhsPhgTUuuiB8EiwWCzY3N9FoNOae4zcP/X7/WNQ38TyPbDaLdDqNXq+HSCSCZDIJv99/6Pc4l8shGo2iXq/vRJctdsB3wlDpud0Y70yyouwuDp/Gy2kcgUAAvV5vrpZ4oxixsSyLWCw2V7F4Pp9HIBDQtfhXNsFcVgRBQCqVgtvtPhBl2I9siKgX7XZbsUeYFyN3My3C2+kwxg0MHtVhZ4SC8ElgGAbxeBwWi2VhlgWrXBg+GAyQz+eRSqXQaDT2TCeYZLxVvV6H1WqF1WoFIWTukVh6QIWTTuy+wKolWOLxOEqlkqZt/XrBcRxcLtdMRna1Wg1ms1mzuX6HYYSo3Sx0Oh0lQjfJcdPTMFUQBJTLZcVheF6MYkMwDpvNBrPZbBhfsFEDgx0OB3ieh9TvALWH0a6VDFMQPilra2sIhUJIpVKqdidPgpHF+ywMh0OUSiWkUilUKhWsra1hc3MTkUhkqs3rYDBAs9lEMBhErVabu5ZRL6hw0hG1IzxymmsST5ZxsCy78DSBjM/nQ7/fn6rgWo66Lao41el0GuaCNynlchmNRgObm5tTXfj0EInySJV4PK7a98UoNgSHEQ6HF5pKGoXX68X6+jry+Tyq5RKiP7sRzF8+BvhfT4bjQ2chdP9HAckY545J4TgOm5ubKJVKuo6ckiTJUF2dsyCKIiqVijJnzul0IplMIhaLzRRN2/9d73Q6SzM6hwonnRg110yNk6TZbEYkEpk5zWW0Op14PH748MddCIKAQqEwd5H8PCyTg7gkSchkMjCZTDMJE7fbrbnje6VSgcfjUU3oGM2GYBwMwyAUChluQoA8MNj+/Q/C+v2/AyPsdN6xIg9mzlEYixKJ7COzFgVBQDabNZRYNRqSJKFWqyGdTiOXy8FmsyGZTCKRSMDhcMy1ucnn80o6r9vtzv18ekKFk07s/3LO4uU0DrmrZRZDPaMJp0nNMRdRDD4Ki8WyFKm63VYDs3ZzOp1OTeuceJ4Hz/Oqdpg1m01DFy/vxul0YjgcGur7CAAY9mC/7xM4eEmbfRTGousD5TE0shWDluUOoigasjFhHIQQNJtNZDIZbG9vK95Y6+vrcLlcqoibVqulDJMGdrpLlyVNB1DhpCu7P3DzdtbtZ21tDZIkTR0RMJpwAnbEiN/vR6Ew3mI/m80iFAoZosZi0ReBo2g0GigUCkgmk3MVqLIsq9nfKUkScrmc6tHDdru9kNq3WYnFYsjn88b6PLULOyMvRjHjKAwtP0vT4HQ6lXIHrTYFPM8bvjCcEIJ2u42trS1lGLQ8gNfj8agaCRIEAZVKRalhFEURhJCliAzLUOGkE/tPEjabTXXX4EgkglqtNpUQUqPDTwtkY7RmswkA4IciMtUu+OFOnt1uty/Ul2U3Vqt1YQ7Qh3GU1cAsqC34ZfL5PCKRiKo7cyPbEIzDZDLB4/GgVqsteimP4orsGBGOYsZRGCzLGsZ132w2Y3NzE61WS5MZgkYtDJe9lrLZLDKZDPr9PmKxGJLJJHw+nybfG0IItre3lcHMwPJFmwDA+H1/K8T+iJPaOxw5zZVOp5FMJie6UBo5pxwOh3Hq4RQ+ctcWPvXdNHpDEZyFxaueHMK7X3HOopenIBeIG2nyuSAI2NraQjAYVDXi4vF40Gq1VP1bm80mzGaz6oWhRr1gHcXa2hrS6TQ8Ho8xWrMtdkhP+0Mwd/31nnQdAcDMOArDaOOKGIZBLBZDo9FQzFLVioDwPG8oYSAbmg4GA9jtdl0j96VSCT6fb8/rdbtdQznPT8LybMWWnP27GK0iPXLhr1pjWRYJwzD4h5/38JFvnUJvuFMPxg8l3PzDAj5w+4MLXt2jGG30yrRWA9Ngt9tHjuaYFUEQUK1WNTlxGt2GYBzyRdwo3k4AkD3zEojPuFaJPBGLA40nvR7ic9420/MZJVW3H6/Xi0gkohg4qoERop7D4RCFQgGpVAq1Wg2BQACbm5sIh8O6iaZutwtBEPa4ci9TJ91uDLCdOT7sju5oGemx2Wzw+/0TG9TJu79Ff7n3ww9FfPLbqZH3feLbp3DVC04HZ1l8XtxIaYdyuYzBYIDNzU1NPmNyPdeoLtFpGRW2V5NlsCEYh2wIaIQarVarBStnh/m89wDPfxvQLoBxReAgJmTzhZlGpxjpO7Mfm82Gzc1NZLNZOBwOQ0WLpkEQBNTrdXQ6HVgsFvh8PtW80aZFFEUUCgWcOHFiz+3VanWhXdGzYqwrJUU13G43zGbzRLUSWtRbqUGp1VciTfvpDkSUWsYpal90gfi8VgPTYLfbVTEQLJfLWFtb02THuyw2BIcRCoVQLpcXKjAkSUKlUnk0IrhrFIbVagXHcWg0GlM/r9FSdfuRLQtkcT/rWgVB0PVzKEkSqtUq0uk0CoUC7HY7kskk4vH4QtPWo7zZ5K7yZfyeUuG0YLS82IZCIXQ6nSNDzjabTXcn3cMghKDRaICvF8GZRwsAu4VFyG2cmiK1U1jToIbVwDR4PB6laH9Wer0eBoOBZqm0ZbIhGAfDMAiHwwv1dpJniI0T4rLj87TWKkaOOO0mEAjA5/MhnU7PtLnUo85OkiQ0Gg2k02lsb2/DYrFgY2MDiUQCTqdz4XWslUpFma+6m2q1ikAgsKBVzQcVTgtETS+nccTjcRQKhUO9hoxiSSAPiMxkMpAkCY85kcTvP/sxIx/7qrPDKOa2DSP4FuUgrpbVwDTM21knSZLmc86MkOJSA4fDAVEUF/L97HQ6MJlMh36u5Dlw2Wx2quc2ao3TKBwOBzY2NpDL5abeMGglnAghaLVaitcSIQTr6+vY2NhQOpKNAM/z6Ha7BzZ0ckffMtY3AbTGaaHIFyAtO2dYlkUikcDW1tbYupdFttOLooharaZ0pe3v8Lj2RWcC2Klp6g5E2C0sXves03ZuJxJKpRIGgwHC4fBCvVI4jjswBFVLCCHI5/OKOZ2eJ0qGYZTU5CyvK0cxtKqpW0YbgsOIRqPIZDKa1a2NghCCYrF4oCZlFLtTdpNGEBmGMcyop0kwmUxIJpMoFovodruIRCITvRf9fl+1DlRZbNTrdYiiCJfLhXg8bthUl+zNtrm5eeC+brdrGDuZWaDCaYHItUVaf4AsFgvC4TCy2SwSicSB+/XenRBC0Ol0lPorn8+HQCAwch0mlsFbzn8srnrB6Sg2efD1Is58zEn5XkSjUYiiuHABpWeNk1ZWA9MgR9imff1GowGr1app+mJZbQjGYTKZsLa2pmtqo1AoIBwOT3xuCAaDSKVScLlcE13IlyVVtxuGYRCJRNBqtZBOp5FIJCba9M57fuV5HtVqFYIgwOFwIBKJGMOm4ggO2yAta1G4jPGP/gqjhZfTOOTp5uVyGcFgUJfX3M9gMEC1WkW/3596t8RZTEgGnCgIjgNOvCaTMQSUyWTSvCi50+mgVCohkUgs1DXd7XajXC5PJZyGwyHq9TqSyaSGK9upb1rWTqhxyN5OXq9X84smz/OQJGmqDd3ulN0kXXbLlKrbj9vths1mw9bWFsLh8Nh00zx/3+5zJcdxhpmSMCn1en3sBkkURTAMY9hI2SRQ4bRALBYLhF4LqPGPuPNqu0v2+/3IZrMj6z/mSb0chly42Gw2YbFYEAgE5gpd+/1+FIvFkZGzRQso2c9Jq6Jkra0GpmFaHzJ5ErpW1gO7WWYbgsOIRqPI5XIztf9PipwCnkXcTpOyW8aI026sVqtiWdDr9UZGAofD4VRiR95YdLtdWK1W+P1+Q5nqTspgMECz2Rz7Oa1UKku/saHCaVFIIpg7rkP8ux8CBH7HWO43rgCe/06A1U6Jx2IxpNNpxSdGRq5zUuuL2u12Ua1WIUkS1tbWkEwmVblgWiwWZbbRuOdblIByOp2oVquqCydJkrC9va1E6YwCy7ITR9hkx2CtoyWrYEMwDlmYtFotzcR5qVRCIBCYuT5s0pSd0e0IJkGe1FCtVrG1tYV4PK4cN4nn0Tp1CtwRAkEURcVryWw2Y21tDcFgcOEbo1mRN0jr6+sj/wZCCHq9HsLh8AJWpx5UOC2KO64D7vxrsHgknDvsAnf+9c7/v/BPNXtZhmGwvr6OdDqNzc1N5Ysud9bNI5xkF2i5xiQajWpyofR4PGg0GlhbWzv0cXoLKKvVqrobfL/fRzabRSwWM9ygULfbjXa7fWR0odfrYTgc6nKyXAUbgsOQhYnT6VS9+L3f7yvfkVmZNGW3zKm6/fj9ftjtdqRSKcQiETQ/9jFUP3UzCM+D4TgMLr0EoauvBvOIkJSj8K1WCyzLYm1tDX6/f2nF0m7y+TyCweDY836n01nqonCZ1Wg7WTaGPeB7NwDYf+IgIN+7Yed+DTGZTIjFYnvGsszq5SR7LsmGay6XS7Hy1yq64PV6p2oLlgVUIpFAvV5HOp3WzMZAzYvBIqwGpsHlch1Zo6eH9cBuVsWGYByyt1OhUFD1eeUUnRrv0yTGmMueqtuPbDS59Zd/hcrffxTkkfML4XlU/v6jKP3t36HZbCr2AXI37Pr6Olwu10qIplarBYZhDv3+1Wq1pU/TAUsmnK6//nowDINrrrlGuY0Qgve85z2KM+rznvc8PPDAA4tb5CS0CzsRphEwwy6yv7gPmUwGmUwGxWIRzWYTg8FA1Ysyx3FYW1tTzPWmtSTY77m0vr6ORCKhiy8HwzAwm81TR3f0EFAWi2VuawdCCHK5HPr9PjY2NgybepKL4Q9DjpbpYQ2wajYE43A4HCCEqPrZrVar8Hq9qn3WjjLGXIVU3X6Y4RDCV7868r7KzTdD6HSRSCSwsbEBj8ezEmJJRhAEVCqVQ0e6CIIAhmFW4vu5NH/BD37wA3zsYx/Dk570pD23/8Vf/AX+5m/+Bh/60Ifwgx/8ANFoFC960YvQarUWtNIJcEWUYZkHsDgQP+Mpym7E4/EoNvpbW1uKoMrlcqjVauB5fmZBJQ9bbDQaOztAnsdgawvSmBOyKIool8tIpVKo1+sIhUJIJpPw+Xy6fxn8fj8qlcpMv6ulgHI6nXMN/BUEAalUCm63e6p28EVhNpvHmqvW63VwHKdbtGzVbAgOIxqNIp/Pq7KZGg6H6HQ6R6a+p+EoY0yjf65nYVAoKJGmA/A8XKKwEqJhP5POnFxmp/D9LEWNU7vdxgUXXIAbb7wR733ve5XbCSH44Ac/iHe96114xSteAQC4+eabEYlE8LnPfQ5veMMbFrXkw7HYdwrB7/xr7E7XEQDSU98A0yPddQzDjLzwEEIgCAJ4nker1UK5XFZOoCzLKr/DcdyRO8hwOIz0ww+jc+PH0bnlFvyq3wfDcfA/kpcHy07suaQnHMeh3+/P1QmoRQ2U0+lEPp+f6SJkFKuBaXC73Wi1WgecgYfD4aGdNVqwijYE42BZFj6fD5VKZW57kVwup0kqdRZjzGVBFEV0u110u13F1Z0VRYDjgBHiibHbYV6QDYzWyI0fh52zVqUoXGYphNOVV16Jl7zkJfjN3/zNPcLp1KlTyOfzOO+885TbbDYbnvvc5+I73/nOWOHU7/f3jDCYd+7WTDz/nQAA6bsfBiv0AIsD5Gl/iPTJ1yBxRDs1wzCwWCywWCwHCmFFUQTP8+B5HvV6fU843GazKYLKYrEoDtDcP/4jqh//uPI4OS/faXfAXnShYR1qXS7XTCaM+1FTQE2SvhqFkawGpsHpdGJ7e3uPcJJ3oOM6a7RiVW0IxuH1epFOp6due99NvV6H0+nUTKhPa4xpREaJJJPJBIfDgbW1NVitVuVzXrz0ElT+/qMHnsN/ycVgDVinOC/dbheCICjZi3GocZ42EoYXTrfeeit+9KMf4Qc/+MGB+/L5PAAcyKtGIhGkUqmxz3n99dfjf/yP/6HuQqeFNQEv/FNkT/s9rK9ZAFcErMWOpCgqrrSzXARMJhOcTueBzgVCCPr9vuJCK9cHkX4fnU/dPPK5+l/+Ms58yx8Z9gu/traGbDar2hdSLQE1jSeWUa0GJmVUd1SxWEQgENDV3XiVbQgOIxaLIZ/PzxTZE0URjUZDU0PSaY0xF81ukSTXKrIsO1IkjSJ09dUAgOrNnwbp9QCOQ0CO3q8YoiiiUChMNJanVquN9N5bVgwtnDKZDN70pjfhtttuO3LQ5G6Oumi94x3vwJvf/GblZ71TCjKEEBAzB/gefW15JtI84mkU49J+/UwGD40ZIEp6PbTSabhPP92QuXn5Qqn2RXNeAcVxHHieP7LexshWA9NgtVrRb9dhG9bRZd0QRVF3S4BVtyEYh8Vigd1uR7PZPHLXvx95JIbWUUGjpuzmFUmjYEwmhK+5BsHLL4dQLiPX6yF0+ulLFUWelGw2i3g8fuTfJgg7tV1GvIbMiqGF0z333INisYhzzjlHuU0URXzrW9/Chz70Ifz85z8HgANttMVi8dDqfpvNZghHVkmSRl7wtRJPo7CEQmA4bmRRI2O3Y+hyIZvNKlEFq9UKh8MBu91uiHlJPp8PtVpNkzEyswoo2UH8MOHUaDSU3f5SR0okEYEf3wDzDz4GCD1wJg6xZ1wFRLU1ct1Pu91eqR3tNAQCASUdNunFqdVqwWq16nYeXHTKTguRdBgsx8G6vg5vvY5ms2kowagGlUoFLpdros/PKjiF72fxV75DeOELX4j7779/z22vfe1r8bjHPQ5ve9vbcNpppyEajeL222/H2WefDWCnzuGb3/wm/uf//J+LWPJUCIIw9iSil3hiHykEH5eXD+5KHxFCMBgM0O12USwWlVoes9kMh8MBh8Ohe1Gz0+lUpUD2MKYVUHa7fWwHieyXI/u4LP1O9I7rYP7u34J5pMmBFXldjFx3c1xsCMYhD5/N5/MTpXslSUKlUhk5tV4r9EzZ6S2SDsPj8WB7e3ulhBPP8+h2uxO9j7JtxmGBjGXE0MLJ7XbjrLPO2nOb0+lEIBBQbr/mmmtw3XXX4YwzzsAZZ5yB6667Dg6HA695zWsWseSpEEXx0KiNXuJp0rw8wzBKtG53MfBwOES320WlUlFqp+QTlcPh0PREJacg9w/+1YJJBdS4C7ggCNja2kIwGFzaQklCCCRJgiiKEPk2uO9+WBFNux61Y/D6nLdoPn8ROF42BOOw2+2o1+sTHQu9UnT70SJlZySRNAq5BlCLOaCLQJIk5HK5iUX3qhrSGlo4TcJb3/pW9Ho9XHHFFajVanja056G2267bSnqHQ6LOMnoIZ725+WrkgR3KKSMCDgKi8UCr9e752Qon9AajYbSicIwDOx2OxwOBziOU+1EctjgXy2YREANRIJUuY2I1w7OYlKidIu2GtgjfEQRgiAo/7/732H+QCaTCSaTCZZOFnZhjMv9sLtj9Oo7oc0fsovjZENwGJFIRBmlNO671el0YDKZFlZTtztlxwg8UHt44gHnRhdJ4/B4PCuTrpNF96TR3VqthvX1dY1XpT8MWZWBQXMgf6gbjcbUBZaTkslkDoQ2a7UaLBbLRIpcnLPbbhpEUdQkpC5JEnq9Hrrd7h7DSY7jlLqpWdMt6XQaoWgc5fYAIbcNnEW/OordAioQDOEjd6bx8TsfAi9IsFtMePU5Ebz23DDWE0cXUh6FmsJn9z+z2bzn54nWOewBf3HaSBd8YnGAeetDukSc0um0pp1hy0Sr1QLP8wiFQgfuI4Tg4YcfxokTJxYqMAZ8D73/8264H/g02DEDzg8TSVpHsdVG7pxdhq7Cw6jX6xgOhyM/W6MQBAGFQmFpag+n0QFLH3FaZgRBmHjnp2fBuMlkAsuyc/nDjIJl2QNWCXIOvNvtolarzVSELkoEn76vhs/c/YAiVl73rJO49kVnwsRqf3LdHYH6f//pPnzq7rxyX28o4hPfy8Jhd+At64xmwsdqtU4vfOblECPXxuMvgIOYoLWr0nG1IRiH2+1WLnD7v7uFQsEQbvTWb/8VLD/+GJRVDLsgd/7VjtHu2VcCWI5I0qSsQrpuMBhM3X1eqVRWxil8P1Q4LZCjapz2o6d4CgaDKJfLmg9nldN3u+sypi1C/8DtD+Lj39lSfu4NRXzojl8CAN5y/mNVWacseHaLnv3/3+kPceuPRg9f/fidD+HlZ3KwmVnjCB81eMTIFd+7YSfyZHGA+Y0r4Hr2W5HN5+F2uw+4iqvJcbUhOIxYLIZsNrsThRv2gHYBvNkLSZJ0n0wvTzkQBAHD4RDDXgu+737owKwvBoDz/pvh/K/vBmPVft6l3ixzuo4Qgmw2O5Wh7aoWhctQ4bRAJqlx2o9e4slms2E4HEKSJN27laYpQh9KO6JkFDd9+xT+8LknYTUxY8XO7v8/CpZ9VPTs/n+LxQKTyYQW6YMXRkeKeEGCzRvChn/FLgqPGLniOW9B9hf3IX7GUwCLHWYAGxsbqFaryGQymjnPH2cbgnGYzWY47Rz4b/wJuHtvAoZdWE0cYk+/Eoi+SxWbCEmSdoTQrn9y9HTUeuRJBw6xtZOeGwEz7AKdImA9Mff6jMYyd9fl83kEg8GpNvmtVmulNzRUOC2QWVuo9RJPgUAAlUpl4py21owqQn+41AIvjBY9vYGInz20hQ2/Y6zYkW9nWXbuKE+MNYMzsyPXY7ewCLkX7x2mGRY7iDcJkbVCviwzDINAIACXy4VMJoNQKKRqxOO42xAchv8nHwXu/jvlZ1bkgbv+BmCYsTYRhBCIojhSEO1PHbMsu0cQ2e12WCyWo79HDutOTdOoujizHXCGsSSx1qlY1nRdq9UCwzBTd8bV6/WVLAqXocJpSdFDPDmdTmWAsFG/7NE1B+wWE3rDgztdh9WEJ56xqUuhuCAIKOa2cdHTErjx25kD97/qKWG0GzVwKzroE3jUMX2/OLLZbNjc3EQ+n0er1UIkElHl80RtCMYw7IH53g0j7iAg370BpcdfAgHmI6NDdrsdHo8HZrNZve//mLo4gMHw3Ncjmy0utV3HYSxbuk4QhJn8vobDobIhXVWocFpi9BBPcpfB2tqa6s+tBtwjheByTdNuLnvmSV1E026rgbdvJEEkgs/+MIveUILDasJlz9wpVG/Ua9ja2kIikTCsEJ0Hh8OBdrs9MqrEMAxisRharRZSqRTi8fjcn1dqQ/AooigqsyiF0q8QHhHRAQBG6MLL9mAOnb64ovpH6uLI927YSc890lVnff47scmwKJVKqNVqiEajC7XuUJtlStfJg7pnOVetclG4DBVOS47W4kmewG5U4QQA1/zmGag36vjyT2voDcQ9YkVrKpUKeJ7f451z8dk+/NFv/doBawS/3w+O45BKpbC+vm6IkTVqwnEcKpXKoY9xu92w2+3IZrNzF44PBgPNrTmMhCRJijjq9/tKnR+wkwriOA42mw2e5ONALHYwwxEeWxYHbIEksMhOxEfq4rrn/CGGtW2srZ+p2FYwAMLhMIbDIfL5PGw2G0Kh0EpsNJYpXVcqleDz+aYWrvIw+WWevTkJq3XmPqZoKZ7krrdutwuHw5iFzY16DX98/uPwJ//NhVKrr4uPkyRJyGazcDgce4qT5YJ/u9WMDf/Br5f8+K2tLUQikZVKNTEMc6iNgozZbJ67cHxVbQjkC48skGQPI+BRl3ybzQan0wmLxTLyAiyKIhqPuwDe+2/a5+rO7KTJdPDWmgSz3Y3OIDZyPRaLBRsbG2i320ilUiuTvluGdF2324UgCDN5Gq56UbgMFU4rgpbiKRAIYHt725Amg5Ikod1uI5lMgmEYXbrWBoMBstnsSOFTq9WOjKJYLBYkk0lks1k4nU5N2/UXwSQ76nkLx5fZhkC229gtjmTByTCMMppkFg8jxSj3Je8Fs7am2ERIZjuEc/8AVtk+wgCYzWYIgnDoY1wuF5xO58qk74yerhNFEYVCASdOnJjp91e9KFyGCqcFoUW4VivxJDtLGzE1UiqVEAwGdQt9t1otVKtVbGxsjIx4dLvdiQYOsyyL9fV1lEqlhc0O0wKr1YrBYDDR1HRg9sJxo9sQyP5FclqN5/k90ThZHHk8HthsNlXe+wPTBR6xiUC7AMYZRjZbRBLMAQ+lRcGy7MgC9f0wDLMy6Tujp+uy2Szi8dkmHAyHQ5jN5pUuCpehwmlBaJVq0Eo8hUIhXefBTYJcEKuHyRohBMViEYQQJbq1n36/P/VFMBQKodVqIZ1OY319fenTT3a7Hb1eb2LhBExfOG4kG4L94mi3H5jFYlHSan6/X9P1jh3JZLEDvhNgsDPLLp/PIx6Pa7aOaZj24rwq6TujpusqlQpcLtdU3939v7/qReEyVDgtiFnMLydFC/FksVgUw0ijXNwLhYIuokkURWxtbcHn8x2a969WqzOdONxuN6xWK9LpNOLx+MwnLiNgt9tRLpdnaiaYtHBcbxuC3R1r/X5/T3rJbDbDZrPBbrdjbW1tId+NSedY2u12NBoNdDod3R3E1WTZ03dGTNfJY69mnacn1+Yt87lrGqhwWhDTjluZFi3EUyAQQLlcNoSNvlwXovUXtdfrIZ/PH3kM5bqVWY+zzWZDMpnE1tYW/H7/0tbvWCyWI+tWDuOownF+KOI/MyWcvh5WY7kKE3esPeJrZBSmHf4diUSQSqXmGqhtBJY5fWe0dJ0kScjlclP7Ne2m2WzOVEy+rBjnDHDM0DLiJKO2eHI4HCiVSob4whcKBUSjUU1fo1qtotvtYnNz88iLTLfbnXsXL79fhUIBPM/rWrtlJEYVjnN2Bz5w+4O46a5T6A3FmQY5H9axxrKsMubH5XKpa/qoEdOKJmDn2BotZTcPy5q+M1K6Tq6xnEdINxqNmaNVywgVTjowqkVbFEVdwppqi6e1tTXU6/WFdoL1ej3F3VgL5KGWHMdN3CEipwzmhWEYRKNR1Ot1bG9vIx6PL11kQC76nXdjsLtw/C//v//Ep+7OK/eNG+SsZceakZhFNMkYLWWnxkZs2dJ3RknX1et1WK3WuVLfg8FgKTYaakKFk07s/1AJgqCbL5Ka4snj8SCVSmFtbW1hX5RisahZy+twOMT29jbC4fDE7w8hBJIkqZrCWVtbg81mU94zI18E9jNu9MosMAwDXzCMW++9b+T9N931EF7xOAcsu7SlFh1rRmIe0SRjlJSdyWSCJEmqRN+XKX1nhHTdYDBAs9mcO1J0nIrCZZZrK7ukjIs46VkrIYun7e3tPSmKaZEHPnY6HRVXNzntdhsOh0OTNGe73UY2m8X6+vpUolar/L7dbsf6+jq2t7fR7Y4eoWFEHA6Hqusttfrgh2MGOQ8lmJw+bGxsKP8ikQi8Xi84jjPkRXMe1BBNwN6U3SKZxMtpWuT0ncPhQCqVQrvdVvX51UJO1y0COao+q/XA7ueZxn5kVaDCSSf2fzgX0Z2mlnjy+/2oVqsqrmwyCCEol8sT+SRN+7zFYhGtVgvJZHJqQdtoNDQrjDSbzdjc3EStVlvIMZ8FjuPQ7/dVe76Q2wb7GCd4h9WEiNcYTthao5ZokpGjTYvaBAHaCCcZl8uFzc1NdLtdZDKZPQX/RmCRwimfzyMYDM69eTdKnZbeUOGkA6PCsYsK0aohnliWhcViUfXiOAn1eh1er1fV4yaKIjKZDGw2G2Kx2NTPLYoiWJbVNN3BMAwSiYQy5mWSsSaLZNLRK5MiD3IehV6DnBeN2qJJJhKJoFQq7fGf0hMthRPwaPouGo0in88rXmxGYHe6Tk9arZaSOZiXRqNBhRNFG4zyRZVRQzyFQiGUSiWVVzYeQggajYaqw4Z5nkc6nVZSO7NQr9d1G4AcDAbh8XiQTqcnclxeNGp+7q990Zm46JwQ7NYdkeSwmnDV80/XZZDzotFKNAGLT9lpLZxkjJq+0zvqJAgCKpWKKpYyx7EoXIYWh+uE0T5c8xaMyyFeQRB0qdWSU3RqHcd6vY5WqzWR1cBhtNtt+P1+VdY0CS6XSzHLjMVihp1CPu3olSMhEt7w9Dje9fJz8OMHH8aTzzxBI00qscguO72Ek4zRuu/07K4jhGB7exuJREKV8+hxLAqXoREnHTCC79Eo5o08BYNBlMtlDVa2F1EU0e12VQkty0WRw+EQ6+vrc4mmwWAwdkK9llitVmxubqJYLC6sRuIo5NEraiFH9jiLCetrHKwm432f1EYP0SSzqJSd3sIJMFb6Ts90XalUgs/nU0UoHteicBkqnHTAaKm63cwjnjiOw2Aw0PxkWyqVEA7P7xQtCAJSqRS8Xq8qbcq1Wk3XaNNuWJbFxsYGer2eoeo2ZNQWTrujIXa7fam6DGdBT9EELC5lx7LswuqrjJK+0yNd1+12IQiCak0sx7W2SYYKJ50wYsRJZh7x5PP5UKvVNFrZjq/ScDicezZZp9PB1tYWEomEaukInucXmiqTL3Y2mw1bW1sLuwCNYt7RK7uRXfbl75DT6VxoJ5jW6C2aZIzQZbcIFt19p7VwEkURhUIBsVhMteekwomiOfujAZIkGU5IzSqeXC4XWq2WZhGPYrE4VyGjbGHQaDSwubmpWj2D3oNmD0OOoKVSKcO1XKtBrVbb41Qv10+tIosSTTKL7rJbFItM32mdrlPDr2k38kxOo13D9IQKJ53Y/SHT2/xyUmYRTwzDwOPxoNVqqb6efr+vjMmYBUmSsLW1BZPJpOqJA9iZY7fIsTP74TgOGxsb2N7eNkzEQB69Mi/7ReqqnrAXLZqAxXfZLZpFpe+0ijpVKhW4XC5Va5GOc1G4DBVOOrC/OFyPAb+zMot40ipdVygUZq5t6vf7SKVSCIVCqgscQghEUTTcGBTZLLPRaKBSqSx6OcrolXkYV4BvMpl0LyrWEiOIJpnjmrLbjd7pOy2EE8/z6Ha7qp7/CCEYDocL/4wuGiqcdGB/CNaoESeZacUTwzDgOE71YmCbzTbTcWo0GigUCkgmk5rUILXbbcNOYGcYRpl6v729vdCicTVGr+xP06n53EbBSKJJRs+UndqGqWqhZ/pO7XSdJEnI5XJIJBKqPJ/Mca9tkqHCSSf2p+qMGnGSmVY8qW1NUC6XEQqFpvodQgjy+Tz6/T42NjY0O8Z6ml7OSiAQwNraGlKp1MIiM2qMXun3+yPF76oUiBtRNAH6puyMHj3UK32nZtQpl8shGo2qPtFAy/FSywQVTjowKlVn5IiTzDTiyWQygWVZVULazWYTLpdrqi+9IAhIp9NwOp0Ih8Oa1cHIO/BFTpSfFKfTiUQigUwmM3fKbBbmjSTwPD+2NkPNrr1FYVTRJKNXym4RXk6zoHX6Ti3hVK/XYbVaVW9e6ff7x74oXMb4Z/8VYFSqzugRJ5lpxJMaUSdCCKrV6lT+SN1uF1tbW4jH43C73XO9/lEsW6jaYrFgc3MTpVIJjUZjIWuYVTyNS9OtAkYXTTJ6pOyWRTgB2qbv1EjXDQYDNJtN1QehA7QofDdUOOnEMkacZCYVTzabDcPhcK6TrNytNumuplKpoFarqWo1cBitVktzcaY2LMtifX0d/X4f+Xxe13qSWa0DZGfiw0SF2WxeSvuFZRFNgD4pu2USTjJape/miTrJUxHU7iCWn1sQBMN/XvWCCicdGBVxWoZUz24mFU+BQGDmji5JktButyfKoctWAwzDqDZ76Sj2GzEuE/JO2eFwIJPJ6ObTM6uDeK/Xg8PhOPQxy1jntEyiSUbrlN0yCicZtdN38winfD6PYDCoyaa8Xq8vVaRda5br6r3E7L/YLuPFdxLx5HQ60e12Z4pqlEqliQb5DgYDpNNpBAIBXUeeTJtCNCIejwfhcBipVEoXE8lZhdMkabpl66xbRtEko2XKbpmFE6Bu+m7WdF2r1QLDMJp1+zabTVoUvgsqnHTAqEN+Z2ES8eT1eqeupxFFEf1+/8hxKK1WC7lcDhsbG7o7dxvJLXweOI5DMplELpfT3OBvliJu2SfrqJ2z2WxWxWBTD5ZZNAHapuyW6X08DLXSd9NGnQRBQKVSmWvCwmH0+33YbLaVuYapARVOOmBEj5J5OEo8zSKcCoXCoV98QggKhQI6nQ6SyaTuxfWHdXgtI/J72Gq1VLWRUIPdA30nwejfr2UXTTJapeyM6uM0K/Om76YRToQQbG9va1quQIvCD0KFk06smlo/TDwxDDPVBPvBYABJksYKE/nCY7fbEY1GF3Isa7Xa0qfp9sMwDGKxGFiWxdbWlqbmftNEFKbxyTL63LpVEU0yx3WW3bTMk76bJl1XKpXg8/k0a4yRJAmCIBhuSsKiocJJB1ZpN7Wbw8RTIBCYOJJxWLSp1+shk8kgFostLMc+SYfXMuP3++H3+zUzy5xm9AohBJIkTRxRNHKB+KqJJoDOspuWWdN3k0Sdut0uBEHQ9LzYaDQMb/a7CKhw0gk5SrJqImqceDKZTDCbzUdGA3q9Hsxm88gdTbVaRaVSQTKZXOiFZ9rU0TLicDiQSCSwtbWl6ugc+bknjT5Oa/cwa/G51qyiaJKhs+ymZ9r03VHCSRRFFAoFxGIxtZe6h2azuXT2K3pAhZMO7C4OJ4QsnRXBUYwTT6FQCKVS6dDfLRaLBwb5ynl7QgjW19cXfrzq9frKGjHuxmKxIJlMKt5YajFNxGlag1GTyWS4tNEqiyYZmrKbnmnSd0el67Tya9qNXNe5amUmarBaV3CDsvvDv2zml5MySjxZLBZIkjS2vqXdbsPhcOxJywyHQ6RSKfh8PkMUJEqSNFXqaNmRzTIFQUAul1MlQjrpiXeecTZGieQeB9EEqJ+yY1n22IiwSdN346JOlUoFLpdL82YVWhQ+HiqcdEK+eCzTuJVpGSWexo1hIYSgXC7vGQ3QbreRzWaxvr5+pPmhXjSbzWNp/BYKheByuZBOp1VrFT9K3MzqFTNNREtLjotoklEzZbfsXk6zcFT6bpRw4nke3W5X8wi4vOGlReGjocJJB3an6lY14iSzXzzZ7XbwPH/goik70cqtyMViEa1WC8lk0lDH5zgbv7ndbkSjUaTTafT7/bmea5Lut1mPtRGMMI+baJJRK2V3HIUTcHj6TknXDbpA7WFI/Q5yuRwSiYTm65qms/U4QoWTDuwWDasccZLZL57W1tZQr9eV+wkhSreGKIrIZDKw2WyIxWKGyqfLo3GMtCa9sdlsSCaTyOfzaLVaMz/PUUXc8xzrRReIH1fRBKiXsjuuwklmZPpOEhH+6UeBvzgN+F9PBv7yNKw/+Cmw0D4tvYwzOfWECied2B1xWnXhBOwVTxzHodlsojcQkKl2sZUrIhAIoN/vI51OIxKJGDIdNsnYj+OA/F52Oh2USqWZ6omOEjfz7HAXWR9znEWTjBopu+MunGR2p+8a//R22O7+MBhh53vDCjws3/1fwB3XaboGWhR+NMbJiaww+yNORkpFaYl8wT31cAqf+XENt977M/BDCZyZwYX/JYELn+zD5ubmwrvmxtHpdGhx5CMwDINoNIp6vY7t7W3E4/Gp3rejRq90Op25DEYZhoEkSbp+lqhoepRIJIJUKqWIqGkxm83U3uARGIZB2OcG+dktYA5ElwjwvRuA57wFsGgz/qlSqRzodKbsxZhXrBXkuEWcZEwmE77yiz4+dXce/HAnKsALBB//zha+9J9dw4om2fCS7rr2sra2hkAggHQ6PfckeBn5OzHPsdY7XUdF017mTdnRiNNeSCsPZjjm8zzsAu2CJq9Li8Inw5hXrRVjd3G43rviRcMPRXzy26mR933y2w+DHxpzuGe1Wl25EStqYbfbsb6+ju3t7amKsseNXlEjJep0OnUrEKeiaTTzpOymHcuzqvA8j2w2i636EMQ8JqJkcQAubQb6HhfPunk5PlfwBbK/JuQ4RTFKrT56Y8RRdyCi1JqvW0sLCCHKRHDKaMxmMzY3N1Gr1VCtVif6nXG2Ab1eD3b7fGkHvSwJqGg6nFm77I7TOXE/oiiiXC4jlUqh0WggFAph47QzwTz9SgD7jwsD/MYVmqXpWq0WXC6XJs+9SlDhpBPH9cQQdFnBmUd/zBxWE0Ju44mTXq9nGB8pI8MwDBKJBCRJQjabPbJofJRtwGAwgMVimfv7IdtaaAkVTUdDZ9lNBiEErVYLmUwGuVwOdrsdyWQSkUjk0TTZ898JPPuPIMmRJ4sDePYf7dyuAb1eDxzHHdtr1TQcjyrlBWMUV2O9GQ6HKGS3cdF/SeDG72QO3H/ZM0+Csxiv3qtWq40dOkw5SDAYRLvdRjqdxvr6+tgaPo7jDpihqtm5KKd7tKghpKJpcux2OxqNxrGY8Tgt/X4f1WoVg8EAbrcbiURifOkGawJe+KfInvZ7WF+z7KTnNIo0ATvlCfS8NxlUOOnEcVPxrVYLlUoFiUQCF5stYE0sbvn+FroDEXYLi9c96zRc+6IzF73MAxBCjlXno1q4XC5YrVak02nEYjFwHHfgMaO+A/1+f+RjZ8HhcKDX66meaqCiaXrm7bJbJSRJQq1WQ7vdhs1mQyAQmOpzRMwc4NvQcIWPjpai573JoEdJB3YXh686hBAUCjsdH5ubm+j1ejCbWLzjJWfh2vMej0KjB6FdxWNOJBe80tFQ47fZsVqt2NzcxNbWFtbW1sa6gMvfB9kvRi2cTifq9bqqwomKptnYnbKLx+MT/96qnCsJIeh0OsqwbJ/PB7/fb9i/rVarUafwKTjeWwGdkFN1q56ykwf0Op1ORKNRAECxWFTCv5zFhM2gCxYDf+oajYYhzTiXBZZlsbGxgV6vh0KhcOAzv3v0itoGo1arde7RMLuhomk+pu2yM5lMSz/odzAYIJ/PK2OK4vE4NjY24HK5ZhJNsqO+1rTbbVoUPgU04qQTskHfqno47U7NycWNpVIJgUDgwBffZrMZsmtNboc+7qmFeZGjDY1GA1tbW3vqOGS/JVlAqSlI1NzNU9GkDtOk7GQvp2U7R0qShHq9jlarBYvFgkAgoNq5TRAEzT2VaFH49NArhA7I4edlPCkcBSEE+XwenU4Hm5ubypd8MBgoBZD78Xg8aDQaei/1SOT5eRR18Hq9CIVCSKVSSpTJbrej3q7jV+VfwWRT/7ughpEiFU3qMU2X3TKZYMqpuK2tLWxvb8NsNiOZTCIej6u6IdRjKHylUqETEqaERpx0QE5XrFrR8XA4xPb2NgKBwB6BRAg5dIo3x3EolUp6LXNiWq0Wkklj1l4tKxzHYWNjA1tbW/AH/fjUg5/Cpx/4NPpSH5yJw0VPuAhXPuVKmFh1RJRseTCuvuooqGhSn0m77JZBOA2HQ1SrVfA8D4fDgVgspulmeDgcanrNkCQJhJCVui7pAT1aOrFqEadRqTmZer0Oj8cz9svIMMxCZosdxnA4VMVPiHIQ2Szz+juvx+dPfV65nRd53Hj/jQCAN/76G1V5LafTiXK5PJNwoqJJOyZJ2ZnNZl3H5kwKIQSNRgPNZhMmkwmBQEC3tn1BEDS1dKhWq9QpfAaMcdVacVYp4jQuNScjCAKazeaRKS+Xy2WooZ70BKItfbGPr6a/OvK+z/zHZ8AL6rh+WyyWmWboUdGkLZOk7IwWcer1etje3kYmkwHDMNjY2EAikVDNPmMStE7VUa+t2aDCSSdWIeK0v2tuVHQmn88jFosdGblxu91oNptaLXVqeJ6fe+wHZTzlXhm8OFoc9YQeyr3yyPtmYdqoIRVN+nBUl53ZbFZtcPSsCIKAYrGIVCqFVquFSCSCZDIJr9e7kGi0lsKp2+3CbrfTKPsMLHf4Y0lYhYjTYak5mXa7DavVOtHFx2w2G2aoJ8/zuu4ijxuCIGBYH8LG2tCXDtoFcCwHO1FPtMoX4Em6kaho0pfDUnYsyy7EsoUQgmaziUajAZPJBJ/Ph3A4rPs6RqGlr1W1WlVsYyjTQSNOOrGsEaejUnMykiShXC4jFApN/NxG2GECOycQv9+/6GWsHIIgIJvN7pgghuO4+NcuHvm4i55wEaS+hFQqhWq1OreXj9PpnCgNTEWT/hhplh3P88hms8hkMiCEYH19HYlE4ljMqRRFkRaFzwE9ajpCCDFMMfQkjOuaG0WhUEA4HJ5qd+TxeNBsNhfaCksImTg6QZkMOd0hSRJCoZDSnn3lU64EANzywC3gJR52sx0XPv5CpauOEIJ2u42trS1YLBYEg8GZ3heHw4FisXhonR0VTYtjkbPsRFFErVZDp9MBx3EIhULH8ruvtvnscYMKJ8pIJknNyfA8D0LI1Ds1p9OJWq22UOHU6XSoY65KjBNMMibWhD94wh/gFfFXAE4gaA+CMz+aImUYBm63G263G/1+X3kuv98/1QX2qCJjKpoWj56z7GRBXq/XwTAMfD4fAoHAsa7t6XQ61LtpDqhwouxh/6y5o04ucipvFv8jhmFACFnofKparTbVLC3KQY4STLvp9XrwOr1HRjBtNhsSiQQkSUKlUkG5XIbb7cba2trEF9pRnysqmozBrLPspqHf76NarSpGvLsd7JcBreq9aFH4/FDhRFGYJjUnU6lU4Pf7Zz4hyRPtF1FXINfSLFvdmVGYRjDJ9Hq9qQpvWZZFKBSaOo1ns9kwGAz2rImKJmMxKmUnb6ZmvahLkoRarYZ2uw2bzYZAILC077Uoipqcm6rVKmKxmOrPe5ygwokCYLrUnMxwOESv10MwGJz5dT0eDyqVykKEEx3oOxuzCCaZWTtLR6XxRFFEIBAYmcZzOByoNqsQORFBexAWxkJFkwHZn7IzmUxTz2eTx5/UajUAgM/ng9/vX/qIihZz6uSicLpZnA8qnI4506bmdpPL5ebeuVit1oV11rVaLWxsbCzktZeReQSTmoxK47lcLvh8PrAsC1ES8YkHP4FbfnaLMtrlZfGX4S3PeAsVTQZjf8pOrk+bRDAMBgNUq1X0+324XC7E4/GVEgRaeDjRDmJ1oMJJJxZZxzOOWVJzMvV6HU6nU5UdEcuymoWlxyFbQxjtPTEiagkmtb8D49J4X9j+Am564CblcbzI44uZL8L7M69qo10o6rE7ZWcSRPCpFGwnToAd4a0mSRLq9TparRYsFgsCgcDCBLzWaDGnrtvtzpUhoOxAhZNO6C0MjmKW1JyMKIqo1+vY3NxUZS1utxutVuvIMS1qQttxj0btCFO/39fkIrc7jdfoNPCZf/3MyMd95j8+gz940h/s6eSjGINwMIhfvfd9EP/xH0F4HgWOg//SSxC6+mqAZdHtdlGr1UAIgdfrRTKZXPlNjyAIqk4z6HQ6cDgcK3/c9IAKJ50wivnlPKk5mUnHqkyK2+1GNpvVVTjRndd4tErJ9Xo9zR3aW1JrpDs5sDPa5f6H7kfUrq1bsjzEWv5+7P550vvm+V35vmWi/KEPQbj1VuVnwvOo/P1H0el0wF54IRwOB2KxmCHOoXqhdqquWq3SDmKVoMJJJ4wwbmWe1JxMp9OByWRSNXLAsuzcbtHTIEc+lu3iojVa1zBN21E3C0F7EJyJGzkXz26244mnPVHTiJPcQi7bbMj/Rt122H2SJM38u3qOLVFDIBKeR/VTN498/v4/fBln/tEfjUzbrTpqbrZFUQTDMMdKeGoJFU46seiI0zypORlCCEqlkmoput3YbDbNUjn7qVar1PxtF3oVfeuxebAwFrws/jJ8MfPFA/dd+PgLNU/TLWvEZ1bmFYiEEAzLZRB+9ABo0utBKJdhXV/X7W8yEmp9jmhRuLpQ4aQToigupIhRjdScTLFYRCgU0uSi4PF40Gg0NI9IEEIwGAxodxWM0yWnFrJP01ue8RZ4f+bFZ/7jM+gJPdjNdvzeGb+Hl/pfCkmSlsoE0ejsjirNimSzocxxI8UTY7fDTFPqc0EIQbfbnWqOKOVwqHDSCUEQdPcqUiM1J9Pv9yEIgmazpTiOQ6lU0uS5d9PtdnWfj2U0FiGYtO4q3W9u+cZffyP+4El/gHKvrIx26fV6SKfT2NjYoCkLA8FyHLwXXYj6jR8/cJ/ngguOZZpOTeg5T32ocNIJvWuc1EjNyRBCkMvlNPU8kneuWkcEarUaolFti4ONyiIjTFqmYcc5gnNmDuvuR1M8drsdsVgMmUwG6+vrC685pOwgiiJ6v/M78LMsap++BaTXA2O3w3fxRei+7GXgeV7zpoJVhhaFqw89c+iEXnYEaqbmZKrVKtbW1jRfv8vlQqfTmTs6Ng656Pa4XTCNkJLTqqNu2jEqsnmmLJ7UdmamTAchBNvb20hsbMB27bUI/eEfQiiXYQ4GwXIcJElCOp1GNBo9VuJJrQJ/WhSuDTTZrxN6GGAOh0OkUik4nU5Eo1FVXk8QBHQ6HV2sAtxuN5rNpmbP32w24fF4NHt+oyEIArLZLPL5PAKBANbX1xdWx9Tr9VT1pAFmnz1nsViwsbGBra0tDAYDVddEmY5CoQCfz6d8LlmOg3V9XUnPsSyLZDKJfD4PfkwB+SqilhWBPEuUoi5UOK0IrVZrZ+eWSKgasVFjrMqkmM1miKKo2fM3Go1jIZyMJJhk1E5Vzzuw12w2I5lMIpvNHqsLspFoNBpgWfbI89VxFE9qzKkjhCxsgPqqQ4XTkkMIQT6fR6fTwebmpqqph2azCY7jdE1nmM1mTWbXiaIIlmVXuqPKiIJJC+YVTTImkwnJZBKFQgG9Xk/FFVKOgud5NJvNiTu9jpt4UiPi1Ol0aFG4RqzuVeQYoEVqTkaSJFSrVd3dtT0ejybpunq9rqszuZ4YXTCpmaZWSzTJyBfkUqmETqejwgopRyGKInK5HBKJxFSfi+MkntSYU0fHSmkHFU5LilapOZl8Po9IJKK7kZ/T6US321X9edvt9srtvowumGTU6qhTWzTJMAyDjY0N1Go1tFot1Z6XchBCCLa2tpBIJGaK/h4X8TRvqk4QBFoUriHHq71oBdCia24/3W4XDMOoXsw7CQzDKI7Cav1tg8EAFotlZdycjdAlNw1qdNRpJZpkGIZBIpFANpuFKIorG51cNLLIn+c9lMXTKnfbzRtxotMRtIVGnHRALRGgZWpOhhCCYrGISCSi+nNPisPhULXmpFarrURnybJEmPYzb0ed1qJJhmEYxONx8DyParWq2escV2q1GsxmM1wu19zPteqRp3n87OSi8EVsfI8LhhZO119/PZ761KfC7XYjHA7j5S9/OX7+85/veQwhBO95z3sQj8dht9vxvOc9Dw888MCCVjwaNbyDtE7NyZRKJQQCgYUWUcvjV9Ri2Q30llUwyczTUaeXaJJhGAbRaBSCIOjiZH9c6PV66HQ6qo79WHXxNOvGuNPpqCJOKeMxtHD65je/iSuvvBLf+973cPvtt0MQBJx33nl7ijj/4i/+An/zN3+DD33oQ/jBD36AaDSKF73oRYaqVSCEzJxr1rJrbj+DwQD9fl9TYTYJVqtVtc66Zd55Lbtgmhe9RdNuwuEwWJZV0uKU2REEAYVCQRP36lUXT7NAi8K1x9A1Tv/yL/+y5+dPfvKTCIfDuOeee/Cc5zwHhBB88IMfxLve9S684hWvAADcfPPNiEQi+NznPoc3vOENi1j2AWaNOKk5a+4o5LEqiURC09eZFJZlVXFbr1armg8OVptlq2E6jFnT1IsUTTKBQAC1Wg25XE6z1PiqM28x+CQch5qnSZGLwlfZdsUILNXRldM3cr3KqVOnkM/ncd555ymPsdlseO5zn4vvfOc7Y5+n3++j2Wzu+acls0Sc9ErNydTrdXg8HsOMI3G73XNHDQkhEEVxacZqrGKEaZaOOiOIJhmfzweHw4Ht7W3VxmAcJ3K5HEKhkObfQRp52qFSqdCicB1YGuFECMGb3/xmPOtZz8JZZ50FYKdDA8CBQuZIJKLcN4rrr78eXq9X+afl8FpguoiTnqk5GUEQ0Gw2DdVJ5Ha70W6353qOdru9FLn+VRRMMtN21BlJNMl4vV6sra1ha2uLiqcpqFarsNlsutmArIp4kiRppugmIQQ8zy9tacIysTTC6aqrrsJPfvITfP7znz9w3/4P2VHpgXe84x1oNBrKv0wmo/p6969nkoiTHl1zo8jn84jFYoZKRbAsC0mS5noOo5terrJgkpmmxsyIoknG5XIhEAggnU7P/bk8DnS7XfR6Pd2jH6sgnmZ1DV+WjeIqYIy8zBFcffXV+NrXvoZvfetbWF9fV26PRqMAHr3wyxzVTm+z2XS9QE0ScWq1WqhUKkgkErqmltrtNqxWq+EuVMDO+zSreaJ8cTNirn+VapiOYtKOOiOLJhmHw4FwOIx0Oo2NjQ1qLjiG4XCIYrGIzc3Nhbz+stc8zWp+WavV9lwfKdphvKvKLgghuOqqq/CVr3wF//Zv/4aTJ0/uuf/kyZOIRqO4/fbbldsGgwG++c1v4hnPeIbeyx0LIWTsBXwRqTkZSZJQLpdVbRFWk3lsCRqNBrxer8ormo/jEGGahWUQTTJ2ux2xWAyZTAaCICx6OYaDEILt7W2sr68vNIK9zJGnWcwvBUGAyWQy5EZxFTF0xOnKK6/E5z73OfzTP/0T3G63Urfk9Xpht9vBMAyuueYaXHfddTjjjDNwxhln4LrrroPD4cBrXvOaBa9+L6NOInp2zY2iUCggHA4bKkW3G47jZvbSabVamteuTcpxijDtZpKOumUSTTI2mw2JRAKZTAbr6+tL03ygB9vb2wiHw4ZoMlnWyJMgCFPXKdGicH1Z/Kf7ED7ykY8AAJ73vOftuf2Tn/wkLr30UgDAW9/6VvR6PVxxxRWo1Wp42tOehttuu23hXkRHsajUnAzP8yCEwOFw6P7ak8IwDBiGmdpFV959LVoQHlfBJHNUmnUZRZOMxWLBxsYGMpnMUq5fC8rlMhwOh6HOKcsonqatcZKLwhc57eG4wRDaJoJmswmv14tGowGPx6Pqc/NDEd/50U/xjF8/C5zFtGfW3CKG6AI7X7RUlc5K5AAALAZJREFUKoVkMmn40K48pmEaIVwsFuF2uxfWXXLcBZPMYe/dMoum3YiiiEwmszQXZa1ot9toNpuamFyqgSRJSyOetra2EI/HJz43t1otDIfDlRgrtUim0QGGjjgtM6JE8IHbH8RNd51CbyjC/s85XPqMJP77Y+0Ih4ILjYhVKhX4/X7DiyZgx5agUChMdbx6vd5CTC+pYNrLuPdhVUQTAJhMJiSTSWQyGYTD4WPZCj4cDlGpVJBMJhe9lLEsU+TpsJrYUdCicP0x/pVzSfnA7Q/iQ3f8Er2hCADoDUV85Jun8A8/7y1UNA2HQ/R6PdUja1phNpshiuLEj+d5XnfBQou+RzOqo26VRJOMfFEulUp7xkEdByRJUox6F50aP4plLhgfx3A4pEXhC4AebQ3ghyJuuuvUyPtu/m4a/HByIaA28viIZcJsNk88u65Wq+kWsqaCaTpWUTTJMAyDjY0N1Go1Q83J1BK5gy4SiRiiGHwSVk080aLwxUCFkwaUWn0l0rSf7kDEA7/KoF6vo9/v6+pE3Gg04HQ6l64LyOPxTDQWhxCCwWCg+UWZCqaj2d9Rt8qiSYZhGCQSCTSbTdTr9UUvR3PK5fJCawlnZVXEEyEE/X7f0GnHVWU5tglLRshtg91iGimeHFYTzkxGQYQB6vU6BoOBcp/NZoPdbofdbld9ByeKImq12sJM6ebB6XSiVqsdubPqdDqajnegNUyTs7uj7jiIJhmGYRCPx1EoFFCtVle2YLfVakEQBMN6wB3FMtU8jaPVahm+e3xVocJJAziLCa971kl86I5fHrjvsmeehNvBAeD21BnJu4der4disajU9TAMo4gpjuOmzmVLPA+hXEZxMEA0Hjd8HcIoGIYBIeRIX6B6vb7HQV4tqGCaDn4o4pe5GiJe+7ESTTIMwyAajaJYLKJUKi2tuBjHYDBAtVo1dDH4JBhRPImiOPE5vl6v06LwBUGFk0Zc+6IzAQCf+PYpdAci7BYWr3vWacrt+2EYBhzHgeM4+Hw+5XZJktDr9dDtdlGtVpXUntlsVgSV1Wo9OK9PFFH6u79D9VM3g/A8YLPB+tpLYbv6ajBLOCrC4XCg1+uN9YiRJAmSJKk6BoMKpunY30nKWVj897P8eNfvPOXYiKbdhMNhVCoVFAqFlfHYkYvBk8nkUm7C9mM08TTpuBVaFL5YqI8TtPdxyte7kLp1nLapnpO13B3X6/VGpvs6N96I2sduPPB7gcvfgPA116i2Dr0YDAaoVCpjI0r1eh0Mw6gyZoUKptn4q//v5yOjrFc9/3S85fzHLmBFxqBWq4HneV0Hd2sBIUSxXVi0wFAbo/g8dTod9Pv9I1O8+Xwea2trK/c+LJJpdACVqxrDWUw4EXLDzKirTy0WCzweDyKRCDY2NrCxsYH19XV4PB4M2x3Ubv70yN+r3vxpSEtYEGm1Wg/trGs2m3OLXlr0PT2SJKHf76NUrePjdz408jGf+PaphXaSLhqfzweHw4Ht7W1dm0HUplgswuv1ruTF2igF45PMqaNF4YuHpup0YpbRIbO8BsdxYEUBhX5/5GNIr4fMT34Cy/o6nE4nnE7nUrUSi6J4IB0n1wXMupunEaaDEEIgiiKGwyEGg4Hy3/2eWizLwmKxoNgRwQvSyOfqDkT85BcP48z4zkxGNdOpy4LX64XJZMLW1tbCB+DOgtzVarTB2WpihLSdIAhHNriosUmkzMdyXDFXALlGR8uuLxlzMAiG43Zqm/bB2O3YeNKTIJpM6Ha7SiE6IQRWqxVOpxMOh8OQFzePx4NWq4W1tbU9t9dqtT11YZNyXAUTIQTD4XCPKBoOh5CkvcLHZDLBarXCYrHA6XRibW1t7AxAz1AEZ2ZHiifOzOCxyRgYSUA+n1eMMT0eD5xO59KJiFlxuVxgWRbpdBobGxtLU5/S7/dRr9cNMzRbSxYtniaZU9doNI7Fe2FkqHDSCafTqfgoac0AgPm//TcMP//5A/f5L7kYLMeBxc7uUd5ByhfTTqeDfD6vXERtNhucTifsdvvCT/QulwvZbPaAcOp0OlOZwK2yYJIk6UCUaFSK02KxKKLI4XDAYrHM/P4SQlAtFXDBuTHc9L3tA/e/9pkn0KiW4XK5EH+ks3M4HKLVaqFWq4EQAo7j4Ha7wXHcSgsph8OBcDisiCcjblB2I4oistksNjc3V/p92c0ixdNRwklO5R2X98KoUOGkEzabDf0x6TM1GQ6HyOfzOPnOd6Di9aB686dBej0wdjv8l1yM0NVXj/w9hmFgtVphtVqV6I2cS+90OsoFDgDsdjscDgfsdruuX2CWZQ9ERWTDy0nWscyCadrUmcVigc1mg8vlgsVi0ex9EkURW1tbCAQCeOfLYuDsHG668yHwAoHDasJlzzyJa190Jkwsg1qthnQ6jXg8DovFAr/fD7/fr3zOms0mSqUSGIaB0+mE2+1eOrPWSbDb7YjFYshkMlhfXzdsqlx2Bp9m4OyqsCjxdJTlSrlcpk7hBoB21UHbrrrdpNNpTb1P5M6Q3Sdj2cfJHAyCVeHLTwhR7BF6vR6AR72mnE4nbDabpmKqUChgbW1NET35fB4+n+9QEWR0wTRL6kz+77jUmR7wPI9cLrfHo6nb7aJSbwGcGyG3DZxlb0RlOBwil8vB5XLB5/ONXDshBJ1OB81mE4IgwGQywe12K6muVWE4HCo1T0YUiPl8XhGwxxW9u+0ymczYNBwhBOl0eilNjJeBaXSAMbc6K4rFYtFsJAghBFtbW4hGo3t2sCzHwaqiSRrDMHA4HHv8lGSvqWazqUTVWJaFw+FQRryodXH3eDxoNBoIh8NKpGKcEDKCYJokdcYwDMxms2qpMz1oNBpoNBrY3Nzcs85ms4lwYLyQtVgsSCaTB6JPu2EYBi6XCy6XC8BOVKvVaiGbzUKSJFitVng8Ht0jnmpjsViwsbGBTCZjOIPQRqMBlmWPtWgCFl/ztBtaFG4cqHDSEZfLhU6no8kJMpfLwe/3L+SLzbKs0qEnI4qiYtopCwWTyTR3Jx/HcSiVSgAw1hBTD8Fk1NSZ1hBCUCwWAQAbGxsH/o7BYDDR8fb5fHC5XEdGn4Cdz83a2ppS2zYYDNBsNlGpVEAIgcPhgNvtNlwkcRLMZjOSySQymczCL8wyPM+j2WxSV+pHMIp4okXhxoEKJx1xOBxKaklNSqUSOI5TduhGQE6v7N6xDodDdLtdlEolCIIwUycfwzBgGAbd/hAPPJzH408mlPvUEkxadJ2tApIkYWtra09Twf77p4mSTRJ9GoXVakUwGATwaOq4VqthMBiAZVm4XK6lsj0wmUyKeAqHwwsdmiuKInK53LEqBp8EPcTTYVUzg8GAFoUbCCqcdMRkMh2ISMxLo9GAKIpLMQ/LYrHM3cknSgQ3/6iKz/zwp+CHEuyWB/HaZ2zi1U/0gAE5UjCtaupMawaDAba3txGLxcZeNNrt9kzifZro0372p44lSUK73V462wP5wpzJZBAIBHTpvt2PnO5PJBLH+rM+Dq3F0yiPOplKpUKLwg0EFU46o6YRZrfbRavVQiKROPrBBmSWTr4P35nGx7+7pTxHbyjihm8+BEk6gbec/1gMh0M0Go1jkzrTg3a7jUqlgmQyeWgUp9VqIRqNzvQas0af9sOyLDwej1ILsky2BwzDYGNjA9vb25AkSff6Itkx30i1VkZDS/E0bk4dIWTiFDhFH6hw0hm1jDAHgwGKxeLKhdR3DzuWkdMx1UZ77FiPT303jVc90Qu3gzs2qTM9qFQq6Pf7Ew11PWzHPCnzRJ9GsWy2BwzDIJFIIJvNQhTFA55lWlGr1WA2mw2V7jcqWomncR5OjUZjpR3blxEqnHRGDSNMURRXakL5OOQC806ng8FggFxriL4wug6AFySYXX4E/QeLxSnTQwhBNpuF3W5HPB4/8vHD4VA18aFW9Gk/+0W5bHsg19wZxfaAYRjE43EUCgVUq9UjB77OS6/XQ6fTocXgU6CFeBo3p67ZbNKicINBhZPOzGuEKU8oTyQSS1P8OgmEEPA8j06ng16vB0KI0q3n9/thsVgQFiTYLb9Ab8TAWIfVhJCbhrLVQBAEbG1tIRQKTSzwtWiVVjv6tB8j2x4wDINoNIpisYhSqaRZDaMgCCgUCpr6y60qaosnQRAONAYMBgNaSmBAqHDSmXm+AHLxZjgcXvo6BLkovNPpQBRFJRrgdDoRCARGHifOYsLrnnUSH7rjlwfuu+yZJw+YLVKmp9froVAoIJFITBXl6Xa7mkRGtIo+jcKItgfhcBiVSgWFQgGRSETV56bF4POjpngalaqjReHGhAqnBTCrEWahUIDH4xnpXWRk9qfc5K41p9OJaDQ6VeTs2hedCQD4xLdPoTsQYbeY8MonBZTbKbNTr9fRarWQTCanupASQo4cFTEvWkefRmEU24NAIIBarYZcLodoNKra353NZhEKhQxT37WsqCWe5HSxjNx1vOyb5FWEjlyBfiNXZFqtFgRBmMrPqVqtLoXtwGEpNzVdxPmhiFKrj5Dbhm6rAQCa14KsKoQQFAoFsCyLUCg09fsji2K9Ppu1Wg3NZlPT6NNRyLYHrVZLN9uDRqOhdNHO+xrVahWEEBrNUJF5x7PsH7dSr9fBMAwtDNcJOnLF4ExrhNlut8Hz/ERFunozS8pNDTiLCRuPFIJzfj8ymYwyK48yOXKjwdra2sybhmazqbqp62EsIvq0n0XYHni9XphMJmW+3azPK8+ZXFYbE6Oids1To9GgtWcGhQqnBTCNESbP86hWq4boqlAz5aY28Xgc6XQaJ06coIWUEyKbWsbj8bkE5yI8ZvSsfZp0PXrYHsgdf+l0GhsbG1PXJg2HQ8XGhKI+aomnfr8Pq9VKz2UGhQqnBTGJEeZwOEQ+n1+I7cAkXW5G+lKbTCaEw2Hk83nEYrFFL8fwtFotVKvVI00tj0ItM9dZMUL0aT9a2x44HA6Ew2FFPE36/hFCsL29PVe0inI0aoinSqWi1NdRjAcVTgviKCNMSZKUk5weF6ZFpdzUxOl0KnUnx32q+2HIF3A1BPmsY1bUxGjRp/1oYXtgt9sRi8WUuphJxNP29jbC4fDMA7YpkzOteNrdXEEIgSAItCjcwNBv0II4zAhTbhOORqOanOSMnHKbl3A4jFQqBbvdTi8Q+5AjDk6nU7VC7nnGrKiNEaNPo1DL9sBmsyGRSCCdTmN9ff1QsVgul/fM9KNozzTiaXdHXb1epwXhBodeWRbEYUaYuVwOfr9fFTfaZUu5zYvsupzNZrGxsbFSf9s8DIdDJeKg5sVTjTEramL06NMo5rE9sFgs2NjYUExxR0Up2u02BoOBIZtLVp1JxdPuOXXNZpMWhRscKpwWxLgLeqlUAsdxM6c/ViHlNi9y6oPWCezQ7XZRLBanNrU8CjXHrKjNskSf9sMwzJ7IkGx7kM/nx9oemM1mJJNJZDKZRy/Owx7QLmBg9aFcrtFi8AUyiXiSzS9pUfhyQIXTArFYLGh2emj0CUJuG/rd9lReTauccpuXtbU1bG1tged5VSeYLxu1Wg2dTkeTYdBajFlRk2WMPu1nGtuDZDKJTOphxB+8GZYffgwYdmE2cUj+xhVgkn8CMMf3fLBojhJPw+EQNpsNlUrF8F59FCqcFoYoEXzihyXc8v37wAsSOAuL331yEO/+b+eMfPxxS7mpQTweRyqVwubm5rEbKUEIQT6fh9ls1mx4q1ZjVtRmWaNPozjK9iD0wMdh/uGHlcezIg98+wMAywIv/NMFrpxymHgSBAEOh2NPyo5iXKhwWhAfuP1B3PjtjPIzP5Tw6R8W4XH/Am85/7E05aYCLMsiGo0il8sdK7M/URSxtbUFv9+vWXehHmNW1GQVok/7OWB7MOgC930CB98RAum7H0b2Ma8Ga3WAYRiwLKv8m/RnhmGW5v02KiPF07AHqXIKHbKuNAxQjA0VTguAH4q46a5TI+/7+J2/wsvP5ODkrDTlpgJ2ux0WiwWNRuNYdKr0+31ks9m5TS2PotfrHZjkvgysUvRpP0ynCAi9kfexQg8JjwmSNwJJkiBJEgghyv/LPw+HwwO37f7/qde0S3RNI9Z2C7VVeX9kFPH08Cms//LTMP/gY4gOu5BMHJhnXAU8/50AS8/5RoYKpwVQavXRG452DucFAps3hJiftg2rRSgUQjqdhsPhWPoow2E0m03U6/W5TS0nfS09x6yoySpGnwAArghgcQDD7sH7LA4w7p1NmJ4bsVHi7DCxNur+aVkGscayLDYf+izwnQ8+epvIA3f+9c4PNK1qaKhwWgAhtw12i2mkeOLMDAbNMuqsR5fJ68cBhmGQSCSwvb29EBd2rSGEoFQqQZIk3SwYFjFmRW1WLvpksQO/cQXInX+1L13HAL9xxc79OsMwDBVro34W+7B+/4aRaVV87wbgOW9ZyPtFmQwqnBYAZzHhdc86iQ/d8csD9/3+sx+Dk8n1Ay3I8oiGpT6xLxCz2Qyfz4dSqYRwOLzo5aiG7DDvdrt1q49Y9JgVNVm16BP/9DdD6LThuv/TO5Eni2NHND3/nYtemm4YXayJogipcgrR4ei0KoZdoF0AfCd0Wz9lOqhwWhDXvuhMAMAnvn0K3YEIh9WEy555Ete+6EywLDOyBXlrawuEENhsNng8HlUnrx8HPB4P2u02ut3uSjgoD4dDxWFez3ojI4xZUZtViT6VqzVEX/w+4Lfes3PxdUVo5EIHDhNrgiAo3dDycHebNwZitoMZVZNmcey8bxTDwpBZ4pIrRrPZhNfrRaPR0N2Xhh+KKLX6CLlt4CxH75AIIcqIhl6vpxjmud1uOttoAiRJWgmLAnlo7Pr6uu6jZba3t1e6aaFWq6HZbC5d9EkQBORyOWxsbCx6KccWSZLQ6/XQ6XTQ7/dBCFG89ZxO597v6v/9fx+padp9CWaAZ/8RrXFaANPoABpxWjCcxYSNKQrBGYaBzWZTTNIIIeh2u6hUKhgOh0eOaDjusCyLWCyGbDarmb+R1lSrVfR6PU1MLSfBaGNW1GZZo0/UKV9f5E1sp9NBt9tVUth2ux0ejwc2m+3wz42cPv3eDcc2rbqs0IgTFhtxUht5REOr1dpTH+V0Opc6wqI25XIZZrN5qXxTCCHI5XKw2WwIBAILWcNwOES5XEYsFlvI6+vNskSfCCFIp9N0tIqGCIKgTGoYDocAdmaOOp1OOByO2c+vj4zHoWnVxUIjTseYcSMatre3aX3ULgKBADKZDBwOx1KkOGVTy0AgsND6IqOPWVGbZYk+1ev1pdoEGB152HKn0wHP8yCEwGQywel0IhgMqiuiLXZaCL5kUOG04uwf0SDXR8kjGo5rfRTDMIjH49ja2lpYymtSeJ5X3M8X/T4ty5gVNVmGzrtms4lkMrnoZSwlcsebPKlBTsLY7Xa4XC4Eg0FDnx8o+kOF0zGC1kftxWw2IxgMolAoIBqNLno5I2k0Gmg0GoYoZl+2MStqY9ToU6fTgcPhMMRaloHdw9GHwyEIIbBadyY1xGKxY3Huo8wHFU7HGIZhlG4P4NH6qP3+UatcH+VyudBqtdDpdJTjYAQIISgWiwCgm6nlUSzrmBU1MWL0qVKpHKtZjNMwaji6yWSCw+GA3+9feASXspxQ4URROK71UdFoFKlUChzHGWK3KUkStra24PV6DTVfb5nHrKiNUaJPg8FAd7NHI7O/yw0AHY5OUR0qnChjOS71UQzDKBYFi/bAGQwG2N7eRiwWU6beG4VVGLOiJkaIPpXLZSX1ftyQJEmpSxoMBgB23hOXy7XSPmOUxUOFE2UiVr0+Sm4rrlarCyt+brfbqFQqugzpnZZVGrOiNouKPkmSBEEQln7jMgmjUm4sy8LpdMLv98NisdBoEkU3qHCizMQq1kf5/X5kMhk4nU7dIyuVSgX9ft+wQ4hXccyKmiwi+rRIka81cpdbt9tVxpTQlBvFKFDhRFGFVamPisfjyGQyulkUEEKQzWZht9sRj8c1f71ZabVahu08NBJ6RZ8IIWi32yvhFC5JktLl1u/3AewIUafTiXA4rPtIIQrlKOgnkqIJy1ofZTKZEAqFkM/nNXfHFgQBW1tbCIVChuroG8Wqj1lREz2iT61WaymNSAkh6Pf7SspNTgE7HA6sra3BarUaemNFoQBUOFF0YNnqo5xOpzK2xu12a/IavV4PhUIBiURi4e3sRzEcDg2/RiOiZfSpVqstheGlIAhKAbeccpPrCX0+39Kk8SmU3VDhRNGdZaiPCofDSKVSsNvtqqcK6vU6Wq0WksnkUlw4jtuYFTXRIvrE8/zRA2QXgCRJypiSfr8PQgjMZjNNuVFWDvpJpiwcI9ZHMQyDRCKhWBSo8bqEEBQKBbAsi/X1dcNd+MZxHMesqI2a0adyubzwejM5/b7bM4llWdjtdng8HkMKOwpFLahwohiO/fVR/X4frVYLpVIJwE4qTY/6KIvFAo/Hg0qlMncRriiK2N7extra2lJFb477mBU1USP6JAgCAOgevREEYc+YEuDRlNva2tpSRE4pFLWgwoliaBiGAcdxihmk3vVRa2tr2NraAs/zMxtSyqaW8Xh86Qwk6ZgV9Zkn+lSpVBAIBDRdHyFESbnxPK+MKXE6nQgGg7TejXLsocKJslQsoj4qHo8jlUrNNGi31WqhWq0a0tRyEuiYFW2YJfokm0BGIhHV1kEIUTyTOp0OCCEAALvdDpfLhWAwSKONFMo+qHCiLDV61EexLItoNIpcLjfVMNVSqQRBEAxrajkJdMyKtkwTfarX61hbW5vr9URR3JNyI4TAarXC6XQiFostpbinUPSGCifKSqFVfZTdbofVakWj0Thy8C4hBNvb23A6nUs9R4yOWdGHo6JPEs9DKJfR6HaxecYZEz/vqDElJpMJDocDfr/fcB5qFMqyQIUTZWVRuz4qGAwinU7D4XCMTasMh0Nsb28jHA7D4XCo+vfoDR2zoi/7o09rHg/KH/oQqp+6GYTnAY5D6dJLELr6ajAjPq/7u9wAOqaEQtECKpwox4Z566Nki4Lt7W0kk0n0BQmlVh8htw2cxYRut4tisbgUppaTQMes6M/u6NOv3vc+CJ+/9dE7eR6Vv/8oACD4xjcqdUmDwUD5XZfLhWg0SlNuFIqGUOFEObbMUh9lNpvh8a7hz//xXtz6oyJ6QxF2iwkXnBvDRWf7dZtxpwd0zMri8NrtKHz1H0feV/nUzej/9m/D9UhK2mKxrMxnjkJZBqhwolAeYdL6qJu+n8cnv59Tfq83FPHx726B4zi8JbkaFzA6ZmV+CCGQJAmiKI78t/u+/Ui53E56bhQ8j4jNBivtdqRQFgIVThTKCMbVR20XSvj4nb8a+Tuf+PYpXPWC08FZlj9Kc5zHrBBCDgibSQXPbhiGAcuyMJlMe/5ZrdY9P7MseyBiJIVCeJDjRoonxm6HeU5DVgqFMjtUOFEoEyDXR5n7DHiBjHxMdyDixw8+jLhnb7cSy7Iwm80wmUwj/2u0zjVe4PFQ5SGcdfKsRS9lKowgeNSC5Tj4L71EqWnajf+Si8HOaMZKoVDmhwonCmUKQm4b7BYTesODF1+H1YQnn3liT8Rpd7pGEATlv/1+X/lZ7oDaD8MwY8WWfPFW88ItSiI+fN+HccvPbgEv8uDu5nDREy7ClU+5EiZWuyjaKgkeNQldfTUAoHrzp0F6PTB2O/yXXKzcTqFQFgMVThTKFHAWE173rJP40B2/PHDfZc88eSBNJ4sf+QI+DbKg2C24eJ7fc9s4Zolyffi+D+PG+29UfuZFXvn5jb/+xrHro4JHGxiTCeFrrkHw8sshlMswB4M00kShGAAqnCiUKbn2RWcC2Klp6g5EOKwmXPbMk8rtasEwDMxm89QDXWeJcvXFPj79wKdHPt8tD9yCF/tfDJtpr4M4FTz6wHIcrOvri14GhUJ5BCqcKJQpMbEM3nL+Y3HVC07f4+NkFGaJcm21ttCX+iPv4yUeXIDDuptevCkUCsVYVakUyhLBWUzY8DsMJZpmJWgPgjONTgPZzXYE7bSLi0KhUAAqnCgUCgDOvFMIPooLH38hODOtraFQKBSApuooFMojXPmUKwEAn/mPz6An9GA323Hh4y9UbqdQKBQKwBBCRpvSHCOazSa8Xi8ajcaxNf2jUGR4gUe5V95J39FIE4VCOQZMowNoxIlCoeyBM9NCcAqFQhkHrXGiUCgUCoVCmRAqnCgUCoVCoVAmhAonCoVCoVAolAmhwolCoVAoFAplQqhwolAoFAqFQpkQKpwoFAqFQqFQJoQKJwqFQqFQKJQJocKJQqFQKBQKZUJWRjjdcMMNOHnyJDiOwznnnIM777xz0UuiUCgUCoWyYqyEcPrCF76Aa665Bu9617tw77334tnPfjZ+67d+C+l0etFLo1AoFAqFskKsxKy6pz3tafj1X/91fOQjH1Fue/zjH4+Xv/zluP7664/8fTqrjkKhUCiU48s0OmDpI06DwQD33HMPzjvvvD23n3feefjOd74z8nf6/T6azeaefxQKhUKhUChHsfTCqVwuQxRFRCKRPbdHIhHk8/mRv3P99dfD6/Uq/zY2NvRYKoVCoVAolCVn6YWTDMMwe34mhBy4TeYd73gHGo2G8i+TyeixRAqFQqFQKEuOedELmJdgMAiTyXQgulQsFg9EoWRsNhtsNpsey6NQKBQKhbJCLH3EyWq14pxzzsHtt9++5/bbb78dz3jGMxa0KgqFQqFQKKvI0kecAODNb34zLrroIpx77rl4+tOfjo997GNIp9O4/PLLJ/p9ubGQFolTKBQKhXL8kK//kxgNrIRwetWrXoVKpYI///M/Ry6Xw1lnnYVvfOMb2NzcnOj3W60WANAicQqFQqFQjjGtVgter/fQx6yEj9O8SJKEbDYLt9s9tqB8WprNJjY2NpDJZKg31AKgx39x0GO/OOixXxz02C8ONY49IQStVgvxeBwse3gV00pEnOaFZVmsr69r8twej4d+iRYIPf6Lgx77xUGP/eKgx35xzHvsj4o0ySx9cTiFQqFQKBSKXlDhRKFQKBQKhTIhVDhphM1mw5/92Z9Rv6gFQY//4qDHfnHQY7846LFfHHofe1ocTqFQKBQKhTIhNOJEoVAoFAqFMiFUOFEoFAqFQqFMCBVOFAqFQqFQKBNChROFQqFQKBTKhFDhpBE33HADTp48CY7jcM455+DOO+9c9JJWjuuvvx5PfepT4Xa7EQ6H8fKXvxw///nP9zyGEIL3vOc9iMfjsNvteN7znocHHnhgQSteTa6//nowDINrrrlGuY0ed23Z3t7GhRdeiEAgAIfDgac85Sm45557lPvp8dcGQRDwJ3/yJzh58iTsdjtOO+00/Pmf/zkkSVIeQ4+9OnzrW9/Cb//2byMej4NhGPzjP/7jnvsnOc79fh9XX301gsEgnE4nXvayl2Fra2v+xRGK6tx6663EYrGQG2+8kfzsZz8jb3rTm4jT6SSpVGrRS1spzj//fPLJT36S/PSnPyX33XcfeclLXkKSySRpt9vKY97//vcTt9tNvvzlL5P777+fvOpVryKxWIw0m80Frnx1uPvuu8mJEyfIk570JPKmN71JuZ0ed+2oVqtkc3OTXHrppeT73/8+OXXqFPnXf/1X8stf/lJ5DD3+2vDe976XBAIB8s///M/k1KlT5Etf+hJxuVzkgx/8oPIYeuzV4Rvf+AZ517veRb785S8TAOSrX/3qnvsnOc6XX345SSQS5Pbbbyc/+tGPyPOf/3zy5Cc/mQiCMNfaqHDSgP/yX/4Lufzyy/fc9rjHPY68/e1vX9CKjgfFYpEAIN/85jcJIYRIkkSi0Sh5//vfrzyG53ni9XrJ3//93y9qmStDq9UiZ5xxBrn99tvJc5/7XEU40eOuLW9729vIs571rLH30+OvHS95yUvIZZddtue2V7ziFeTCCy8khNBjrxX7hdMkx7lerxOLxUJuvfVW5THb29uEZVnyL//yL3Oth6bqVGYwGOCee+7Beeedt+f28847D9/5zncWtKrjQaPRAAD4/X4AwKlTp5DP5/e8FzabDc997nPpe6ECV155JV7ykpfgN3/zN/fcTo+7tnzta1/Dueeei1e+8pUIh8M4++yzceONNyr30+OvHc961rPwf//v/8WDDz4IAPjxj3+Mu+66Cy9+8YsB0GOvF5Mc53vuuQfD4XDPY+LxOM4666y53ws65FdlyuUyRFFEJBLZc3skEkE+n1/QqlYfQgje/OY341nPehbOOussAFCO96j3IpVK6b7GVeLWW2/Fj370I/zgBz84cB897try0EMP4SMf+Qje/OY3453vfCfuvvtuvPGNb4TNZsPFF19Mj7+GvO1tb0Oj0cDjHvc4mEwmiKKI973vfXj1q18NgH729WKS45zP52G1WuHz+Q48Zt5rMRVOGsEwzJ6fCSEHbqOox1VXXYWf/OQnuOuuuw7cR98LdclkMnjTm96E2267DRzHjX0cPe7aIEkSzj33XFx33XUAgLPPPhsPPPAAPvKRj+Diiy9WHkePv/p84QtfwGc+8xl87nOfw6/92q/hvvvuwzXXXIN4PI5LLrlEeRw99vowy3FW472gqTqVCQaDMJlMBxRtsVg8oI4p6nD11Vfja1/7Gu644w6sr68rt0ejUQCg74XK3HPPPSgWizjnnHNgNpthNpvxzW9+E3/7t38Ls9msHFt63LUhFovhCU94wp7bHv/4xyOdTgOgn3st+eM//mO8/e1vx+/93u/hiU98Ii666CJce+21uP766wHQY68XkxznaDSKwWCAWq029jGzQoWTylitVpxzzjm4/fbb99x+++234xnPeMaCVrWaEEJw1VVX4Stf+Qr+7d/+DSdPntxz/8mTJxGNRve8F4PBAN/85jfpezEHL3zhC3H//ffjvvvuU/6de+65uOCCC3DffffhtNNOo8ddQ575zGcesN148MEHsbm5CYB+7rWk2+2CZfdeNk0mk2JHQI+9PkxynM855xxYLJY9j8nlcvjpT386/3sxV2k5ZSSyHcFNN91Efvazn5FrrrmGOJ1O8vDDDy96aSvFH/7hHxKv10v+/d//neRyOeVft9tVHvP+97+feL1e8pWvfIXcf//95NWvfjVtDdaA3V11hNDjriV33303MZvN5H3vex/5xS9+QT772c8Sh8NBPvOZzyiPocdfGy655BKSSCQUO4KvfOUrJBgMkre+9a3KY+ixV4dWq0Xuvfdecu+99xIA5G/+5m/Ivffeq9j6THKcL7/8crK+vk7+9V//lfzoRz8iL3jBC6gdgZH58Ic/TDY3N4nVaiW//uu/rrTIU9QDwMh/n/zkJ5XHSJJE/uzP/oxEo1Fis9nIc57zHHL//fcvbtEryn7hRI+7tvzv//2/yVlnnUVsNht53OMeRz72sY/tuZ8ef21oNpvkTW96E0kmk4TjOHLaaaeRd73rXaTf7yuPocdeHe64446R5/dLLrmEEDLZce71euSqq64ifr+f2O128tKXvpSk0+m518YQQsh8MSsKhUKhUCiU4wGtcaJQKBQKhUKZECqcKBQKhUKhUCaECicKhUKhUCiUCaHCiUKhUCgUCmVCqHCiUCgUCoVCmRAqnCgUCoVCoVAmhAonCoVCoVAolAmhwolCoVAoFAplQqhwolAox4KbbroJ5513nuav88///M84++yzlfllFApltaDCiUKhrDz9fh/vfve78ad/+qeav9ZLX/pSMAyDz33uc5q/FoVC0R8qnCgUysrz5S9/GS6XC89+9rN1eb3Xvva1+Lu/+ztdXotCoegLFU4UCmVpKJVKiEajuO6665Tbvv/978NqteK2224b+3u33norXvayl+257XnPex6uueaaPbe9/OUvx6WXXqr8fOLECbz3ve/FxRdfDJfLhc3NTfzTP/0TSqUSfud3fgculwtPfOIT8cMf/nDP87zsZS/D3XffjYceemj2P5ZCoRgSKpwoFMrSEAqF8IlPfALvec978MMf/hDtdhsXXnghrrjiikPrl+68806ce+65M73mBz7wATzzmc/Evffei5e85CW46KKLcPHFF+PCCy/Ej370I5x++um4+OKLsXte+ubmJsLhMO68886ZXpNCoRgXKpwoFMpS8eIXvxivf/3rccEFF+Dyyy8Hx3F4//vfP/bx9Xod9Xod8Xh85td7wxvegDPOOAPvfve70Wq18NSnPhWvfOUrceaZZ+Jtb3sb/uM//gOFQmHP7yUSCTz88MMzvSaFQjEuVDhRKJSl46/+6q8gCAK++MUv4rOf/Sw4jhv72F6vBwCHPuYwnvSkJyn/H4lEAABPfOITD9xWLBb3/J7dbke3253pNSkUinGhwolCoSwdDz30ELLZLCRJQiqVOvSxgUAADMOgVqvtuZ1l2T3pNQAYDocHft9isSj/zzDM2Nv22w9Uq1WEQqEJ/hoKhbJMUOFEoVCWisFggAsuuACvetWr8N73vheve93rDqTJdmO1WvGEJzwBP/vZz/bcHgqFkMvllJ9FUcRPf/pTVdbI8zx+9atf4eyzz1bl+SgUinGgwolCoSwV73rXu9BoNPC3f/u3eOtb34rHP/7xeN3rXnfo75x//vm466679tz2ghe8AF//+tfx9a9/Hf/5n/+JK664AvV6XZU1fu9734PNZsPTn/50VZ6PQqEYByqcKBTK0vDv//7v+OAHP4hbbrkFHo8HLMvilltuwV133YWPfOQjY3/v9a9/Pb7xjW+g0Wgot1122WW45JJLcPHFF+O5z30uTp48iec///mqrPPzn/88LrjgAjgcDlWej0KhGAeG7E/yUygUygryu7/7uzj77LPxjne8Q9PXKZVKeNzjHocf/vCHOHnypKavRaFQ9IdGnCgUyrHgL//yL+FyuTR/nVOnTuGGG26goolCWVFoxIlCoVAoFAplQmjEiUKhUCgUCmVCqHCiUCgUCoVCmRAqnCgUCoVCoVAmhAonCoVCoVAolAmhwolCoVAoFAplQqhwolAoFAqFQpkQKpwoFAqFQqFQJoQKJwqFQqFQKJQJocKJQqFQKBQKZUL+fx9IRHgIcdkwAAAAAElFTkSuQmCC", - "text/plain": [ - "
" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "import matplotlib.pyplot as plt\n", - "\n", - "inv = {v: k for k, v in meta[\"label_maps\"][\"cell_type\"].items()}\n", - "codes = g.categorical_codes[:, 0].numpy()\n", - "pos = g.pos.numpy()\n", - "\n", - "fig, ax = plt.subplots(figsize=(6, 6))\n", - "# Draw every edge: this is a Delaunay triangulation, and subsampling the\n", - "# edges would make it look like something else. edge_index stores both\n", - "# directions, so keep one per pair.\n", - "from matplotlib.collections import LineCollection\n", - "\n", - "src, dst = g.edge_index.numpy()\n", - "keep = src < dst\n", - "segments = [[pos[s], pos[d]] for s, d in zip(src[keep], dst[keep])]\n", - "ax.add_collection(LineCollection(segments, linewidths=0.4, colors=\"0.8\", zorder=1))\n", - "for c in np.unique(codes):\n", - " m = codes == c\n", - " ax.scatter(pos[m, 0], pos[m, 1], s=18, label=inv[int(c)], zorder=2)\n", - "ax.set_xlabel(\"x (um)\")\n", - "ax.set_ylabel(\"y (um)\")\n", - "ax.set_title(f\"patch {sample_entry['patch_id']} - cellular graph\")\n", - "ax.legend(fontsize=7, markerscale=0.8, loc=\"upper right\")\n", - "ax.set_aspect(\"equal\")\n", - "plt.tight_layout()" - ] - }, - { - "cell_type": "markdown", - "id": "8a65eabff63a45729fe45fb5ade58bdc", - "metadata": {}, - "source": [ - "## Next\n", - "\n", - "Vignette 2 pretrains an encoder on these graphs without labels, then uses the\n", - "learned embeddings to discover spatial niches.\n" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.11" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/vignettes/02_pretraining_and_niches.ipynb b/vignettes/02_pretraining_and_niches.ipynb deleted file mode 100644 index ccc6d2d..0000000 --- a/vignettes/02_pretraining_and_niches.ipynb +++ /dev/null @@ -1,1069 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "7fb27b941602401d91542211134fc71a", - "metadata": {}, - "source": [ - "# 2 - Self-supervised pretraining and niche discovery\n", - "\n", - "Vignette 1 turned three MIBI-TOF images into 44 patch graphs. Here we:\n", - "\n", - "1. pretrain a graph encoder with **BGRL** (`experiment=grass_mil_pretrain_a3`), which needs no labels,\n", - "2. fine-tune the MIL model (`regime=mil`) and export per-instance embeddings,\n", - "3. run the interpretability report to discover **spatial niches**.\n", - "\n", - "---\n", - "\n", - "### Read this before interpreting any number below\n", - "\n", - "This dataset has **three images from two donors**. That is enough to\n", - "demonstrate the mechanics and enough for unsupervised niche discovery, but it\n", - "is *not* enough for supervised evaluation:\n", - "\n", - "- Splitting by sample (the correct choice) puts one image in each of\n", - " train/val/test.\n", - "- The only genuine bag-level variable is donor identity, and after that split\n", - " the training image carries a single class.\n", - "\n", - "So the fine-tuning step below trains on a **synthetic** target -- is a patch\n", - "epithelial-rich -- derived from the cell-type composition. That makes it a\n", - "positive control: it shows the training loop, the SSL-to-finetune transfer and\n", - "the attention head all working, and a pipeline that could not recover it would\n", - "be broken. It is not a biological finding, and the metrics are not a benchmark.\n", - "\n", - "The niche analysis at the end stands on its own, because it is unsupervised.\n" - ] - }, - { - "cell_type": "markdown", - "id": "acae54e37e7d407bbb7b55eff062a284", - "metadata": {}, - "source": [ - "## Setup\n" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "id": "9a63283cbaf04dbcab1f6479b197f3a8", - "metadata": { - "execution": { - "iopub.execute_input": "2026-10-13T18:06:17.038316Z", - "iopub.status.busy": "2026-10-13T18:06:17.038225Z", - "iopub.status.idle": "2026-10-13T18:06:17.208367Z", - "shell.execute_reply": "2026-10-13T18:06:17.207510Z" - } - }, - "outputs": [], - "source": [ - "import json\n", - "import os\n", - "import re\n", - "import pathlib\n", - "import subprocess\n", - "import sys\n", - "import glob\n", - "\n", - "import numpy as np\n", - "import pandas as pd\n", - "\n", - "WORK = pathlib.Path(\"vignette_work\").resolve()\n", - "RAW, PROCESSED = WORK / \"data\" / \"raw\", WORK / \"data\" / \"processed\"\n", - "assert PROCESSED.exists(), \"Run vignette 1 first.\"\n", - "env = dict(os.environ, PROJECT_ROOT=str(WORK))\n", - "\n", - "DATA_ARGS = [\n", - " \"data=spatial_omics\",\n", - " f\"data.raw_manifest_path={RAW}/manifest.csv\",\n", - " f\"data.processed_dir={PROCESSED}\",\n", - " \"data.h5ad.coord_source=obsm\",\n", - " \"data.h5ad.coord_key=spatial\",\n", - " \"data.h5ad.categorical_label_columns=[cell_type]\",\n", - " \"data.coord_scale_um=0.39\",\n", - " \"data.sample_unit=tile\",\n", - " \"data.tiling.tile_size_um=100\",\n", - " \"data.tiling.stride_um=100\",\n", - " \"data.min_cells=20\",\n", - " \"data.tiling.min_cells=20\",\n", - " \"data.graph_builder.kwargs.edge_features=[log_distance,neighbor]\",\n", - " \"data.graph_builder.kwargs.neighbor_cutoff_um=20\",\n", - "]\n", - "# Supervised steps split the three images into train / val / test (one each).\n", - "SPLIT_ARGS = [\n", - " \"data.split.train_val_test_split=[0.34,0.33,0.33]\",\n", - " \"data.split.split_by=sample\",\n", - "]\n", - "# Ego-graph settings shared by every model below. The preprint uses depth 4;\n", - "# the demo images are small, so 2 hops keep the CPU runtime short.\n", - "EGO_ARGS = [\n", - " \"data.sampler.runtime.depth=2\",\n", - " \"data.sampler.runtime.subgraph_batch_size=32\",\n", - "]\n", - "\n", - "\n", - "def run(args, tag):\n", - " out = subprocess.run([sys.executable, \"-m\"] + args, capture_output=True, text=True, env=env)\n", - " print(f\"{tag}: exit {out.returncode}\")\n", - " if out.returncode:\n", - " print(out.stdout[-2500:], out.stderr[-2500:])\n", - " raise RuntimeError(f\"{tag} failed - see output above\")\n", - " return out" - ] - }, - { - "cell_type": "markdown", - "id": "8dd0d8092fe74a7c96281538738b07e2", - "metadata": {}, - "source": [ - "## 1. BGRL pretraining (no labels)\n", - "\n", - "BGRL learns node representations by matching two augmented views of the same\n", - "graph. Nothing here uses a label. `experiment=grass_mil_pretrain_a3` is the\n", - "preprint's recipe (the \"A3\" encoder): roots drawn from the union of all\n", - "regions, momentum 0.99 -> 1.0, degree-importance node drop plus Gaussian noise\n", - "on the cell-size feature. This cohort keeps `cell_size` in `obs`, not in the\n", - "marker matrix, so the noise column list is emptied; everything else is the\n", - "default. 200k steps become 300 for a CPU demo, and every tile is used\n", - "(`data.split.train_val_test_split=[1,0,0]` from the experiment preset).\n" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "id": "72eea5119410473aa328ad9291626812", - "metadata": { - "execution": { - "iopub.execute_input": "2026-10-13T18:06:17.210182Z", - "iopub.status.busy": "2026-10-13T18:06:17.210046Z", - "iopub.status.idle": "2026-10-13T18:06:37.834013Z", - "shell.execute_reply": "2026-10-13T18:06:37.833596Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "BGRL pretraining: exit 0\n", - "pretrained encoder: logs/pretrain_a3/runs/2026-10-13_20-06-18/checkpoints/last.ckpt\n" - ] - } - ], - "source": [ - "run(\n", - " [\n", - " \"grass_mil.train\",\n", - " \"experiment=grass_mil_pretrain_a3\",\n", - " \"trainer=cpu\",\n", - " \"trainer.precision=32\",\n", - " *DATA_ARGS,\n", - " *EGO_ARGS,\n", - " \"data.force_precompute=true\", # all tiles label-free (the preset's [1,0,0] split)\n", - " \"task.total_steps=300\",\n", - " \"task.warmup_steps=30\",\n", - " \"task.lr_warmup.warmup_steps=30\",\n", - " \"task.augmentation.feature_noise_columns=[]\",\n", - " \"callbacks.model_checkpoint.every_n_train_steps=100\",\n", - " f\"paths.log_dir={WORK}/logs\",\n", - " ],\n", - " \"BGRL pretraining\",\n", - ")\n", - "\n", - "bgrl_ckpt = sorted(\n", - " glob.glob(f\"{WORK}/logs/pretrain_a3/runs/*/checkpoints/last.ckpt\"), key=os.path.getmtime\n", - ")[-1]\n", - "print(\"pretrained encoder:\", os.path.relpath(bgrl_ckpt, WORK))" - ] - }, - { - "cell_type": "markdown", - "id": "8edb47106e1a46a883d545849b8ab81b", - "metadata": {}, - "source": [ - "## 2. Fine-tune the MIL model\n", - "\n", - "Pretraining gave an encoder. Fine-tuning attaches the MIL head and trains it\n", - "on a bag-level label: each patch is a bag, the ego-graphs inside it are the\n", - "instances, and per-class attention decides which of them the prediction\n", - "leans on. `regime=mil` is the GRASS-MIL default (MLP instance head, attention\n", - "at twice the learning rate, one region per bag, region-level selection); here a\n", - "bag is a tile, so `task.bag_key=patch_id`.\n", - "\n", - "This cohort has no usable clinical label -- see the note at the top -- so we\n", - "assign a **dummy target**: is this patch epithelial-rich? It is derived from\n", - "the cell-type composition the model can see, which makes it a *positive\n", - "control*: a working pipeline must recover it. It demonstrates the training\n", - "loop; it is not a finding.\n", - "\n", - "Deriving it per patch rather than per image matters. Both classes then occur\n", - "in every image, so the split can stay at the sample level -- the correct\n", - "choice, since patches from one image share a donor and a batch.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "id": "10185d26023b46108eb7d9f57d49d2b3", - "metadata": { - "execution": { - "iopub.execute_input": "2026-10-13T18:06:37.835602Z", - "iopub.status.busy": "2026-10-13T18:06:37.835513Z", - "iopub.status.idle": "2026-10-13T18:06:39.239558Z", - "shell.execute_reply": "2026-10-13T18:06:39.238788Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - " patches positive\n", - "sample_id \n", - "point16 16 6\n", - "point23 12 5\n", - "point8 16 11\n" - ] - } - ], - "source": [ - "import torch\n", - "\n", - "index = json.loads((PROCESSED / \"processed_index.json\").read_text())\n", - "entries = index[\"entries\"] if \"entries\" in index else index\n", - "meta = json.loads((PROCESSED / \"metadata.json\").read_text())\n", - "epithelial_code = meta[\"label_maps\"][\"cell_type\"][\"Epithelial\"]\n", - "\n", - "rows = []\n", - "for e in entries:\n", - " g = torch.load(e[\"path\"], weights_only=False)\n", - " fraction = float((g.categorical_codes[:, 0] == epithelial_code).float().mean())\n", - " rows.append({\"patch_id\": e[\"patch_id\"], \"fraction\": fraction, \"sample_id\": e[\"sample_id\"]})\n", - "\n", - "patches = pd.DataFrame(rows)\n", - "# Split at the median so the classes are balanced.\n", - "patches[\"epithelial_rich\"] = (patches[\"fraction\"] > patches[\"fraction\"].median()).astype(int)\n", - "patches[[\"patch_id\", \"epithelial_rich\"]].to_csv(RAW / \"patch_labels.csv\", index=False)\n", - "\n", - "print(\n", - " patches.groupby(\"sample_id\")[\"epithelial_rich\"].agg(patches=\"size\", positive=\"sum\").to_string()\n", - ")" - ] - }, - { - "cell_type": "markdown", - "id": "f6fa52606d8c4a75a9b52967216f8f3f", - "metadata": {}, - "source": [ - "Both classes appear in every image, so no split is degenerate.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "f5a1fa73e5044315a093ec459c9be902", - "metadata": { - "execution": { - "iopub.execute_input": "2026-10-13T18:06:39.240835Z", - "iopub.status.busy": "2026-10-13T18:06:39.240738Z", - "iopub.status.idle": "2026-10-13T18:06:49.440027Z", - "shell.execute_reply": "2026-10-13T18:06:49.439459Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "fine-tuning: exit 0\n", - "fine-tuned checkpoint: logs/train/runs/2026-10-13_20-06-41/checkpoints/last.ckpt\n", - "test/acc: 0.0\n", - "test/loss_region: 0.8508459329605103\n" - ] - } - ], - "source": [ - "LABEL_ARGS = [\n", - " f\"data.graph_labels.label_file={RAW}/patch_labels.csv\",\n", - " \"data.graph_labels.id_column=patch_id\",\n", - " \"data.graph_labels.tasks=[epithelial_rich]\",\n", - " \"data.graph_labels.scope=patch\",\n", - "]\n", - "BAG_ARGS = [\n", - " \"task.bag_key=patch_id\", # one tile = one bag\n", - " \"data.sampler.runtime.subsample_count=32\", # ego-graphs per bag per epoch\n", - " \"data.val_sampler.runtime.subsample_count=64\",\n", - " \"task.region_accumulation.hyperbatch_size=4\",\n", - "]\n", - "\n", - "out = run(\n", - " [\n", - " \"grass_mil.train\",\n", - " \"regime=mil\",\n", - " \"trainer=cpu\",\n", - " \"trainer.precision=32\",\n", - " f\"model.init_from_ckpt={bgrl_ckpt}\",\n", - " *DATA_ARGS,\n", - " *SPLIT_ARGS,\n", - " *EGO_ARGS,\n", - " *LABEL_ARGS,\n", - " *BAG_ARGS,\n", - " # Labels are attached during precompute, so the graphs are rebuilt.\n", - " \"data.force_precompute=true\",\n", - " f\"paths.log_dir={WORK}/logs\",\n", - " \"trainer.max_epochs=8\",\n", - " \"trainer.min_epochs=2\",\n", - " ],\n", - " \"fine-tuning\",\n", - ")\n", - "\n", - "ft_ckpt = sorted(glob.glob(f\"{WORK}/logs/train/runs/*/checkpoints/*.ckpt\"), key=os.path.getmtime)[\n", - " -1\n", - "]\n", - "print(\"fine-tuned checkpoint:\", os.path.relpath(ft_ckpt, WORK))\n", - "\n", - "# Surface the held-out metrics Lightning printed, so the run is visible\n", - "# rather than just its exit code. Rich draws a boxed table with ANSI colour and\n", - "# redraws the progress bar, so strip the codes and keep the final values only.\n", - "ansi = re.compile(r\"\\x1b\\[[0-9;]*[A-Za-z]\")\n", - "seen = {}\n", - "for line in out.stdout.splitlines():\n", - " clean = ansi.sub(\"\", line)\n", - " for key in (\"test/auroc\", \"test/bal_acc\", \"test/acc\", \"test/loss_region\"):\n", - " if key in clean:\n", - " numbers = re.findall(r\"\\d+\\.\\d+\", clean)\n", - " if numbers:\n", - " seen[key] = numbers[-1]\n", - "for key, value in seen.items():\n", - " print(f\"{key}: {value}\")" - ] - }, - { - "cell_type": "markdown", - "id": "cdf66aed5cc84ca1b48e60bad68798a8", - "metadata": {}, - "source": [ - "The `test/` metrics printed above are on the held-out image. With three\n", - "images and a synthetic target they are a smoke test of the loop, not a\n", - "benchmark -- but the loss falling and the encoder transferring from\n", - "pretraining is what a real fine-tuning run looks like.\n", - "\n", - "Swap `patch_labels.csv` for genuine clinical labels and the same command is\n", - "the real workflow.\n" - ] - }, - { - "cell_type": "markdown", - "id": "8763a12b2bbd4a93a75aff182afb95dc", - "metadata": {}, - "source": [ - "## 3. Export per-instance embeddings\n", - "\n", - "`grass-mil-predict` runs inference on the held-out split and writes the tables\n", - "the interpretability suite consumes. Enabling the ShaDow sampler means each\n", - "*instance* is an ego-graph around one cell, so the exported table describes\n", - "local neighbourhoods rather than whole patches.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "id": "7623eae2785240b9bd12b16a66d81610", - "metadata": { - "execution": { - "iopub.execute_input": "2026-10-13T18:06:49.441611Z", - "iopub.status.busy": "2026-10-13T18:06:49.441529Z", - "iopub.status.idle": "2026-10-13T18:06:54.524179Z", - "shell.execute_reply": "2026-10-13T18:06:54.523552Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "predict: exit 0\n", - "logs/predict/runs/2026-10-13_20-06-51/predictions/instance_table.csv\n", - "instances: 1023\n", - "embedding dims: 256\n", - "composition : ['comp_Imm_other', 'comp_Tcell_CD4', 'comp_Endothelial', 'comp_Tcell_CD8', 'comp_Fibroblast', 'comp_Myeloid_CD11c', 'comp_Myeloid_CD68', 'comp_Epithelial']\n", - "attention cols: ['attention', 'attention_c0', 'attention_logit_c0', 'attention_c1', 'attention_logit_c1']\n", - "spatial edges: 5700\n", - "cells: 17709 | cell-cell edges: 77456\n" - ] - } - ], - "source": [ - "run(\n", - " [\n", - " \"grass_mil.inference.predict\",\n", - " f\"ckpt_path={ft_ckpt}\",\n", - " # Must match the checkpoint's architecture and ego-graph depth.\n", - " \"regime=mil\",\n", - " \"trainer=cpu\",\n", - " \"trainer.precision=32\",\n", - " *DATA_ARGS,\n", - " *SPLIT_ARGS,\n", - " *EGO_ARGS,\n", - " *LABEL_ARGS,\n", - " \"task.bag_key=patch_id\",\n", - " \"aggregation.bag_scope=patch\",\n", - " \"interpretability.enabled=true\",\n", - " \"interpretability.spatial.enabled=true\",\n", - " \"interpretability.cells.enabled=true\",\n", - " \"embeddings.enabled=true\",\n", - " f\"paths.log_dir={WORK}/logs\",\n", - " ],\n", - " \"predict\",\n", - ")\n", - "\n", - "instance_table = sorted(\n", - " glob.glob(f\"{WORK}/logs/predict/runs/*/predictions/instance_table.csv\"), key=os.path.getmtime\n", - ")[-1]\n", - "inst = pd.read_csv(instance_table)\n", - "print(os.path.relpath(instance_table, WORK))\n", - "print(\"instances:\", len(inst))\n", - "print(\"embedding dims:\", sum(c.startswith(\"inst_emb_\") for c in inst.columns))\n", - "print(\"composition :\", [c for c in inst.columns if c.startswith(\"comp_\")])\n", - "print(\"attention cols:\", [c for c in inst.columns if c.startswith(\"attention\")])\n", - "spatial_table = sorted(\n", - " glob.glob(f\"{WORK}/logs/predict/runs/*/predictions/spatial_table.csv\"),\n", - " key=os.path.getmtime,\n", - ")[-1]\n", - "print(\"spatial edges:\", len(pd.read_csv(spatial_table)))\n", - "cell_table = sorted(\n", - " glob.glob(f\"{WORK}/logs/predict/runs/*/predictions/cell_table.csv\"),\n", - " key=os.path.getmtime,\n", - ")[-1]\n", - "cell_edge_table = sorted(\n", - " glob.glob(f\"{WORK}/logs/predict/runs/*/predictions/cell_edges.csv\"),\n", - " key=os.path.getmtime,\n", - ")[-1]\n", - "print(\n", - " \"cells:\", len(pd.read_csv(cell_table)), \"| cell-cell edges:\", len(pd.read_csv(cell_edge_table))\n", - ")" - ] - }, - { - "cell_type": "markdown", - "id": "7cdc8c89c7104fffa095e18ddfef8986", - "metadata": {}, - "source": [ - "> **Spatial edges.** The run above exports only the instance table. Adding\n", - "> `interpretability.spatial.enabled=true` also writes a `spatial_table` of\n", - "> edges between neighbouring instances, which the Moran's I, Ripley and\n", - "> neighbourhood-enrichment plugins need. Enable it if you want those analyses;\n", - "> the niche discovery below does not require it.\n" - ] - }, - { - "cell_type": "markdown", - "id": "b118ea5561624da68c537baed56e602f", - "metadata": {}, - "source": [ - "## 4. Discover niches\n", - "\n", - "The report reduces the embedding space, clusters it into niches, orders them\n", - "by composition, and renders an HTML document. Niches are named in dendrogram\n", - "order; unassigned instances are labelled `Background`.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "id": "938c804e27f84196a10c8828c723f798", - "metadata": { - "execution": { - "iopub.execute_input": "2026-10-13T18:06:54.526143Z", - "iopub.status.busy": "2026-10-13T18:06:54.526036Z", - "iopub.status.idle": "2026-10-13T18:07:04.844533Z", - "shell.execute_reply": "2026-10-13T18:07:04.843901Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "interpretability report: exit 0\n", - "report: report/interpretability_report/report.html\n", - "niche_label\n", - "-1 257\n", - " 1 430\n", - " 2 336\n" - ] - } - ], - "source": [ - "run(\n", - " [\n", - " \"grass_mil.interpretability.report_cli\",\n", - " f\"data.instance_table={instance_table}\",\n", - " f\"data.spatial_table={spatial_table}\",\n", - " f\"data.cell_table={cell_table}\",\n", - " f\"data.cell_edge_table={cell_edge_table}\",\n", - " \"interpretability/plugins@plugins=full\",\n", - " # min_cluster_size defaults to 0.5% of the instances, a good rule at\n", - " # cohort scale. Here 0.5% of ~1000 is 5, which shatters the embedding\n", - " # into ~90 niches of a handful of cells each, so pin it for the demo.\n", - " \"clustering.params.min_cluster_size=60\",\n", - " f\"paths.log_dir={WORK}/logs\",\n", - " # report.yaml gets its own hydra run dir; pin it so outputs land\n", - " # in the work directory instead of ./outputs//