diff --git a/.github/workflows/code-quality-main.yaml b/.github/workflows/code-quality-main.yaml index 844d64f..925ef03 100644 --- a/.github/workflows/code-quality-main.yaml +++ b/.github/workflows/code-quality-main.yaml @@ -13,10 +13,10 @@ jobs: steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v5 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: "3.12" diff --git a/.github/workflows/code-quality-pr.yaml b/.github/workflows/code-quality-pr.yaml index 4e07b23..76161a6 100644 --- a/.github/workflows/code-quality-pr.yaml +++ b/.github/workflows/code-quality-pr.yaml @@ -16,10 +16,10 @@ jobs: steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v5 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: "3.12" diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index c130688..6e2b78a 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -18,14 +18,14 @@ jobs: steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v5 - name: Set up conda env (spatial-pinned) uses: conda-incubator/setup-miniconda@v3 with: activate-environment: spatial environment-file: environment.ci-spatial.yaml - auto-activate-base: false + auto-activate: false mamba-version: "*" channels: pyg,pytorch,conda-forge,defaults channel-priority: flexible @@ -61,14 +61,14 @@ jobs: steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v5 - name: Set up conda env (spatial-pinned) uses: conda-incubator/setup-miniconda@v3 with: activate-environment: spatial environment-file: environment.ci-spatial.yaml - auto-activate-base: false + auto-activate: false mamba-version: "*" channels: pyg,pytorch,conda-forge,defaults channel-priority: flexible @@ -77,7 +77,40 @@ jobs: shell: bash -el {0} run: | pip install pytest-cov[toml] - pytest --cov src + pytest --cov grass_mil - name: Upload coverage to Codecov - uses: codecov/codecov-action@v3 + 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 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.11" + + # A plain pip install, deliberately outside conda: 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. + - 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/Makefile b/Makefile index 38184df..df29c6f 100644 --- a/Makefile +++ b/Makefile @@ -20,11 +20,8 @@ sync: ## Merge changes from main branch to your current branch git pull git pull origin main -test: ## Run not slow tests - pytest -k "not slow" - -test-full: ## Run all tests +test: ## Run the test suite pytest train: ## Train the model - python src/train.py + grass-mil-train diff --git a/README.md b/README.md index a732ab8..20926fb 100644 --- a/README.md +++ b/README.md @@ -14,10 +14,10 @@ Geometric deep learning framework for spatial omics data, built around a Hydra-c Core architecture: -- Data pipeline: `src/data/` + `configs/data/` -- Model/runtime pipeline: `src/models/` + `configs/model/` + `configs/task/` -- Interpretability pipeline: `src/interpretability/` + `configs/interpretability/` -- Entrypoints: `src/train.py`, `src/eval.py` +- Data pipeline: `src/grass_mil/data/` + `configs/data/` +- Model/runtime pipeline: `src/grass_mil/models/` + `configs/model/` + `configs/task/` +- Interpretability pipeline: `src/grass_mil/interpretability/` + `configs/interpretability/` +- Entrypoints: `src/grass_mil/train.py`, `src/grass_mil/eval.py` - Configuration system: `configs/` (Hydra groups) ## Interpretability @@ -32,7 +32,7 @@ Core architecture: Quick run: ```bash -python src/interpretability/report_cli.py \ +grass-mil-report \ data.instance_table=/absolute/path/to/instance_table.csv \ data.spatial_table=/absolute/path/to/spatial_table.csv ``` @@ -42,16 +42,30 @@ python src/interpretability/report_cli.py \ ### 1) Install ```bash -cd /Users/lovrorabuzin/Projects/grass-mil_unification/grass-mil +git clone git@github.com:BoevaLab/grass-mil.git +cd grass-mil python -m venv .venv source .venv/bin/activate -pip install -r requirements.txt + +# Editable install for development; drop -e for a plain install. +pip install -e . + +# Optional extras: interpretability plots, report rendering, spatial statistics, +# and the h5ad/SingleCellExperiment readers. +pip install -e ".[interpretability,report,spatial-stats,io]" ``` -Set project root for Hydra path resolution: +Installing exposes the CLI entrypoints (`grass-mil-train`, `grass-mil-eval`, +`grass-mil-loocv`, `grass-mil-predict`, `grass-mil-report`). The Hydra config +tree ships inside the package, so it resolves from any working directory. + +Hydra's `paths.root_dir` defaults to the `PROJECT_ROOT` environment variable +when set (`rootutils` exports it automatically from the `.project-root` marker +in a source checkout) and otherwise falls back to the current working +directory. Set it explicitly to place `data/` and `logs/` elsewhere: ```bash -export PROJECT_ROOT=/Users/lovrorabuzin/Projects/grass-mil_unification/grass-mil +export PROJECT_ROOT=/path/to/your/workspace ``` ### 2) Prepare New Data @@ -67,16 +81,19 @@ Optional: - `region_id` - `polygons_path` +For the full onboarding checklist and dataset schema details, see +[`docs/dataset_preparation.md`](docs/dataset_preparation.md). + ### 3) Run Default Supervised Pipeline ```bash -python src/train.py task=finetune_mean model=supervised_module data=spatial_omics +grass-mil-train task=finetune_mean model=supervised_module data=spatial_omics ``` ### 4) Evaluate A Checkpoint ```bash -python src/eval.py \ +grass-mil-eval \ ckpt_path=/absolute/path/to/checkpoint.ckpt \ task=finetune_mean \ model=supervised_module \ @@ -86,6 +103,7 @@ python src/eval.py \ ## Documentation Index - Full tracked-file map: [`docs/repository_map.md`](docs/repository_map.md) +- Dataset onboarding and preparation checklist: [`docs/dataset_preparation.md`](docs/dataset_preparation.md) - End-to-end workflows (new data, SSL pretrain->finetune, model adjustments, inference): [`docs/workflows_train_infer.md`](docs/workflows_train_infer.md) - Inference utility suite (prediction artifacts, metrics, embeddings): [`docs/inference_utility_suite.md`](docs/inference_utility_suite.md) - Interpretability suite architecture and usage: [`docs/interpretability_suite.md`](docs/interpretability_suite.md) @@ -93,16 +111,18 @@ python src/eval.py \ - Exhaustive hyperparameter reference for repo configs: [`docs/hyperparameter_reference.md`](docs/hyperparameter_reference.md) - Capability modes and operational options: [`docs/capabilities_and_modes.md`](docs/capabilities_and_modes.md) - Existing model contracts: [`docs/model_component_contracts.md`](docs/model_component_contracts.md) -- Existing legacy gap audit: [`docs/legacy_gap_audit.md`](docs/legacy_gap_audit.md) +- Method fidelity map (method -> code -> test): [`docs/method_fidelity.md`](docs/method_fidelity.md) +- Fidelity status and bugs found while closing gaps: [`docs/legacy_gap_audit.md`](docs/legacy_gap_audit.md) ## Supported Training/Inference Modes | Mode | Task Config | Model Config | Entrypoint | Primary Output | |---|---|---|---|---| -| Supervised mean aggregation | `task=finetune_mean` | `model=supervised_module` | `src/train.py` | Region/sample-level logits via mean bag aggregation | -| Supervised MIL attention | `task=finetune_mil` | `model=supervised_module` | `src/train.py` | Attention-weighted bag logits (+ optional MIL auxiliary terms) | -| SSL pretraining (BGRL) | `task=pretrain_bgrl` | `model=bgrl_module` | `src/train.py` | BGRL-pretrained encoder checkpoint | -| Checkpoint evaluation | any compatible task/model | matching training stack | `src/eval.py` | Test metrics from selected checkpoint | +| Supervised mean aggregation | `task=finetune_mean` | `model=supervised_module` | `src/grass_mil/train.py` | Region/sample-level logits via mean bag aggregation | +| Supervised MIL attention | `task=finetune_mil` | `model=supervised_module` | `src/grass_mil/train.py` | Per-class attention-weighted bag logits | +| Cox survival with MIL | `task=finetune_survival` | `model=supervised_module` | `src/grass_mil/train.py` | Scalar log-hazard per bag, scored by c-index | +| SSL pretraining (BGRL) | `task=pretrain_bgrl` | `model=bgrl_module` | `src/grass_mil/train.py` | BGRL-pretrained encoder checkpoint | +| Checkpoint evaluation | any compatible task/model | matching training stack | `src/grass_mil/eval.py` | Test metrics from selected checkpoint | | Advanced prediction (Python API) | compatible with `SupervisedModule.predict_step` | `model=supervised_module` | `Trainer.predict(...)` | `bag_ids`, `bag_logits`, `row_region_ids`, `row_sample_ids`, optional `bag_targets`, optional `bag_attention`, optional `instance_*`, optional `embedding_*` | ## Required Inputs (Manifest Summary) @@ -129,20 +149,20 @@ See full input contracts and file examples in [`docs/workflows_train_infer.md`]( ### Switch aggregation regime ```bash -python src/train.py task=finetune_mean model=supervised_module -python src/train.py task=finetune_mil model=supervised_module +grass-mil-train task=finetune_mean model=supervised_module +grass-mil-train task=finetune_mil model=supervised_module ``` ### Run SSL pretraining ```bash -python src/train.py task=pretrain_bgrl model=bgrl_module +grass-mil-train task=pretrain_bgrl model=bgrl_module ``` ### Change backbone ```bash -python src/train.py task=finetune_mil model=supervised_module \ +grass-mil-train task=finetune_mil model=supervised_module \ model.encoder.conv_type=gine \ model.encoder.use_edge_attr=true \ model.encoder.edge_attr_dim=2 @@ -151,13 +171,13 @@ python src/train.py task=finetune_mil model=supervised_module \ ### Force data re-precompute ```bash -python src/train.py data.force_precompute=true +grass-mil-train data.force_precompute=true ``` ### Resume training ```bash -python src/train.py ckpt_path=/absolute/path/to/last.ckpt +grass-mil-train ckpt_path=/absolute/path/to/last.ckpt ``` ## Where Outputs Go @@ -176,9 +196,20 @@ Data precompute outputs: ## Known Gaps / Non-Goals -Current repository behavior intentionally excludes: +`grass-mil` implements the method, and only the method. It is deliberately not +a reproduction harness for any particular study, so the following live outside +this repository: + +1. Cohort-specific experiment configs and dataset preprocessing. +2. Non-GNN baselines (cell-type composition models, frozen-encoder probes) and + the benchmark tables that compare against them. +3. Cross-validation campaign harnesses, hyperparameter sweep recipes, and + experiment-tracking orchestration. + +The one deliberate exception is the interpretability suite, which is +first-class here because the method's attribution claim is part of the method. -1. Curated first-class hyperparameter sweep recipes in `configs/hparams_search/`. -2. Turnkey dataset-specific benchmark packs beyond the generic Hydra/config contracts. +For what is implemented and the tests that hold it in place, see +[`docs/method_fidelity.md`](docs/method_fidelity.md). All currently supported capabilities are documented in [`docs/workflows_train_infer.md`](docs/workflows_train_infer.md), [`docs/hyperparameter_reference.md`](docs/hyperparameter_reference.md), and [`docs/capabilities_and_modes.md`](docs/capabilities_and_modes.md). diff --git a/configs/data/sampler/shadow_custom.yaml b/configs/data/sampler/shadow_custom.yaml deleted file mode 100644 index 6206abf..0000000 --- a/configs/data/sampler/shadow_custom.yaml +++ /dev/null @@ -1,13 +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 diff --git a/configs/data/sampler/shadow_native.yaml b/configs/data/sampler/shadow_native.yaml deleted file mode 100644 index b1df0f3..0000000 --- a/configs/data/sampler/shadow_native.yaml +++ /dev/null @@ -1,13 +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 diff --git a/configs/interpretability/plugins/default.yaml b/configs/interpretability/plugins/default.yaml deleted file mode 100644 index 3681346..0000000 --- a/configs/interpretability/plugins/default.yaml +++ /dev/null @@ -1,41 +0,0 @@ -enabled: - - cluster_profiles - - attention_attribution -params: - cluster_profiles: - variance_estimator: unbiased - attention_attribution: - attention_column: attention - score_column: score - variance_estimator: unbiased - neighborhood_enrichment: - label_column: cluster_label - # One of: zscore, obs-exp, log2fc. - enrichment_mode: zscore - # n_perms=0 uses analytical expected/std (runtime warning emitted). - # Set n_perms>0 for Squidpy-style permutation expected/std + p-values. - n_perms: 0 - diff_neighborhood_enrichment: - label_column: cluster_label - condition_column: condition - permutation_group_column: sample_id - # One of: zscore, obs-exp, log2fc. - enrichment_mode: zscore - # n_perms=0 computes pairwise differentials without permutation p-values. - n_perms: 0 - filtration_curves: - # Distances and thresholds are interpreted in micrometers. - threshold_start: 0.0 - threshold_stop: 55.0 - threshold_count: 500 - 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 diff --git a/configs/model/attention/gated.yaml b/configs/model/attention/gated.yaml deleted file mode 100644 index 99fd8ae..0000000 --- a/configs/model/attention/gated.yaml +++ /dev/null @@ -1,5 +0,0 @@ -_target_: src.models.components.attention.AttnNetGated -input_dim: 128 -hidden_dim: 64 -dropout: false -n_classes: 1 diff --git a/configs/model/attention/gated_projected.yaml b/configs/model/attention/gated_projected.yaml deleted file mode 100644 index 76f09b5..0000000 --- a/configs/model/attention/gated_projected.yaml +++ /dev/null @@ -1,6 +0,0 @@ -_target_: src.models.components.attention.AttnNetGatedProjected -input_dim: 128 -projection_dim: 64 -hidden_dim: 16 -dropout: false -n_classes: 1 diff --git a/configs/model/components.yaml b/configs/model/components.yaml deleted file mode 100644 index e2d1729..0000000 --- a/configs/model/components.yaml +++ /dev/null @@ -1,56 +0,0 @@ -# @package _global_ - -model: - flags: - use_attention: true - use_ssl: false - - encoder: - _target_: src.models.components.backbones.EncoderConfig - input_dim: 64 - hidden_dim: 128 - out_dim: 128 - num_layers: 3 - dropout: 0.1 - conv_type: gin - norm: batchnorm - jk: last - act: relu - pooling: mean - set2set_steps: 3 - gat_heads: 4 - use_edge_attr: false - edge_weight_index: 0 - edge_attr_dim: null - - attention: - _target_: src.models.components.attention.AttnNetGatedProjected - input_dim: 128 - projection_dim: 64 - hidden_dim: 16 - dropout: false - n_classes: 1 - - ssl: - method: bgrl - predictor: - _target_: src.models.components.ssl.MLPPredictor - input_size: 128 - output_size: 128 - hidden_size: 512 - - graph_head: - _target_: src.models.components.heads.GraphPredictionHead - input_dim: 128 - output_dim: 1 - hidden_dim: 128 - num_layers: 2 - dropout: 0.0 - - node_head: - _target_: src.models.components.heads.NodePredictionHead - input_dim: 128 - output_dim: 1 - hidden_dim: 128 - num_layers: 2 - dropout: 0.0 diff --git a/configs/model/heads/node.yaml b/configs/model/heads/node.yaml deleted file mode 100644 index 9b0be66..0000000 --- a/configs/model/heads/node.yaml +++ /dev/null @@ -1,6 +0,0 @@ -_target_: src.models.components.heads.NodePredictionHead -input_dim: 128 -output_dim: 1 -hidden_dim: 128 -num_layers: 2 -dropout: 0.0 diff --git a/configs/model/loss/categorical_bce.yaml b/configs/model/loss/categorical_bce.yaml deleted file mode 100644 index b497fe8..0000000 --- a/configs/model/loss/categorical_bce.yaml +++ /dev/null @@ -1,2 +0,0 @@ -_target_: src.models.components.losses.WeightedBCEWithLogitsLoss -pos_weight: null diff --git a/configs/model/loss/categorical_ce.yaml b/configs/model/loss/categorical_ce.yaml deleted file mode 100644 index 5a4cee7..0000000 --- a/configs/model/loss/categorical_ce.yaml +++ /dev/null @@ -1,2 +0,0 @@ -_target_: src.models.components.losses.WeightedCrossEntropyLoss -class_weight: null diff --git a/configs/model/loss/regression_huber.yaml b/configs/model/loss/regression_huber.yaml deleted file mode 100644 index e66c563..0000000 --- a/configs/model/loss/regression_huber.yaml +++ /dev/null @@ -1,2 +0,0 @@ -_target_: src.models.components.losses.WeightedHuberLoss -delta: 1.0 diff --git a/configs/model/loss/regression_mse.yaml b/configs/model/loss/regression_mse.yaml deleted file mode 100644 index 81184cc..0000000 --- a/configs/model/loss/regression_mse.yaml +++ /dev/null @@ -1 +0,0 @@ -_target_: src.models.components.losses.WeightedMSELoss diff --git a/configs/model/loss/survival_coxsgd.yaml b/configs/model/loss/survival_coxsgd.yaml deleted file mode 100644 index 5543e12..0000000 --- a/configs/model/loss/survival_coxsgd.yaml +++ /dev/null @@ -1,3 +0,0 @@ -_target_: src.models.components.losses.CoxSGDLoss -top_n: 5 -regularizer_weight: 0.05 diff --git a/docs/capabilities_and_modes.md b/docs/capabilities_and_modes.md index cfa548a..9d384b4 100644 --- a/docs/capabilities_and_modes.md +++ b/docs/capabilities_and_modes.md @@ -7,7 +7,8 @@ This document catalogs operational modes and non-default capabilities in `grass- | Regime | Task | Model | Primary Use | |---|---|---|---| | Supervised mean | `finetune_mean` | `supervised_module` | Stable baseline graph-level classification/regression/survival | -| Supervised MIL | `finetune_mil` | `supervised_module` | Region-level attention aggregation + optional MIL auxiliaries | +| 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 | ## 2) Aggregation And MIL Feature Modes @@ -17,8 +18,6 @@ This document catalogs operational modes and non-default capabilities in `grass- | 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 | -| Node auxiliary loss | `task.node_aux.enabled=true` | Adds node-level auxiliary objective in binary MIL training | -| Entropy regularization | `task.entropy_reg.enabled=true` | Adds entropy-based regularization term | ## 3) Backbone/Encoder Modes @@ -36,7 +35,8 @@ Notable behavior: 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`, `attention`, `set2set`. +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 @@ -115,8 +115,8 @@ Debug presets: Common usage: ```bash -python src/train.py debug=fdr -python src/train.py debug=limit +grass-mil-train debug=fdr +grass-mil-train debug=limit ``` ## 9) Checkpoint Initialization And Resume Modes @@ -131,9 +131,9 @@ python src/train.py debug=limit ## 10) Inference Modes -1. CLI evaluation mode: `src/eval.py` with required `ckpt_path`. +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/inference/predict.py` for prediction export, metrics, and embeddings. +3. Inference utility CLI mode: `src/grass_mil/inference/predict.py` for prediction export, metrics, and embeddings. Predict outputs include: @@ -153,7 +153,7 @@ Metrics available through the inference utility suite: - regression: R2, MAE, RMSE - survival: concordance index (c-index) -Validation (during `src/train.py`) includes: +Validation (during `src/grass_mil/train.py`) includes: - step-level `val/loss` logging - epoch-end region-level aggregation of buffered validation outputs diff --git a/docs/dataset_preparation.md b/docs/dataset_preparation.md new file mode 100644 index 0000000..9f042c1 --- /dev/null +++ b/docs/dataset_preparation.md @@ -0,0 +1,138 @@ +# 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 [`configs/data/spatial_omics.yaml`](/Users/lovrorabuzin/Projects/grass-mil_unification/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. + +## 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/hyperparameter_reference.md b/docs/hyperparameter_reference.md index 1abbc74..bbb7002 100644 --- a/docs/hyperparameter_reference.md +++ b/docs/hyperparameter_reference.md @@ -20,7 +20,7 @@ This reference covers repo-defined configuration keys and behavior for: Override syntax examples use: ```bash -python src/train.py key.path=value +grass-mil-train key.path=value ``` ## 1) Top-Level Train/Eval Composition Keys @@ -29,33 +29,33 @@ python src/train.py key.path=value | Key | Default | Type | Valid Values | Active When | Example Override | |---|---|---|---|---|---| -| `defaults` | composition list | list | Hydra config groups | always | `python src/train.py trainer=gpu` | -| `task_name` | `train` | str | any string | always | `python src/train.py task_name=my_run` | -| `tags` | `[dev]` | list[str] | any list | always | `python src/train.py tags="[paper,ablation]"` | -| `train` | `True` | bool | `true/false` | always | `python src/train.py train=false` | -| `test` | `True` | bool | `true/false` | after fit | `python src/train.py test=false` | -| `ckpt_path` | `null` | str/null | path or null | resume training when set | `python src/train.py ckpt_path=/abs/last.ckpt` | -| `seed` | `null` | int/null | integer or null | random seeding | `python src/train.py seed=42` | +| `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 | `python src/eval.py trainer=gpu` | -| `task_name` | `eval` | str | any string | always | `python src/eval.py task_name=holdout_eval` | -| `tags` | `[dev]` | list[str] | any list | always | `python src/eval.py tags="[final_eval]"` | -| `ckpt_path` | `???` | str | required absolute/relative checkpoint path | always | `python src/eval.py ckpt_path=/abs/model.ckpt` | +| `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 | `python src/loocv.py task_name=loocv_run` | -| `loocv.mode` | `single` | str | `single`, `all` | LOOCV entrypoint | `python src/loocv.py loocv.mode=all` | -| `loocv.fold_index` | `null` | int/null | `>=0` or null | single mode fold selection | `python src/loocv.py loocv.mode=single loocv.fold_index=3` | -| `loocv.per_fold_processed_dir` | `true` | bool | `true/false` | fold isolation | `python src/loocv.py loocv.per_fold_processed_dir=false` | -| `loocv.force_precompute_per_fold` | `true` | bool | `true/false` | leakage-safe fold preprocessing | `python src/loocv.py loocv.force_precompute_per_fold=true` | -| `loocv.summary_filename` | `loocv_summary.json` | str | filename | fold summary artifact naming | `python src/loocv.py loocv.summary_filename=my_summary.json` | +| `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`) @@ -63,39 +63,39 @@ python src/train.py key.path=value | Key | Default | Type | Valid Values | Active When | Example Override | |---|---|---|---|---|---| -| `data._target_` | `src.data.spatial_omics_datamodule.SpatialOmicsDataModule` | str | import path | always | `python src/train.py data._target_=...` | -| `data.data_dir` | `${paths.data_dir}` | str | path | always | `python src/train.py data.data_dir=/tmp/data` | -| `data.raw_manifest_path` | `${paths.data_dir}/raw/manifest.csv` | str | path | always | `python src/train.py data.raw_manifest_path=/abs/manifest.csv` | -| `data.processed_dir` | `${paths.data_dir}/processed/spatial_omics` | str | path | always | `python src/train.py data.processed_dir=/abs/processed` | -| `data.batch_size` | `1` | int | `>=1` | loader batching | `python src/train.py data.batch_size=4` | -| `data.num_workers` | `0` | int | `>=0` | dataloader workers | `python src/train.py data.num_workers=8` | -| `data.pin_memory` | `False` | bool | `true/false` | dataloader | `python src/train.py data.pin_memory=true` | -| `data.coord_scale_um` | `1.0` | float | `>0` | preprocessing (`1.0` means input coords already in `um`) | `python src/train.py data.coord_scale_um=0.5` | -| `data.sample_unit` | `tile` | str | `tile`/`full` | patch generation | `python src/train.py data.sample_unit=full` | -| `data.reducer_scope` | `sample` | str | `sample`/`dataset` | reducer fitting | `python src/train.py data.reducer_scope=dataset` | -| `data.keep_raw_molecular` | `false` | bool | `true/false` | precompute outputs | `python src/train.py data.keep_raw_molecular=true` | -| `data.force_precompute` | `false` | bool | `true/false` | precompute trigger | `python src/train.py data.force_precompute=true` | -| `data.min_cells` | `10` | int | `>=0` | patch filtering | `python src/train.py data.min_cells=25` | -| `data.use_molecular_features` | `true` | bool | `true/false` | loader/reducer behavior | `python src/train.py data.use_molecular_features=false` | +| `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_index` | `python src/train.py data.categorical_features.include_labels="[cell_type,state]"` | +| `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 | `python src/train.py data.split.train_val_test_split="[0.7,0.15,0.15]"` | -| `data.split.split_by` | `sample` | str | `sample`/`region`/`patch` | split semantics | `python src/train.py data.split.split_by=region` | -| `data.split.loocv.enabled` | `false` | bool | `true/false` | LOOCV override path | `python src/train.py data.split.loocv.enabled=true` | -| `data.split.loocv.fold_unit` | `region` | str | `region`/`sample` | LOOCV override path | `python src/train.py data.split.loocv.fold_unit=sample` | -| `data.split.loocv.holdout_id` | `null` | str/null | canonical fold id (`sample` or `sample::region`) | LOOCV override path | `python src/train.py 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 | `python src/train.py 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 | `python src/train.py data.split.loocv.val_ratio=0.2` | -| `data.split.loocv.seed` | `42` | int | integer | LOOCV validation split determinism | `python src/train.py data.split.loocv.seed=7` | +| `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: @@ -114,92 +114,92 @@ Notes: | Key | Default | Type | Valid Values | Active When | Example Override | |---|---|---|---|---|---| -| `data.manifest.sample_id` | `sample_id` | str | manifest column name | manifest parse | `python src/train.py data.manifest.sample_id=sid` | -| `data.manifest.input_path` | `input_path` | str | manifest column name | manifest parse | `python src/train.py data.manifest.input_path=file_path` | -| `data.manifest.input_type` | `input_type` | str | manifest column name | manifest parse | `python src/train.py data.manifest.input_type=source_type` | -| `data.manifest.region_id` | `region_id` | str | manifest column name | optional region grouping | `python src/train.py data.manifest.region_id=roi` | -| `data.manifest.polygons_path` | `polygons_path` | str | manifest column name | optional polygons | `python src/train.py data.manifest.polygons_path=poly_path` | +| `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` | `python src/train.py data.csv.sep=';'` | -| `data.csv.coord_columns` | `[x,y]` | list[str] len=2 | existing columns | `input_type=csv/tsv` | `python src/train.py 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` | `python src/train.py data.csv.cell_id_column=id` | -| `data.csv.categorical_label_columns` | `[]` | list[str] | existing columns | `input_type=csv/tsv` | `python src/train.py 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` | `python src/train.py data.csv.molecular_columns="[gene1,gene2]"` | +| `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` | `python src/train.py data.h5ad.coord_source=obs` | -| `data.h5ad.coord_key` | `spatial` | str/null | obsm key | `coord_source=obsm` | `python src/train.py data.h5ad.coord_key=X_spatial` | -| `data.h5ad.coord_columns` | `[x,y]` | list[str] len=2 | obs columns | `coord_source=obs` | `python src/train.py data.h5ad.coord_columns="[x_um,y_um]"` | -| `data.h5ad.cell_id_column` | `null` | str/null | obs column or null | `input_type=h5ad` | `python src/train.py data.h5ad.cell_id_column=cell_id` | -| `data.h5ad.categorical_label_columns` | `[]` | list[str] | obs columns | `input_type=h5ad` | `python src/train.py data.h5ad.categorical_label_columns="[cell_type]"` | -| `data.h5ad.molecular_layer` | `null` | str/null | layer key or null (`X`) | `use_molecular_features=true` | `python src/train.py data.h5ad.molecular_layer=counts` | -| `data.h5ad.molecular_features` | `null` | list[str]/null | var names subset | `use_molecular_features=true` | `python src/train.py data.h5ad.molecular_features="[CD3,CD8]"` | +| `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` | `python src/train.py data.sce.assay_name=logcounts` | -| `data.sce.coord_source` | `colData` | str | `colData`/`reducedDims` | `input_type=sce/rds` | `python src/train.py data.sce.coord_source=reducedDims` | -| `data.sce.coord_key` | `null` | str/null | reducedDims key | `coord_source=reducedDims` | `python src/train.py data.sce.coord_key=spatial` | -| `data.sce.coord_columns` | `[x,y]` | list[str]/null | colData columns | `coord_source=colData` | `python src/train.py 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` | `python src/train.py data.sce.cell_id_column=cell_id` | -| `data.sce.categorical_label_columns` | `[]` | list[str] | colData columns | `input_type=sce/rds` | `python src/train.py data.sce.categorical_label_columns="[cell_type]"` | -| `data.sce.transpose_assay` | `true` | bool | `true/false` | molecular matrix orientation | `python src/train.py data.sce.transpose_assay=false` | -| `data.sce.molecular_features` | `null` | list[str]/null | rowname subset | `use_molecular_features=true` | `python src/train.py data.sce.molecular_features="[geneA,geneB]"` | +| `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 | `python src/train.py data.graph_builder.name=delaunay` | -| `data.graph_builder.kwargs.edge_features` | `[distance,neighbor]` | list[str] | `distance`, `neighbor` | `delaunay` builder | `python src/train.py data.graph_builder.kwargs.edge_features="[distance]"` | -| `data.graph_builder.kwargs.neighbor_cutoff_um` | `20.0` | float | `>0` | when `neighbor` edge feature enabled | `python src/train.py data.graph_builder.kwargs.neighbor_cutoff_um=50.0` | +| `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 | `python src/train.py data.feature_reducer.name=pca` | -| `data.feature_reducer.fit_mode` | `train_only` | str | `train_only`/`global` | dataset-scope reducer | `python src/train.py data.feature_reducer.fit_mode=global` | -| `data.feature_reducer.kwargs` | `{}` | dict | reducer kwargs | reducer-specific | `python src/train.py data.feature_reducer.kwargs.n_components=64` | +| `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 | Key | Default | Type | Valid Values | Active When | Example Override | |---|---|---|---|---|---| -| `data.sampler.name` | `shadow_custom` | str | `identity`, `shadow_native`, `shadow_custom` | dataloader build | `python src/train.py data.sampler.name=identity` | -| `data.sampler.kwargs` | `{}` | dict | strategy kwargs | strategy-specific | `python src/train.py data.sampler.kwargs='{}'` | -| `data.sampler.runtime.enabled` | `true` | bool | `true/false` | shadow strategies | `python src/train.py data.sampler.runtime.enabled=false` | -| `data.sampler.runtime.depth` | `2` | int | `>=1` | shadow runtime | `python src/train.py data.sampler.runtime.depth=3` | -| `data.sampler.runtime.num_neighbors` | `8` | int | `>=1` | shadow runtime | `python src/train.py data.sampler.runtime.num_neighbors=16` | -| `data.sampler.runtime.subgraph_batch_size` | `32` | int | `>=1` | runtime subgraph batching | `python src/train.py data.sampler.runtime.subgraph_batch_size=64` | -| `data.sampler.runtime.replace` | `false` | bool | `true/false` | root sampling | `python src/train.py data.sampler.runtime.replace=true` | -| `data.sampler.runtime.shuffle_subgraphs` | `true` | bool | `true/false` | training shuffle behavior | `python src/train.py data.sampler.runtime.shuffle_subgraphs=false` | -| `data.sampler.runtime.proportional_root_sampling` | `true` | bool | `true/false` | weighted root sampling | `python src/train.py data.sampler.runtime.proportional_root_sampling=false` | -| `data.sampler.runtime.property_name` | `cell_type` | str | categorical label present in `categorical_slices` | weighted root sampling | `python src/train.py data.sampler.runtime.property_name=cell_type` | -| `data.sampler.runtime.weight_mode` | `inverse` | str | `inverse`, `sqrt_inverse`, `proportional` | weighted root sampling | `python src/train.py data.sampler.runtime.weight_mode=sqrt_inverse` | -| `data.sampler.runtime.min_weight` | `1.0e-6` | float | `>0` | weighted root sampling | `python src/train.py data.sampler.runtime.min_weight=1.0e-4` | +| `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` | int | `>=1` | shadow runtime | `grass-mil-train data.sampler.runtime.depth=3` | +| `data.sampler.runtime.num_neighbors` | `8` | int | `>=1` | shadow runtime | `grass-mil-train data.sampler.runtime.num_neighbors=16` | +| `data.sampler.runtime.subgraph_batch_size` | `32` | int | `>=1` | runtime subgraph batching | `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` | bool | `true/false` | weighted root sampling | `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` | ## 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` | `python src/train.py data.tiling.tile_size_um=300.0` | -| `data.tiling.stride_um` | `200.0` | float | `>0` | `sample_unit=tile` | `python src/train.py data.tiling.stride_um=150.0` | -| `data.tiling.min_cells` | `10` | int | `>=0` | tile retention | `python src/train.py data.tiling.min_cells=20` | -| `data.graph_labels.label_file` | `null` | str/null | CSV path or null | optional graph labels | `python src/train.py data.graph_labels.label_file=/abs/labels.csv` | -| `data.graph_labels.id_column` | `id` | str | column in label file | graph labels enabled | `python src/train.py data.graph_labels.id_column=region_id` | -| `data.graph_labels.tasks` | `null` | list[str]/null | task columns or null for all | graph labels enabled | `python src/train.py data.graph_labels.tasks="[outcome]"` | -| `data.graph_labels.scope` | `region` | str | `region`, `patch`, `sample` | graph labels enabled | `python src/train.py data.graph_labels.scope=sample` | -| `data.transforms` | `[]` | list[dict] | importable transform specs | transform pipeline | `python src/train.py data.transforms='[{name: src.data.components.transforms.CompositionVector, kwargs: {label_attr: label_cell_type}}]'` | +| `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.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`) @@ -215,18 +215,18 @@ Notes: | Key | Default | Type | Valid Values | Active When | Example Override | |---|---|---|---|---|---| -| `task.name` | `finetune_mean` | str | label | always | `python src/train.py task.name=custom_mean` | -| `task.aggregation` | `mean` | str | `mean` | supervised module | `python src/train.py task.aggregation=mean` | -| `task.bag_key` | `region_id` | str | batch attribute name | supervised bagging | `python src/train.py task.bag_key=sample_id` | -| `task.bag_fallback_key` | `sample_id` | str | batch attribute name | bag key missing fallback | `python src/train.py task.bag_fallback_key=patch_id` | -| `task.max_instances_per_bag` | `0` | int | `0` means all, otherwise `>=1` | bag aggregation | `python src/train.py task.max_instances_per_bag=64` | -| `task.instance_sampling` | `all` | str | `all`, `random` | if max instances > 0 | `python src/train.py task.instance_sampling=random` | -| `task.target_type` | `binary` | str | `binary`, `regression`, `survival` | loss routing | `python src/train.py task.target_type=regression` | -| `task.target_columns` | `null` | list[str]/null | label names | target selection | `python src/train.py task.target_columns="[outcome]"` | -| `task.loss` | `categorical_bce` | str | model loss presets | loss build | `python src/train.py task.loss=regression_mse` | -| `task.lr_warmup.enabled` | `false` | bool | `true/false` | scheduler wrapping | `python src/train.py task.lr_warmup.enabled=true` | -| `task.lr_warmup.warmup_steps` | `0` | int | `>=0` | warmup enabled | `python src/train.py task.lr_warmup.warmup_steps=100` | -| `task.lr_warmup.start_factor` | `0.1` | float | `(0,1]` (clamped) | warmup enabled | `python src/train.py task.lr_warmup.start_factor=0.2` | +| `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` @@ -234,50 +234,40 @@ Includes all keys above plus MIL-specific controls: | Key | Default | Type | Valid Values | Active When | Example Override | |---|---|---|---|---|---| -| `task.name` | `finetune_mil` | str | label | always | `python src/train.py task.name=mil_exp` | -| `task.aggregation` | `mil_attention` | str | `mil_attention` | supervised module | `python src/train.py task.aggregation=mil_attention` | -| `task.optimization.backbone_lr` | `null` | float/null | lr value or null | MIL optimizer param groups | `python src/train.py task.optimization.backbone_lr=1e-3` | -| `task.optimization.attention_lr` | `null` | float/null | lr value or null | MIL optimizer param groups | `python src/train.py task.optimization.attention_lr=4e-3` | -| `task.region_accumulation.enabled` | `false` | bool | `true/false` | MIL training only | `python src/train.py task.region_accumulation.enabled=true` | -| `task.region_accumulation.mode` | `manual_region_buffer` | str | currently manual mode in code | MIL accumulation metadata | `python src/train.py task.region_accumulation.mode=manual_region_buffer` | -| `task.region_accumulation.hyperbatch_size` | `8` | int | `>=1` when enabled | MIL accumulation | `python src/train.py task.region_accumulation.hyperbatch_size=4` | -| `task.region_accumulation.flush_on_epoch_end` | `true` | bool | `true/false` | MIL accumulation | `python src/train.py task.region_accumulation.flush_on_epoch_end=true` | -| `task.node_aux.enabled` | `false` | bool | `true/false` | MIL + binary training stage | `python src/train.py task.node_aux.enabled=true` | -| `task.node_aux.target_mode` | `attention_shaped_ti` | str | `graph_label`, `attention_shaped_ti` | node auxiliary targets | `python src/train.py task.node_aux.target_mode=graph_label` | -| `task.node_aux.loss_mode` | `bce` | str | `bce`, `weighted_bce` | node auxiliary loss | `python src/train.py task.node_aux.loss_mode=weighted_bce` | -| `task.node_aux.weight` | `1.0` | float | any float | node auxiliary weighting | `python src/train.py task.node_aux.weight=0.2` | -| `task.entropy_reg.enabled` | `false` | bool | `true/false` | MIL + binary training stage | `python src/train.py task.entropy_reg.enabled=true` | -| `task.entropy_reg.mode` | `attention_shaped_target` | str | `attention`, `attention_shaped_target` | entropy regularization | `python src/train.py task.entropy_reg.mode=attention` | -| `task.entropy_reg.weight` | `0.01` | float | any float | entropy weighting | `python src/train.py task.entropy_reg.weight=0.005` | -| `task.loss_weights.region` | `1.0` | float | any float | MIL total loss composition | `python src/train.py task.loss_weights.region=1.0` | -| `task.loss_weights.node_aux` | `1.0` | float | any float | MIL total loss composition | `python src/train.py task.loss_weights.node_aux=0.5` | -| `task.loss_weights.entropy` | `0.0` | float | any float | MIL total loss composition | `python src/train.py task.loss_weights.entropy=0.1` | +| `task.name` | `finetune_mil` | str | label | always | `grass-mil-train task.name=mil_exp` | +| `task.aggregation` | `mil_attention` | str | `mil_attention` | supervised module | `grass-mil-train task.aggregation=mil_attention` | +| `task.optimization.backbone_lr` | `null` | float/null | lr value or null | MIL optimizer param groups | `grass-mil-train task.optimization.backbone_lr=1e-3` | +| `task.optimization.attention_lr` | `null` | float/null | lr value or null | MIL optimizer param groups | `grass-mil-train task.optimization.attention_lr=4e-3` | +| `task.region_accumulation.enabled` | `false` | bool | `true/false` | MIL training only | `grass-mil-train task.region_accumulation.enabled=true` | +| `task.region_accumulation.mode` | `manual_region_buffer` | str | currently manual mode in code | MIL accumulation metadata | `grass-mil-train task.region_accumulation.mode=manual_region_buffer` | +| `task.region_accumulation.hyperbatch_size` | `8` | int | `>=1` when enabled | MIL accumulation | `grass-mil-train task.region_accumulation.hyperbatch_size=4` | +| `task.region_accumulation.flush_on_epoch_end` | `true` | bool | `true/false` | MIL accumulation | `grass-mil-train task.region_accumulation.flush_on_epoch_end=true` | ## 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 | `python src/train.py task=pretrain_bgrl` | -| `task.name` | `pretrain_bgrl` | str | label | always | `python src/train.py task.name=ssl_run` | -| `task.aggregation` | `mean` | str | `mean` | BGRL module ignores bag loss semantics but keeps task schema | `python src/train.py task.aggregation=mean` | -| `task.bag_key` | `region_id` | str | batch attr | task metadata | `python src/train.py task.bag_key=sample_id` | -| `task.bag_fallback_key` | `sample_id` | str | batch attr | task metadata | `python src/train.py task.bag_fallback_key=patch_id` | -| `task.max_instances_per_bag` | `0` | int | `0`/`>=1` | task metadata | `python src/train.py task.max_instances_per_bag=0` | -| `task.instance_sampling` | `all` | str | `all`/`random` | task metadata | `python src/train.py task.instance_sampling=all` | -| `task.target_type` | `binary` | str | task labels | task metadata | `python src/train.py task.target_type=binary` | -| `task.target_columns` | `null` | list/null | optional | task metadata | `python src/train.py task.target_columns='[label]'` | -| `task.loss` | `categorical_bce` | str | loss label | task metadata | `python src/train.py task.loss=categorical_bce` | -| `task.lr_warmup.enabled` | `true` | bool | `true/false` | optimizer scheduler warmup | `python src/train.py task.lr_warmup.enabled=true` | -| `task.lr_warmup.warmup_steps` | `1000` | int | `>=0` | warmup enabled | `python src/train.py task.lr_warmup.warmup_steps=500` | -| `task.lr_warmup.start_factor` | `0.1` | float | `(0,1]` | warmup enabled | `python src/train.py task.lr_warmup.start_factor=0.2` | -| `task.warmup_steps` | `1000` | int | `>=0` | BGRL momentum scheduler | `python src/train.py task.warmup_steps=500` | -| `task.total_steps` | `10000` | int | `>0` preferred | BGRL momentum scheduler | `python src/train.py task.total_steps=20000` | -| `task.momentum` | `0.99` | float | `[0,1]` practical range | target network EMA | `python src/train.py task.momentum=0.996` | -| `task.momentum_min` | `0.99` | float | `[0,1]` practical range | cosine momentum floor | `python src/train.py task.momentum_min=0.9` | -| `task.drop_edge_p1` | `0.2` | float | `[0,1]` | view-1 edge dropout | `python src/train.py task.drop_edge_p1=0.1` | -| `task.drop_edge_p2` | `0.2` | float | `[0,1]` | view-2 edge dropout | `python src/train.py task.drop_edge_p2=0.3` | -| `task.drop_feat_p1` | `0.2` | float | `[0,1]` | view-1 feature dropout | `python src/train.py task.drop_feat_p1=0.1` | -| `task.drop_feat_p2` | `0.2` | float | `[0,1]` | view-2 feature dropout | `python src/train.py task.drop_feat_p2=0.3` | +| `defaults` | `override /scheduler: cosine_step` | Hydra defaults entry | scheduler group override | pretrain task composition | `grass-mil-train task=pretrain_bgrl` | +| `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` | ## 4) Model Config Group (`configs/model/**/*.yaml`) @@ -285,45 +275,44 @@ Includes all keys above plus MIL-specific controls: | Key | Default | Type | Valid Values | Active When | Example Override | |---|---|---|---|---|---| -| `model._target_` | `src.models.supervised_module.SupervisedModule` | str | import path | supervised runs | `python src/train.py model=supervised_module` | -| `model._recursive_` | `false` | bool | `true/false` | Hydra instantiate semantics | `python src/train.py model._recursive_=false` | -| `model.flags.use_ssl` | `false` | bool | `true/false` | optional SSL wrapper in supervised module | `python src/train.py model.flags.use_ssl=true` | -| `model.encoder.input_dim` | `0` | int | `>0` or `0` for infer | encoder build | `python src/train.py model.encoder.input_dim=128` | -| `model.encoder.hidden_dim` | `128` | int | `>=1` | encoder | `python src/train.py model.encoder.hidden_dim=256` | -| `model.encoder.out_dim` | `128` | int/null | `>=1` or null | encoder projection | `python src/train.py model.encoder.out_dim=256` | -| `model.encoder.num_layers` | `3` | int | `>=1` | encoder | `python src/train.py model.encoder.num_layers=4` | -| `model.encoder.dropout` | `0.1` | float | `[0,1]` | encoder | `python src/train.py model.encoder.dropout=0.2` | -| `model.encoder.conv_type` | `gin` | str | `gin`, `gcn`, `gat`, `graphsage`, `gine` | encoder | `python src/train.py model.encoder.conv_type=gine` | -| `model.encoder.norm` | `batchnorm` | str | `batchnorm`, `layernorm`, `none` | encoder | `python src/train.py model.encoder.norm=layernorm` | -| `model.encoder.jk` | `last` | str | `last`, `concat`, `max`, `sum` | encoder | `python src/train.py model.encoder.jk=concat` | -| `model.encoder.act` | `relu` | str | `relu`, `gelu`, `leaky_relu` | encoder | `python src/train.py model.encoder.act=gelu` | -| `model.encoder.pooling` | `mean` | str | `sum`,`mean`,`max`,`attention`,`set2set` | graph embedding | `python src/train.py model.encoder.pooling=set2set` | -| `model.encoder.set2set_steps` | `3` | int | `>=1` | `pooling=set2set` | `python src/train.py model.encoder.set2set_steps=4` | -| `model.encoder.gat_heads` | `4` | int | `>=1` and divides hidden dim | `conv_type=gat` | `python src/train.py model.encoder.gat_heads=8` | -| `model.encoder.use_edge_attr` | `false` | bool | `true/false` | `gcn`/`gine` meaningful | `python src/train.py model.encoder.use_edge_attr=true` | -| `model.encoder.edge_weight_index` | `0` | int | column index | `gcn` with edge attrs | `python src/train.py model.encoder.edge_weight_index=1` | -| `model.encoder.edge_attr_dim` | `null` | int/null | required for `gine` | `conv_type=gine` | `python src/train.py model.encoder.edge_attr_dim=2` | -| `model.attention.attention_type` | `gated_projected` | str | `gated`, `gated_projected` | `task.aggregation=mil_attention` | `python src/train.py model.attention.attention_type=gated` | -| `model.attention.input_dim` | `128` | int | `>=1` | MIL attention | `python src/train.py model.attention.input_dim=256` | -| `model.attention.projection_dim` | `64` | int | `>=1` | projected attention | `python src/train.py model.attention.projection_dim=128` | -| `model.attention.hidden_dim` | `16` | int | `>=1` | MIL attention | `python src/train.py model.attention.hidden_dim=32` | -| `model.attention.dropout` | `false` | bool | `true/false` | MIL attention | `python src/train.py model.attention.dropout=true` | -| `model.attention.n_classes` | `1` | int | `>=1` | MIL attention output channels | `python src/train.py model.attention.n_classes=1` | -| `model.graph_head.input_dim` | `128` | int | `>=1` | graph head | `python src/train.py model.graph_head.input_dim=256` | -| `model.graph_head.output_dim` | `1` | int | `>=1` | graph head | `python src/train.py model.graph_head.output_dim=2` | -| `model.graph_head.hidden_dim` | `128` | int | `>=1` | graph head | `python src/train.py model.graph_head.hidden_dim=256` | -| `model.graph_head.num_layers` | `2` | int | `>=1` | graph head | `python src/train.py model.graph_head.num_layers=3` | -| `model.graph_head.dropout` | `0.0` | float | `[0,1]` | graph head | `python src/train.py model.graph_head.dropout=0.2` | -| `model.ssl.method` | `bgrl` | str | `bgrl` | SSL factory | `python src/train.py model.ssl.method=bgrl` | -| `model.ssl.predictor.hidden_size` | `512` | int | `>=1` | SSL enabled | `python src/train.py model.ssl.predictor.hidden_size=256` | -| `model.loss.loss_type` | `${task.loss}` | str | loss name | loss factory | `python src/train.py task.loss=regression_mse` | -| `model.optim` | `${optim}` | DictConfig | optimizer config object | always | `python src/train.py optim.lr=5e-4` | -| `model.scheduler` | `${scheduler}` | DictConfig | scheduler config object | always | `python src/train.py scheduler.T_max=200` | -| `model.task` | `${task}` | DictConfig | task config object | always | `python src/train.py task=finetune_mil` | -| `model.init_from_ckpt` | `null` | str/null | checkpoint path | optional init | `python src/train.py model.init_from_ckpt=/abs/pretrain.ckpt` | -| `model.init_strict` | `false` | bool | `true/false` | checkpoint init | `python src/train.py model.init_strict=true` | -| `model.encoder_init_map` | `auto_bgrl_or_identity` | str | `identity`, `auto_bgrl_or_identity` | checkpoint init | `python src/train.py model.encoder_init_map=identity` | -| `model.freeze_encoder` | `false` | bool | `true/false` | checkpoint transfer/finetuning | `python src/train.py model.freeze_encoder=true` | +| `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.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` | `64` | int | `>=1` | projected attention | `grass-mil-train model.attention.projection_dim=128` | +| `model.attention.hidden_dim` | `16` | int | `>=1` | MIL attention | `grass-mil-train model.attention.hidden_dim=32` | +| `model.attention.dropout` | `false` | bool | `true/false` | MIL attention | `grass-mil-train model.attention.dropout=true` | +| `model.attention.n_classes` | `1` | int | `>=1` | MIL attention output channels | `grass-mil-train model.attention.n_classes=1` | +| `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` | graph head | `grass-mil-train model.graph_head.num_layers=3` | +| `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` | checkpoint init | `grass-mil-train model.encoder_init_map=identity` | +| `model.freeze_encoder` | `false` | bool | `true/false` | checkpoint transfer/finetuning | `grass-mil-train model.freeze_encoder=true` | ## 4.2 `configs/model/bgrl_module.yaml` @@ -331,17 +320,17 @@ Same encoder block as supervised module plus SSL/pretrain init keys: | Key | Default | Type | Valid Values | Active When | Example Override | |---|---|---|---|---|---| -| `model._target_` | `src.models.bgrl_module.BGRLModule` | str | import path | pretrain task | `python src/train.py model=bgrl_module` | -| `model._recursive_` | `false` | bool | `true/false` | instantiate behavior | `python src/train.py model._recursive_=false` | -| `model.encoder.*` | same as supervised defaults | mixed | see section 4.1 | BGRL encoder | `python src/train.py model.encoder.conv_type=gin` | -| `model.ssl.method` | `bgrl` | str | `bgrl` | SSL module | `python src/train.py model.ssl.method=bgrl` | -| `model.ssl.predictor.hidden_size` | `512` | int | `>=1` | predictor width | `python src/train.py model.ssl.predictor.hidden_size=1024` | -| `model.optim` | `${optim}` | DictConfig | optimizer object | always | `python src/train.py optim.lr=1e-3` | -| `model.scheduler` | `${scheduler}` | DictConfig | scheduler object | always | `python src/train.py scheduler.T_max=20000` | -| `model.task` | `${task}` | DictConfig | pretrain task object | always | `python src/train.py task=pretrain_bgrl` | -| `model.init_from_ckpt` | `null` | str/null | checkpoint path | optional init | `python src/train.py model.init_from_ckpt=/abs/ssl.ckpt` | -| `model.init_strict` | `false` | bool | `true/false` | checkpoint init | `python src/train.py model.init_strict=true` | -| `model.encoder_init_map` | `auto_bgrl_or_identity` | str | `identity`, `auto_bgrl_or_identity` | checkpoint init | `python src/train.py model.encoder_init_map=identity` | +| `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.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 Component Preset Files @@ -349,21 +338,21 @@ Same encoder block as supervised module plus SSL/pretrain init keys: | Key | Default | Type | Valid Values | Active When | Example Override | |---|---|---|---|---|---| -| `model.flags.use_attention` | `true` | bool | `true/false` | component notebook/testing | `python src/train.py model.flags.use_attention=false` | -| `model.flags.use_ssl` | `false` | bool | `true/false` | component notebook/testing | `python src/train.py model.flags.use_ssl=true` | -| `model.encoder._target_` | `src.models.components.backbones.EncoderConfig` | str | import path | hydra instantiate | `python src/train.py model.encoder._target_=...` | -| `model.encoder.*` | same fields as section 4.1 | mixed | see section 4.1 | component testing | `python src/train.py model.encoder.conv_type=gcn` | -| `model.attention._target_` | `src.models.components.attention.AttnNetGatedProjected` | str | import path | component testing | `python src/train.py model.attention._target_=...AttnNetGated` | -| `model.attention.*` | defaults per file | mixed | see attention files below | component testing | `python src/train.py model.attention.hidden_dim=32` | -| `model.ssl.method` | `bgrl` | str | `bgrl` | component testing | `python src/train.py model.ssl.method=bgrl` | -| `model.ssl.predictor._target_` | `src.models.components.ssl.MLPPredictor` | str | import path | component testing | `python src/train.py model.ssl.predictor._target_=...` | -| `model.ssl.predictor.input_size` | `128` | int | `>=1` | component testing | `python src/train.py model.ssl.predictor.input_size=256` | -| `model.ssl.predictor.output_size` | `128` | int | `>=1` | component testing | `python src/train.py model.ssl.predictor.output_size=256` | -| `model.ssl.predictor.hidden_size` | `512` | int | `>=1` | component testing | `python src/train.py model.ssl.predictor.hidden_size=256` | -| `model.graph_head._target_` | `src.models.components.heads.GraphPredictionHead` | str | import path | component testing | `python src/train.py model.graph_head._target_=...` | -| `model.graph_head.*` | see section 4.1 | mixed | MLP head args | component testing | `python src/train.py model.graph_head.num_layers=3` | -| `model.node_head._target_` | `src.models.components.heads.NodePredictionHead` | str | import path | component testing | `python src/train.py model.node_head._target_=...` | -| `model.node_head.*` | graph-head-like args | mixed | MLP head args | component testing | `python src/train.py model.node_head.dropout=0.1` | +| `model.flags.use_attention` | `true` | bool | `true/false` | component notebook/testing | `grass-mil-train model.flags.use_attention=false` | +| `model.flags.use_ssl` | `false` | bool | `true/false` | component notebook/testing | `grass-mil-train model.flags.use_ssl=true` | +| `model.encoder._target_` | `grass_mil.models.components.backbones.EncoderConfig` | str | import path | hydra instantiate | `grass-mil-train model.encoder._target_=...` | +| `model.encoder.*` | same fields as section 4.1 | mixed | see section 4.1 | component testing | `grass-mil-train model.encoder.conv_type=gcn` | +| `model.attention._target_` | `grass_mil.models.components.attention.AttnNetGatedProjected` | str | import path | component testing | `grass-mil-train model.attention._target_=...AttnNetGated` | +| `model.attention.*` | defaults per file | mixed | see attention files below | component testing | `grass-mil-train model.attention.hidden_dim=32` | +| `model.ssl.method` | `bgrl` | str | `bgrl` | component testing | `grass-mil-train model.ssl.method=bgrl` | +| `model.ssl.predictor._target_` | `grass_mil.models.components.ssl.MLPPredictor` | str | import path | component testing | `grass-mil-train model.ssl.predictor._target_=...` | +| `model.ssl.predictor.input_size` | `128` | int | `>=1` | component testing | `grass-mil-train model.ssl.predictor.input_size=256` | +| `model.ssl.predictor.output_size` | `128` | int | `>=1` | component testing | `grass-mil-train model.ssl.predictor.output_size=256` | +| `model.ssl.predictor.hidden_size` | `512` | int | `>=1` | component testing | `grass-mil-train model.ssl.predictor.hidden_size=256` | +| `model.graph_head._target_` | `grass_mil.models.components.heads.GraphPredictionHead` | str | import path | component testing | `grass-mil-train model.graph_head._target_=...` | +| `model.graph_head.*` | see section 4.1 | mixed | MLP head args | component testing | `grass-mil-train model.graph_head.num_layers=3` | +| `model.node_head._target_` | `grass_mil.models.components.heads.NodePredictionHead` | str | import path | component testing | `grass-mil-train model.node_head._target_=...` | +| `model.node_head.*` | graph-head-like args | mixed | MLP head args | component testing | `grass-mil-train model.node_head.dropout=0.1` | ### Encoder preset files (`configs/model/encoder/*.yaml`) @@ -377,7 +366,6 @@ Covered files: All five files expose the same keys: -- `_target_`, `input_dim`, `hidden_dim`, `out_dim`, `num_layers`, `dropout`, `conv_type`, `norm`, `jk`, `act`, `pooling`, `set2set_steps`, `gat_heads`, `use_edge_attr`, `edge_weight_index`, `edge_attr_dim`. Per-file defaults differ mainly in `conv_type`, `use_edge_attr`, and `edge_attr_dim`: @@ -404,7 +392,6 @@ Covered files: Covered files: - `configs/model/heads/graph.yaml` -- `configs/model/heads/node.yaml` Both `graph.yaml` and `node.yaml` expose: @@ -448,20 +435,20 @@ Keys: | Key | Default | Type | Valid Values | Active When | Example Override | |---|---|---|---|---|---| -| `optim._target_` | `torch.optim.AdamW` | str | optimizer class path | always | `python src/train.py optim._target_=torch.optim.SGD` | -| `optim.lr` | `0.001` | float | `>0` | optimizer step | `python src/train.py optim.lr=5e-4` | -| `optim.weight_decay` | `0.01` | float | `>=0` | optimizer step | `python src/train.py optim.weight_decay=1e-4` | +| `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 | `python src/train.py scheduler._target_=...` | -| `configs/scheduler/cosine_epoch.yaml` | `T_max` | `100` | int | `python src/train.py scheduler.T_max=200` | -| `configs/scheduler/cosine_epoch.yaml` | `eta_min` | `1.0e-6` | float | `python src/train.py scheduler.eta_min=1.0e-7` | -| `configs/scheduler/cosine_step.yaml` | `_target_` | `torch.optim.lr_scheduler.CosineAnnealingLR` | str | `python src/train.py scheduler._target_=...` | -| `configs/scheduler/cosine_step.yaml` | `T_max` | `10000` | int | `python src/train.py scheduler.T_max=20000` | -| `configs/scheduler/cosine_step.yaml` | `eta_min` | `1.0e-7` | float | `python src/train.py scheduler.eta_min=1.0e-8` | +| `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` | ## 6) Trainer Config Group (`configs/trainer/*.yaml`) @@ -469,14 +456,14 @@ Keys: | Key | Default | Type | Valid Values | Active When | Example Override | |---|---|---|---|---|---| -| `trainer._target_` | `lightning.pytorch.trainer.Trainer` | str | trainer import path | always | `python src/train.py trainer._target_=...` | -| `trainer.default_root_dir` | `${paths.output_dir}` | str | path | trainer init | `python src/train.py trainer.default_root_dir=/tmp/out` | -| `trainer.min_epochs` | `1` | int | `>=0` | fit loop | `python src/train.py trainer.min_epochs=0` | -| `trainer.max_epochs` | `10` | int | `>=1` | fit loop | `python src/train.py trainer.max_epochs=100` | -| `trainer.accelerator` | `cpu` | str | lightning-supported accelerators | runtime | `python src/train.py trainer.accelerator=gpu` | -| `trainer.devices` | `1` | int/list/str | lightning-supported device spec | runtime | `python src/train.py trainer.devices=2` | -| `trainer.check_val_every_n_epoch` | `1` | int | `>=1` | fit loop | `python src/train.py trainer.check_val_every_n_epoch=2` | -| `trainer.deterministic` | `False` | bool | `true/false` | reproducibility | `python src/train.py trainer.deterministic=true` | +| `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 @@ -494,52 +481,52 @@ Keys: | 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 | `python src/train.py callbacks=none` | -| `callbacks.model_checkpoint.dirpath` | `${paths.output_dir}/checkpoints` | str | path | checkpoint callback enabled | `python src/train.py callbacks.model_checkpoint.dirpath=/tmp/ckpt` | -| `callbacks.model_checkpoint.filename` | `epoch_{epoch:03d}` | str | format string | checkpoint callback enabled | `python src/train.py callbacks.model_checkpoint.filename=best` | -| `callbacks.model_checkpoint.monitor` | `val/loss` | str | logged metric name | checkpoint callback enabled | `python src/train.py callbacks.model_checkpoint.monitor=val/acc` | -| `callbacks.model_checkpoint.mode` | `min` | str | `min`, `max` | checkpoint callback enabled | `python src/train.py callbacks.model_checkpoint.mode=max` | -| `callbacks.model_checkpoint.save_last` | `True` | bool | `true/false` | checkpoint callback enabled | `python src/train.py callbacks.model_checkpoint.save_last=false` | -| `callbacks.model_checkpoint.auto_insert_metric_name` | `False` | bool | `true/false` | checkpoint callback enabled | `python src/train.py callbacks.model_checkpoint.auto_insert_metric_name=true` | -| `callbacks.early_stopping.monitor` | `val/loss` | str | logged metric name | early stopping enabled | `python src/train.py callbacks.early_stopping.monitor=val/acc` | -| `callbacks.early_stopping.patience` | `100` | int | `>=0` | early stopping enabled | `python src/train.py callbacks.early_stopping.patience=20` | -| `callbacks.early_stopping.mode` | `min` | str | `min`, `max` | early stopping enabled | `python src/train.py callbacks.early_stopping.mode=max` | -| `callbacks.model_summary.max_depth` | `-1` | int | integer depth | model summary enabled | `python src/train.py callbacks.model_summary.max_depth=2` | +| `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 | `python src/train.py callbacks.model_checkpoint._target_=...` | -| `model_checkpoint.dirpath` | `null` | str/null | path or null | callback active | `python src/train.py callbacks.model_checkpoint.dirpath=/tmp/ckpt` | -| `model_checkpoint.filename` | `null` | str/null | filename template or null | callback active | `python src/train.py callbacks.model_checkpoint.filename=epoch_{epoch}` | -| `model_checkpoint.monitor` | `null` | str/null | metric name or null | callback active | `python src/train.py callbacks.model_checkpoint.monitor=val/loss` | -| `model_checkpoint.verbose` | `False` | bool | `true/false` | callback active | `python src/train.py callbacks.model_checkpoint.verbose=true` | -| `model_checkpoint.save_last` | `null` | bool/null | bool or null | callback active | `python src/train.py callbacks.model_checkpoint.save_last=true` | -| `model_checkpoint.save_top_k` | `1` | int | integer | callback active | `python src/train.py callbacks.model_checkpoint.save_top_k=3` | -| `model_checkpoint.mode` | `min` | str | `min`, `max` | callback active | `python src/train.py callbacks.model_checkpoint.mode=max` | -| `model_checkpoint.auto_insert_metric_name` | `True` | bool | `true/false` | callback active | `python src/train.py callbacks.model_checkpoint.auto_insert_metric_name=false` | -| `model_checkpoint.save_weights_only` | `False` | bool | `true/false` | callback active | `python src/train.py callbacks.model_checkpoint.save_weights_only=true` | -| `model_checkpoint.every_n_train_steps` | `null` | int/null | integer or null | callback active | `python src/train.py callbacks.model_checkpoint.every_n_train_steps=100` | -| `model_checkpoint.train_time_interval` | `null` | duration/null | Lightning duration format or null | callback active | `python src/train.py callbacks.model_checkpoint.train_time_interval='00:10:00'` | -| `model_checkpoint.every_n_epochs` | `null` | int/null | integer or null | callback active | `python src/train.py callbacks.model_checkpoint.every_n_epochs=1` | -| `model_checkpoint.save_on_train_epoch_end` | `null` | bool/null | bool or null | callback active | `python src/train.py callbacks.model_checkpoint.save_on_train_epoch_end=true` | +| `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 | `python src/train.py callbacks.early_stopping._target_=...` | -| `early_stopping.monitor` | `???` | str | metric name | early stopping enabled | `python src/train.py callbacks.early_stopping.monitor=val/loss` | -| `early_stopping.min_delta` | `0.` | float | any float | early stopping enabled | `python src/train.py callbacks.early_stopping.min_delta=0.001` | -| `early_stopping.patience` | `3` | int | `>=0` | early stopping enabled | `python src/train.py callbacks.early_stopping.patience=10` | -| `early_stopping.verbose` | `False` | bool | `true/false` | early stopping enabled | `python src/train.py callbacks.early_stopping.verbose=true` | -| `early_stopping.mode` | `min` | str | `min`, `max` | early stopping enabled | `python src/train.py callbacks.early_stopping.mode=max` | -| `early_stopping.strict` | `True` | bool | `true/false` | early stopping enabled | `python src/train.py callbacks.early_stopping.strict=false` | -| `early_stopping.check_finite` | `True` | bool | `true/false` | early stopping enabled | `python src/train.py callbacks.early_stopping.check_finite=true` | -| `early_stopping.stopping_threshold` | `null` | float/null | float or null | early stopping enabled | `python src/train.py callbacks.early_stopping.stopping_threshold=0.1` | -| `early_stopping.divergence_threshold` | `null` | float/null | float or null | early stopping enabled | `python src/train.py callbacks.early_stopping.divergence_threshold=10.0` | -| `early_stopping.check_on_train_epoch_end` | `null` | bool/null | bool or null | early stopping enabled | `python src/train.py callbacks.early_stopping.check_on_train_epoch_end=true` | +| `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 @@ -551,10 +538,10 @@ Covered files: | Key | Default | Type | Valid Values | Active When | Example Override | |---|---|---|---|---|---| -| `model_summary._target_` | `lightning.pytorch.callbacks.RichModelSummary` | str | import path | `callbacks=model_summary` | `python src/train.py callbacks.model_summary._target_=...` | -| `model_summary.max_depth` | `1` | int | integer depth | `callbacks=model_summary` | `python src/train.py callbacks.model_summary.max_depth=2` | -| `rich_progress_bar._target_` | `lightning.pytorch.callbacks.RichProgressBar` | str | import path | `callbacks=rich_progress_bar` | `python src/train.py callbacks.rich_progress_bar._target_=...` | -| `configs/callbacks/none.yaml` | empty file | n/a | no fields | `callbacks=none` | `python src/train.py callbacks=none` | +| `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` | ## 8) Logger Config Group (`configs/logger/*.yaml`) @@ -562,90 +549,90 @@ Covered files: | Key | Default | Type | Valid Values | Active When | Example Override | |---|---|---|---|---|---| -| `csv._target_` | `lightning.pytorch.loggers.csv_logs.CSVLogger` | str | import path | `logger=csv` | `python src/train.py logger.csv._target_=...` | -| `csv.save_dir` | `${paths.output_dir}` | str | path | `logger=csv` | `python src/train.py logger.csv.save_dir=/tmp/logs` | -| `csv.name` | `csv/` | str | name | `logger=csv` | `python src/train.py logger.csv.name=metrics` | -| `csv.prefix` | `""` | str | prefix string | `logger=csv` | `python src/train.py logger.csv.prefix=exp1` | +| `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` | `python src/train.py logger.tensorboard._target_=...` | -| `tensorboard.save_dir` | `${paths.output_dir}/tensorboard/` | str | path | `logger=tensorboard` | `python src/train.py logger.tensorboard.save_dir=/tmp/tb` | -| `tensorboard.name` | `null` | str/null | name or null | `logger=tensorboard` | `python src/train.py logger.tensorboard.name=runA` | -| `tensorboard.log_graph` | `False` | bool | `true/false` | `logger=tensorboard` | `python src/train.py logger.tensorboard.log_graph=true` | -| `tensorboard.default_hp_metric` | `True` | bool | `true/false` | `logger=tensorboard` | `python src/train.py logger.tensorboard.default_hp_metric=false` | -| `tensorboard.prefix` | `""` | str | prefix string | `logger=tensorboard` | `python src/train.py logger.tensorboard.prefix=train` | +| `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` | `python src/train.py logger.wandb._target_=...` | -| `wandb.save_dir` | `${paths.output_dir}` | str | path | `logger=wandb` | `python src/train.py logger.wandb.save_dir=/tmp/wb` | -| `wandb.offline` | `False` | bool | `true/false` | `logger=wandb` | `python src/train.py logger.wandb.offline=true` | -| `wandb.id` | `null` | str/null | run ID or null | `logger=wandb` | `python src/train.py logger.wandb.id=my_id` | -| `wandb.anonymous` | `null` | str/null | WandB anonymous mode or null | `logger=wandb` | `python src/train.py logger.wandb.anonymous=allow` | -| `wandb.project` | `grass-mil` | str | project name | `logger=wandb` | `python src/train.py logger.wandb.project=grass-mil` | -| `wandb.log_model` | `False` | bool | `true/false` | `logger=wandb` | `python src/train.py logger.wandb.log_model=true` | -| `wandb.prefix` | `""` | str | prefix string | `logger=wandb` | `python src/train.py logger.wandb.prefix=exp` | -| `wandb.group` | `""` | str | group name | `logger=wandb` | `python src/train.py logger.wandb.group=ablation` | -| `wandb.tags` | `[]` | list[str] | any list | `logger=wandb` | `python src/train.py logger.wandb.tags='[ssl,mil]'` | -| `wandb.job_type` | `""` | str | job type | `logger=wandb` | `python src/train.py logger.wandb.job_type=train` | +| `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` | `python src/train.py logger.mlflow._target_=...` | -| `mlflow.tracking_uri` | `${paths.log_dir}/mlflow/mlruns` | str | URI/path | `logger=mlflow` | `python src/train.py logger.mlflow.tracking_uri=/tmp/mlruns` | -| `mlflow.tags` | `null` | dict/null | tags dict or null | `logger=mlflow` | `python src/train.py logger.mlflow.tags='{stage:dev}'` | -| `mlflow.prefix` | `""` | str | prefix | `logger=mlflow` | `python src/train.py logger.mlflow.prefix=exp` | -| `mlflow.artifact_location` | `null` | str/null | URI/path or null | `logger=mlflow` | `python src/train.py logger.mlflow.artifact_location=/tmp/artifacts` | +| `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` | `python src/train.py logger.comet._target_=...` | -| `comet.api_key` | `${oc.env:COMET_API_TOKEN}` | str | API key/env ref | `logger=comet` | `python src/train.py logger.comet.api_key=...` | -| `comet.save_dir` | `${paths.output_dir}` | str | path | `logger=comet` | `python src/train.py logger.comet.save_dir=/tmp/comet` | -| `comet.project_name` | `grass-mil` | str | project name | `logger=comet` | `python src/train.py logger.comet.project_name=grass-mil` | -| `comet.rest_api_key` | `null` | str/null | key or null | `logger=comet` | `python src/train.py logger.comet.rest_api_key=...` | -| `comet.experiment_key` | `null` | str/null | key or null | `logger=comet` | `python src/train.py logger.comet.experiment_key=...` | -| `comet.offline` | `False` | bool | `true/false` | `logger=comet` | `python src/train.py logger.comet.offline=true` | -| `comet.prefix` | `""` | str | prefix | `logger=comet` | `python src/train.py logger.comet.prefix=train` | +| `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` | `python src/train.py logger.neptune._target_=...` | -| `neptune.api_key` | `${oc.env:NEPTUNE_API_TOKEN}` | str | API key/env ref | `logger=neptune` | `python src/train.py logger.neptune.api_key=...` | -| `neptune.project` | `org/grass-mil` | str | project slug | `logger=neptune` | `python src/train.py logger.neptune.project=org/project` | -| `neptune.log_model_checkpoints` | `True` | bool | `true/false` | `logger=neptune` | `python src/train.py logger.neptune.log_model_checkpoints=false` | -| `neptune.prefix` | `""` | str | prefix | `logger=neptune` | `python src/train.py logger.neptune.prefix=exp` | +| `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` | `python src/train.py logger.aim._target_=...` | -| `aim.repo` | `${paths.root_dir}` | str | path/remote URI | `logger=aim` | `python src/train.py logger.aim.repo=/tmp/.aim` | -| `aim.experiment` | `null` | str/null | experiment name or null | `logger=aim` | `python src/train.py logger.aim.experiment=baseline` | -| `aim.train_metric_prefix` | `train/` | str | prefix | `logger=aim` | `python src/train.py logger.aim.train_metric_prefix=train` | -| `aim.val_metric_prefix` | `val/` | str | prefix | `logger=aim` | `python src/train.py logger.aim.val_metric_prefix=val` | -| `aim.test_metric_prefix` | `test/` | str | prefix | `logger=aim` | `python src/train.py logger.aim.test_metric_prefix=test` | -| `aim.system_tracking_interval` | `10` | int/null | seconds or null | `logger=aim` | `python src/train.py logger.aim.system_tracking_interval=5` | -| `aim.log_system_params` | `true` | bool | `true/false` | `logger=aim` | `python src/train.py logger.aim.log_system_params=false` | -| `aim.capture_terminal_logs` | `false` | bool | `true/false` | `logger=aim` | `python src/train.py logger.aim.capture_terminal_logs=true` | +| `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` | `python src/train.py logger=many_loggers logger.defaults='[csv,tensorboard]'` | +| `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`) @@ -653,18 +640,18 @@ Covered files: | Key | Default | Type | Valid Values | Active When | Example Override | |---|---|---|---|---|---| -| `task_name` | `debug` | str | any string | `debug=default` | `python src/train.py debug=default task_name=dbg` | -| `callbacks` | `null` | null/config | null or callback config | `debug=default` | `python src/train.py debug=default callbacks=default` | -| `logger` | `null` | null/config | null or logger config | `debug=default` | `python src/train.py debug=default logger=csv` | -| `extras.ignore_warnings` | `False` | bool | `true/false` | `debug=default` | `python src/train.py debug=default extras.ignore_warnings=true` | -| `extras.enforce_tags` | `False` | bool | `true/false` | `debug=default` | `python src/train.py debug=default extras.enforce_tags=false` | -| `hydra.job_logging.root.level` | `DEBUG` | str | logging level | `debug=default` | `python src/train.py debug=default hydra.job_logging.root.level=INFO` | -| `trainer.max_epochs` | `1` | int | `>=1` | `debug=default` | `python src/train.py debug=default trainer.max_epochs=2` | -| `trainer.accelerator` | `cpu` | str | accelerator value | `debug=default` | `python src/train.py debug=default trainer.accelerator=cpu` | -| `trainer.devices` | `1` | int | `>=1` | `debug=default` | `python src/train.py debug=default trainer.devices=1` | -| `trainer.detect_anomaly` | `true` | bool | `true/false` | `debug=default` | `python src/train.py debug=default trainer.detect_anomaly=true` | -| `data.num_workers` | `0` | int | `>=0` | `debug=default` | `python src/train.py debug=default data.num_workers=0` | -| `data.pin_memory` | `False` | bool | `true/false` | `debug=default` | `python src/train.py debug=default data.pin_memory=false` | +| `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 @@ -677,48 +664,48 @@ Covered files: | Key | Default | Type | Valid Values | Active When | Example Override | |---|---|---|---|---|---| -| `debug/fdr.defaults` | `[default]` | list | debug group list | `debug=fdr` | `python src/train.py debug=fdr` | -| `debug/fdr.trainer.fast_dev_run` | `true` | bool | `true/false` | `debug=fdr` | `python src/train.py debug=fdr trainer.fast_dev_run=true` | -| `debug/limit.defaults` | `[default]` | list | debug group list | `debug=limit` | `python src/train.py debug=limit` | -| `debug/limit.trainer.max_epochs` | `3` | int | `>=1` | `debug=limit` | `python src/train.py debug=limit trainer.max_epochs=5` | -| `debug/limit.trainer.limit_train_batches` | `0.01` | float | `(0,1]` or int | `debug=limit` | `python src/train.py debug=limit trainer.limit_train_batches=0.1` | -| `debug/limit.trainer.limit_val_batches` | `0.05` | float | `(0,1]` or int | `debug=limit` | `python src/train.py debug=limit trainer.limit_val_batches=0.2` | -| `debug/limit.trainer.limit_test_batches` | `0.05` | float | `(0,1]` or int | `debug=limit` | `python src/train.py debug=limit trainer.limit_test_batches=0.2` | -| `debug/overfit.defaults` | `[default]` | list | debug group list | `debug=overfit` | `python src/train.py debug=overfit` | -| `debug/overfit.trainer.max_epochs` | `20` | int | `>=1` | `debug=overfit` | `python src/train.py debug=overfit trainer.max_epochs=10` | -| `debug/overfit.trainer.overfit_batches` | `3` | int | `>=1` | `debug=overfit` | `python src/train.py debug=overfit trainer.overfit_batches=1` | -| `debug/overfit.callbacks` | `null` | null/config | null or callback config | `debug=overfit` | `python src/train.py debug=overfit callbacks=none` | -| `debug/profiler.defaults` | `[default]` | list | debug group list | `debug=profiler` | `python src/train.py debug=profiler` | -| `debug/profiler.trainer.max_epochs` | `1` | int | `>=1` | `debug=profiler` | `python src/train.py debug=profiler trainer.max_epochs=2` | -| `debug/profiler.trainer.profiler` | `simple` | str | `simple`, `advanced`, `pytorch`, etc. | `debug=profiler` | `python src/train.py debug=profiler trainer.profiler=advanced` | +| `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 | `python src/train.py 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 | `python src/train.py 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 | `python src/train.py -m hydra.sweep.dir=/tmp/multi` | -| `hydra.sweep.subdir` | `${hydra.job.num}` | str | template string | multirun | `python src/train.py -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 | `python src/train.py hydra.job_logging.handlers.file.filename=/tmp/train.log` | +| `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 | `python src/train.py paths.root_dir=/abs/project` | -| `paths.data_dir` | `${paths.root_dir}/data/` | str | path | always | `python src/train.py paths.data_dir=/abs/data` | -| `paths.log_dir` | `${paths.root_dir}/logs/` | str | path | always | `python src/train.py paths.log_dir=/abs/logs` | -| `paths.output_dir` | `${hydra:runtime.output_dir}` | str | hydra runtime path | always | `python src/train.py paths.output_dir=/tmp/out` | -| `paths.work_dir` | `${hydra:runtime.cwd}` | str | working directory path | always | `python src/train.py paths.work_dir=/tmp` | +| `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 | `python src/train.py extras.ignore_warnings=true` | -| `extras.enforce_tags` | `True` | bool | `true/false` | `extras` enabled | `python src/train.py extras.enforce_tags=false` | -| `extras.print_config` | `True` | bool | `true/false` | `extras` enabled | `python src/train.py extras.print_config=false` | +| `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` | ## 13) Repo-Enforced Behavioral Constraints (From Code) @@ -726,19 +713,16 @@ These constraints are not just documentation conventions; they are enforced in r | Constraint | Enforced In | Behavior | |---|---|---| -| `task.target_type` must be one of `binary`, `regression`, `survival` | `src/models/training/builders.py` | raises `ValueError` for unsupported target types | -| `task.instance_sampling` must be `all` or `random` | `src/models/training/builders.py` | raises `ValueError` for invalid sampling mode | -| `task.region_accumulation.hyperbatch_size >= 1` when enabled | `src/models/supervised_module.py` | raises `ValueError` | -| `task.node_aux.target_mode` must be `graph_label` or `attention_shaped_ti` | `src/models/training/loss_utils.py` | raises `ValueError` | -| `task.node_aux.loss_mode` must be `bce` or `weighted_bce` | `src/models/training/loss_utils.py` | raises `ValueError` | -| `task.entropy_reg.mode` must be `attention` or `attention_shaped_target` | `src/models/supervised_module.py` | raises `ValueError` | -| `conv_type=gine` requires `model.encoder.edge_attr_dim` and runtime `edge_attr` | `src/models/components/backbones.py` | raises `ValueError` | -| `conv_type=gat` requires `hidden_dim % gat_heads == 0` | `src/models/components/backbones.py` | raises `ValueError` | -| `model.encoder.use_edge_attr=true` is meaningful only for `gcn`/`gine` | `src/models/components/backbones.py` | warning/ignored for unsupported convs | -| `data.split.train_val_test_split` must sum to `1.0` | `src/data/components/precompute.py` split helpers | raises `ValueError` | -| `data.feature_reducer.fit_mode` must be `global` or `train_only` | `src/data/components/precompute.py` | raises `ValueError` | -| `data.sampler.runtime.weight_mode` must be `inverse`, `sqrt_inverse`, or `proportional` | `src/data/components/samplers.py` | raises `ValueError` | -| `model.encoder_init_map` must be `identity` or `auto_bgrl_or_identity` | `src/models/training/checkpoint_init.py` | raises `ValueError` | +| `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 @@ -750,7 +734,7 @@ Hydra still allows passing additional Lightning kwargs that are not pre-declared Example: ```bash -python src/train.py trainer.log_every_n_steps=1 trainer.enable_progress_bar=true +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 index 6cfbff7..e1dedf2 100644 --- a/docs/inference_utility_suite.md +++ b/docs/inference_utility_suite.md @@ -1,7 +1,7 @@ # Inference Utility Suite This document describes the artifact-oriented inference workflow implemented in -`src/inference/predict.py`. +`src/grass_mil/inference/predict.py`. ## Purpose @@ -16,7 +16,7 @@ The suite extends checkpoint evaluation with: ## CLI Entry Point ```bash -python src/inference/predict.py \ +grass-mil-predict \ ckpt_path=/absolute/path/to/checkpoint.ckpt \ task=finetune_mil \ model=supervised_module \ @@ -38,7 +38,7 @@ Sub-config groups: `aggregation.subsample_fraction` supports efficient inference-time subsampling: -- when `< 1.0`, `src/inference/predict.py` injects predict-time sampler runtime +- 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` @@ -154,8 +154,8 @@ Unit contract: High-value legacy utilities from `working_version` are mapped as: - prediction collectors (`collect_predict_for_all_nodes*`) -> unified - `collect_inference_payload` in `src/inference/collectors.py` -- full-graph aggregation (`full_graph_*`) -> `src/inference/aggregation.py` -- evaluation metrics (`graph_*_evaluate_fn`) -> `src/inference/metrics.py` -- embedding collectors -> `src/inference/collectors.py` + `src/inference/embeddings.py` -- interpretability table bridge -> `src/inference/interpretability_export.py` + `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/embeddings.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 index 23675ff..051ad65 100644 --- a/docs/interpretability_extension_guide.md +++ b/docs/interpretability_extension_guide.md @@ -31,8 +31,8 @@ Examples: Create a module: -- `src/interpretability/core/.py` -- or `src/interpretability/tier2/.py` +- `src/grass_mil/interpretability/core/.py` +- or `src/grass_mil/interpretability/tier2/.py` Guidelines: @@ -52,7 +52,7 @@ def compute_spatial_entropy( node_table: pd.DataFrame, spatial_table: pd.DataFrame, *, - label_column: str = "cluster_label", + label_column: str = "niche_label", ) -> pd.DataFrame: if label_column not in node_table.columns: raise ValueError(f"Missing label column {label_column!r}.") @@ -64,12 +64,12 @@ def compute_spatial_entropy( Add plugin class in: -- `src/interpretability/plugins/builtin.py` +- `src/grass_mil/interpretability/plugins/builtin.py` - or a new plugin module imported at registration time. Plugin responsibilities: -1. Read context state (e.g., cluster labels). +1. Read context state (e.g., niche labels). 2. Call your pure analysis function. 3. Return `PluginResult` with: - `payload` for machine consumption @@ -83,15 +83,15 @@ class SpatialEntropyPlugin(InterpretabilityPlugin): name: str = "spatial_entropy" def required_inputs(self) -> list[str]: - return ["spatial_table", "cluster_labels"] + return ["spatial_table", "niche_labels"] def run(self, dataset, context: PluginContext, **params): table = dataset.instance_table.copy() - table["cluster_label"] = context.state["cluster_labels"] + table["niche_label"] = context.state["niche_labels"] entropy_df = compute_spatial_entropy( table, dataset.spatial_table, - label_column=str(params.get("label_column", "cluster_label")), + label_column=str(params.get("label_column", "niche_label")), ) return PluginResult( name=self.name, @@ -99,7 +99,7 @@ class SpatialEntropyPlugin(InterpretabilityPlugin): sections=[ ReportSection( title="Spatial Entropy", - description="Entropy by cluster.", + description="Entropy by niche.", tables={"entropy": entropy_df}, ) ], @@ -141,11 +141,11 @@ Example: ```yaml enabled: - - cluster_profiles + - niche_profiles - spatial_entropy params: spatial_entropy: - label_column: cluster_label + label_column: niche_label radius: 50.0 ``` @@ -153,7 +153,7 @@ params: If table-only output is insufficient: -1. Add figure builder in `src/interpretability/reporting/plotly_builders.py`. +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. diff --git a/docs/interpretability_suite.md b/docs/interpretability_suite.md index 0bece28..6c3e529 100644 --- a/docs/interpretability_suite.md +++ b/docs/interpretability_suite.md @@ -28,7 +28,7 @@ Key principles: Package roots: -- `src/interpretability/` +- `src/grass_mil/interpretability/` - `configs/interpretability/` - `tests/interpretability/` @@ -37,34 +37,34 @@ Package roots: Core layers: 1. Data contracts and loaders -- `src/interpretability/contracts.py` -- `src/interpretability/core/data.py` +- `src/grass_mil/interpretability/contracts.py` +- `src/grass_mil/interpretability/core/data.py` 2. Core numerical analysis -- `src/interpretability/core/reduction.py` -- `src/interpretability/core/clustering.py` -- `src/interpretability/core/biomarkers.py` +- `src/grass_mil/interpretability/core/reduction.py` +- `src/grass_mil/interpretability/core/clustering.py` +- `src/grass_mil/interpretability/core/biomarkers.py` 3. Tier-2 analyses -- `src/interpretability/tier2/neighborhood.py` -- `src/interpretability/tier2/filtration.py` +- `src/grass_mil/interpretability/tier2/neighborhood.py` +- `src/grass_mil/interpretability/tier2/filtration.py` 4. Plugin orchestration -- `src/interpretability/plugins/base.py` -- `src/interpretability/plugins/registry.py` -- `src/interpretability/plugins/builtin.py` -- `src/interpretability/pipeline.py` +- `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/interpretability/reporting/plotly_builders.py` -- `src/interpretability/reporting/sections.py` -- `src/interpretability/reporting/html.py` -- `src/interpretability/reporting/snapshot.py` -- `src/interpretability/reporting/pdf.py` -- `src/interpretability/reporting/render.py` +- `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/interpretability/report_cli.py` +- `src/grass_mil/interpretability/report_cli.py` ## 3) Normalized Input Contracts @@ -102,7 +102,7 @@ Common optional columns: - `instance_attention_logits` export requires shape `(N, 1)`; multi-column attention logits are rejected as ambiguous. -Note: `cluster_profiles` and `attention_attribution` require `comp_*` columns. +Note: `niche_profiles` requires `comp_*` columns. There is no fallback to one-hot `cell_type` composition. `variance_estimator` convention (biomarker z-scoring): @@ -146,10 +146,10 @@ Public exports: - `compute_filtration_curves(...)` - `prepare_tissue_graph_view(...)` - `build_tissue_graph_figure(...)` -- `fit_cluster_transfer_from_report_bundle(...)` -- `apply_cluster_transfer(...)` -- `save_cluster_transfer_bundle(...)` -- `load_cluster_transfer_bundle(...)` +- `fit_niche_transfer_from_report_bundle(...)` +- `apply_niche_transfer(...)` +- `save_niche_transfer_bundle(...)` +- `load_niche_transfer_bundle(...)` - `run_interpretability_pipeline(...)` - `render_interpretability_report(...)` @@ -161,8 +161,8 @@ from src.interpretability import ( run_clustering, prepare_tissue_graph_view, build_tissue_graph_figure, - fit_cluster_transfer_from_report_bundle, - apply_cluster_transfer, + fit_niche_transfer_from_report_bundle, + apply_niche_transfer, run_interpretability_pipeline, render_interpretability_report, ) @@ -245,7 +245,7 @@ Common parameters: Returns a `ClusteringResult` dataclass with: -- cluster labels +- niche labels - resolved parameters - fitted object @@ -271,8 +271,8 @@ Plugin registration is per pipeline run (fresh `PluginRegistry`). Built-in plugins: -- `cluster_profiles` -- `attention_attribution` +- `niche_profiles` +- `margin_attribution` - `neighborhood_enrichment` - `diff_neighborhood_enrichment` - `filtration_curves` @@ -288,7 +288,7 @@ Built-in plugins: - 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 `cluster_labels` context +- colors nodes by pipeline `niche_labels` context - uses `spatial_table` edges filtered to selected sample nodes ## 7) Pipeline Orchestration @@ -300,14 +300,14 @@ Workflow: 1. Load embedding matrix from `instance_table`. 2. Run configured reduction. 3. Run configured clustering: -- `cluster_on_pca=true`: cluster on PCA projection used for clustering. -- `cluster_on_pca=false`: cluster on raw embedding columns. -4. Build core cluster summaries. +- `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_clusters.csv` +- `artifacts/instance_with_niches.csv` - `artifacts/reduction.csv` (if enabled) -- `artifacts/cluster_labels.csv` (if enabled) +- `artifacts/niche_labels.csv` (if enabled) - `artifacts/pipeline_summary.json` 7. Return `ReportBundle`. @@ -316,11 +316,11 @@ Workflow: - populated with clustering PCA result when `cluster_on_pca=true` - `None` when clustering runs on raw embeddings -## 8) Cluster Transfer Workflow (Library API) +## 8) Niche Transfer Workflow (Library API) -Notebook-parity cluster transfer is available as a library workflow: +Notebook-parity niche transfer is available as a library workflow: -1. Fit transfer bundle on a source interpretability run (cluster labels + clustering feature space). +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. @@ -338,14 +338,14 @@ from pathlib import Path import pandas as pd -from src.interpretability.core.data import load_interpretability_dataset -from src.interpretability.core.transfer import ( - apply_cluster_transfer, - fit_cluster_transfer_from_report_bundle, - load_cluster_transfer_bundle, - save_cluster_transfer_bundle, +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 src.interpretability.pipeline import run_interpretability_pipeline +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( @@ -363,18 +363,18 @@ bundle = run_interpretability_pipeline( }, artifacts_dir=Path("/abs/source_artifacts"), ) -transfer = fit_cluster_transfer_from_report_bundle(dataset, bundle, n_neighbors=15) -save_cluster_transfer_bundle(transfer, "/abs/cluster_transfer_bundle.joblib") +transfer = fit_niche_transfer_from_report_bundle(dataset, bundle, n_neighbors=15) +save_niche_transfer_bundle(transfer, "/abs/cluster_transfer_bundle.joblib") -loaded = load_cluster_transfer_bundle("/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_cluster_transfer(loaded, query_table) +transferred = apply_niche_transfer(loaded, query_table) ``` Notebook helper example for single tissue graph: ```python -from src.interpretability.tier2.tissue_graph import ( +from grass_mil.interpretability.tier2.tissue_graph import ( build_tissue_graph_figure, prepare_tissue_graph_view, ) @@ -387,7 +387,7 @@ view = prepare_tissue_graph_view( id_column="instance_id", x_column="center_x", y_column="center_y", - label_column="cluster_label", + label_column="niche_label", ) fig = build_tissue_graph_figure(view, show_edges=True) fig.show() @@ -424,7 +424,7 @@ Engine: Generate interpretability-ready tables directly from inference: ```bash -python src/inference/predict.py \ +grass-mil-predict \ ckpt_path=/abs/path/model.ckpt \ data=spatial_omics \ interpretability.enabled=true \ @@ -438,23 +438,23 @@ This writes `instance_table.csv` (and optionally `spatial_table.csv`) into Entrypoint: ```bash -python src/interpretability/report_cli.py \ +grass-mil-report \ data.instance_table=/abs/path/instance_table.csv ``` With spatial plugins: ```bash -python src/interpretability/report_cli.py \ +grass-mil-report \ data.instance_table=/abs/path/instance_table.csv \ data.spatial_table=/abs/path/spatial_table.csv \ - plugins.enabled="[cluster_profiles,attention_attribution,neighborhood_enrichment,diff_neighborhood_enrichment,filtration_curves]" + plugins.enabled="[niche_profiles,margin_attribution,neighborhood_enrichment,diff_neighborhood_enrichment,filtration_curves]" ``` -Render a single selected tissue graph colored by cluster labels: +Render a single selected tissue graph colored by niche labels: ```bash -python src/interpretability/report_cli.py \ +grass-mil-report \ data.instance_table=/abs/path/instance_table.csv \ data.spatial_table=/abs/path/spatial_table.csv \ plugins.enabled="[tissue_graph]" \ @@ -465,7 +465,7 @@ python src/interpretability/report_cli.py \ Override algorithms: ```bash -python src/interpretability/report_cli.py \ +grass-mil-report \ data.instance_table=/abs/path/instance_table.csv \ reduction=umap \ clustering=dbscan \ @@ -476,7 +476,7 @@ python src/interpretability/report_cli.py \ Tune HDBSCAN selection behavior explicitly: ```bash -python src/interpretability/report_cli.py \ +grass-mil-report \ data.instance_table=/abs/path/instance_table.csv \ clustering=hdbscan \ clustering.params.min_cluster_size=100 \ @@ -488,7 +488,7 @@ python src/interpretability/report_cli.py \ Enable/disable report outputs: ```bash -python src/interpretability/report_cli.py \ +grass-mil-report \ data.instance_table=/abs/path/instance_table.csv \ report.pdf.enabled=true \ report.html.enabled=true @@ -540,7 +540,7 @@ Parameter contract: From project root: ```bash -pip install -r requirements.txt +pip install -e ".[interpretability,report,spatial-stats,io]" ``` For PDF export, install Playwright browser binaries: @@ -616,3 +616,26 @@ Fix: - 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. diff --git a/docs/legacy_gap_audit.md b/docs/legacy_gap_audit.md index 760a0ba..752a552 100644 --- a/docs/legacy_gap_audit.md +++ b/docs/legacy_gap_audit.md @@ -1,58 +1,74 @@ -# Legacy Gap Audit - -This project intentionally does not copy legacy implementation details directly. -Instead, it ports validated concepts into a clean contract-first architecture. - -## Resolved Legacy Risks - -1. Broken module registration patterns - -- Risk: modules stored in plain Python lists are not registered in `nn.Module`. -- Resolution: always use `nn.ModuleList`/`nn.ModuleDict` for learnable submodules. - -2. Fragile attention assumptions - -- Risk: implicit transpose/shape assumptions lead to silent logic errors. -- Resolution: explicit input-output shape contracts and attention normalization controlled by caller. - -3. Hidden mutable dataset/sampler coupling - -- Risk: mutating dataset-global indices/cache from training code introduces race conditions and stale state. -- Resolution: sampler strategies operate on provided graph units only, no mutation of datamodule internals. - -4. Hardcoded script behavior - -- Risk: project, path, split, and task logic hardcoded in scripts prevents reuse. -- Resolution: all component behavior and toggles are Hydra-configurable. - -5. Over-customized message passing - -- Risk: custom conv implementations increase maintenance burden and bug surface. -- Resolution: use native PyG layers by default; add only thin adapters when strictly needed. - -6. Inconsistent loss interfaces - -- Risk: different tasks using incompatible shape conventions. -- Resolution: standardized loss interfaces with validation for categorical, regression, and Cox objectives. - -7. Optional feature leakage - -- Risk: attention/SSL paths accidentally active by default due to implicit construction. -- Resolution: explicit config flags (`use_attention`, `use_ssl`) govern instantiation. - -8. LOOCV split leakage in legacy scripts - -- Risk: legacy LOOCV scripts computed class weights and dataset-wide metadata before or across fold boundaries, allowing held-out fold information to influence training. -- Resolution: configurable LOOCV split assignment is performed in precompute, reducer `train_only` fitting is guarded for fold-specific preprocessing, and validation strategy is explicit (`heldout_fold_items` or `patches_from_train_items`) with deterministic seeds. - -9. Fragile cluster transfer state assumptions - -- Risk: notebook-era transfer used implicit estimator internals (for example `clusterer._embedding_`) that are absent for some clustering backends and break silently across versions. -- Resolution: explicit transfer bundle contract persists projector + kNN model + metadata, and transfer fitting reads clustering feature-space outputs from `ReportBundle`. - -## Acceptance Criteria - -- Component modules are side-effect free. -- Sampler API is data-layer agnostic. -- Optional modules instantiate only when enabled. -- Tests cover shape contracts and native/custom sampler parity on synthetic inputs. +# Published-Method Fidelity Status + +This document used to assert that a list of "legacy risks" had been resolved by +rebuilding from concepts rather than porting code. That framing hid a real +problem: several load-bearing pieces of the published method were not +implemented here at all, and the audit could not have caught it, because it +checked engineering hygiene rather than fidelity to the method. + +It is now a status record. For the authoritative mapping of method → code → +test, see [`method_fidelity.md`](method_fidelity.md). + +## Gaps that were open and are now closed + +1. **The encoder could not express the published input scheme.** There was no + cell-type embedding, so cell identity — the only node input for some cohorts + — had no path into the model. Closed by `NodeInputEmbedding`. +2. **Edge conditioning was wrong.** `gine` fed the whole edge-attribute vector + into the convolution. The production models condition on the scalar edge + length. Closed by `EncoderConfig.edge_feature_index`. +3. **Attention was single-channel.** Bag pooling squeezed the head to one + channel broadcast across all class logits, so the per-class formulation — + and therefore the exact attribution identity — could not be expressed. + Closed; the identity is now asserted at `atol=1e-12`. +4. **Sampling controls were missing.** No ego-graph radius cutoff and no + interior seeding, so ego-graphs rooted at a section edge were truncated by + the boundary rather than by biology. +5. **SSL views were uniform.** Elementwise masking instead of the + degree-importance drop plus cell-size noise. +6. **No attribution.** Attention lift existed; margins, the identity check and + uncertainty intervals did not. +7. **No spatial statistics.** Moran's I, Ripley's cross-L and cross-space + agreement had no implementation. + +## Bugs found while closing them + +These produced silently wrong results rather than failures, which is why they +survived: + +1. **Sampling discarded cell type.** `_apply_transform_to_each_subgraph` dropped + every node-length tensor as "already represented by `x`" — false for + `categorical_index` and `pos`. +2. **Node codes were batched along the wrong axis.** PyTorch Geometric treats + any attribute whose name contains `index` as an edge-index tensor, so + `categorical_index` was concatenated along the feature dimension. Renamed to + `categorical_codes`; graphs from older runs are migrated on load. +3. **Attention flattening in two directions.** `inference/aggregation.py` + silently reshaped multi-column attention to 1-D while + `interpretability_export.py` raised on it. One gave wrong numbers, the other + crashed. +4. **Asymmetric edge dropping.** The finetune-time transform left + `force_undirected=False` while pretraining used `True`, so the two regimes + augmented differently and message passing became direction-dependent. +5. **Order-dependent Ripley radii.** The radius grid came from whichever region + was iterated first. +6. **A concordance index that synchronised per pair.** O(n²) Python with two + `.item()` calls per pair. + +## Engineering invariants (the original content, still enforced) + +- Learnable submodules live in `nn.ModuleList`/`nn.ModuleDict`. +- Attention shape contracts are explicit; normalisation is caller-controlled and + validated by `grass_mil.contracts.validate_attention_width`. +- Samplers do not mutate datamodule state. +- Behaviour is config-driven, not hardcoded per project. The legacy + `os.environ` control surface is deliberately not ported. +- Native PyG operators are preferred; adapters stay thin. +- Optional components instantiate only when their config is set. + +## Deliberately retained + +- Region accumulation. It is a throughput mechanism rather than an objective, + and the final NSCLC runs use it (`hyperbatch_size: 16`), so it stays. +- The clustering algorithm's own names (`run_clustering`, `n_clusters`, + `min_cluster_size`). Clustering is the method; niches are its result. diff --git a/docs/method_fidelity.md b/docs/method_fidelity.md new file mode 100644 index 0000000..abe44ab --- /dev/null +++ b/docs/method_fidelity.md @@ -0,0 +1,78 @@ +# Method Fidelity + +Where each piece of the GRASS-MIL method lives in this package, and the test +that holds it in place. This is the map to check before trusting a number. + +**The run code is the reference, not the manuscript.** Where `methods.tex` and +the final NSCLC/LUAD run code disagree, the code wins; the known divergences are +listed at the bottom. + +## Core method + +| Component | Definition | Implementation | Test | +|---|---|---|---| +| Cellular graph | Voronoi/Delaunay adjacency, edge distance + neighbour flag | `data/components/graph_builders.py` | `tests/data/` | +| Node inputs | `h⁰ = E[t] + Wₓx` — cell-type embedding **summed** with a projection of continuous features | `models/components/embeddings.py::NodeInputEmbedding` | `test_node_embeddings.py::test_node_embedding_sums_rather_than_concatenates` | +| Cell-type-only cohorts | `x` has zero columns; the embedding is the whole input | same | `test_node_embedding_supports_cell_type_only_cohorts`, `test_encoder_trains_on_cell_type_only_cohort` | +| Message passing | GINE conditioned on the **scalar edge length** only | `models/components/backbones.py::_select_edge_features` | `test_gine_reads_only_the_selected_edge_column` | +| Ego-graph extent | `r_k = radius_per_hop·k + radius_offset` | `data/components/ego_radius.py` | `test_seed_sampling.py::test_resolve_ego_radius_is_linear_in_depth` | +| Interior seeding | exclude the `n_hops` rim around the convex hull | `data/components/seed_sampling.py::convex_hull_interior_mask` | `test_interior_mask_excludes_the_hull_rim` | +| MIL attention | gated attention, softmax over instances **per class** (a binary head has two channels) | `models/components/attention.py`, `models/training/bagging.py` | `test_bagging.py::test_attention_is_normalised_per_class_over_instances` | +| Bag logit | `L_c = Σᵢ A[i,c]·ℓ[i,c]` | `models/training/bagging.py::aggregate_bag_logits_attention` | `test_bag_logit_is_the_attention_weighted_sum_of_instance_logits` | +| Objective | region cross-entropy, **sole** loss | `models/supervised_module.py::_compute_losses` | `tests/test_train_regimes.py` | +| Survival | Cox partial likelihood on a scalar log-hazard | `models/components/losses.py::CoxSGDLoss` | `test_survival_regime_fast_dev_run` | +| SSL pretraining | BGRL, symmetric cosine loss, EMA target | `models/bgrl_module.py` | `tests/test_train_regimes.py` | +| SSL views | degree-importance node/edge drop + Gaussian noise on cell size | `models/training/augmentations.py` | `tests/models/test_augmentations.py` | + +## Attribution + +The interpretability claim rests on one identity. If `identity_residual` is not +~0, every niche-level attribution derived from it is invalid. + +| Quantity | Definition | Implementation | Test | +|---|---|---|---| +| Instance contribution | `M[i,c] = A[i,c]·(ℓ[i,c] − βc)` — bias-free attention-weighted logits | `interpretability/core/attribution.py` | `test_bias_removal_shifts_the_margin_by_exactly_the_bias_term` | +| Margin (any C) | one-vs-rest: `m[i,c] = M[i,c] − mean(M[i,c'≠c])`; reduces to `M[i,1] − M[i,0]` when C=2 | `instance_ovr_margins` | `test_ovr_margin_reduces_to_the_binary_margin_when_c_is_two`, `test_ovr_margin_contrasts_each_class_against_the_mean_of_the_rest` | +| Additive identity | `Σᵢ M[i,c] = L_c` | same, `identity_residual` | `test_contributions_sum_exactly_to_the_bag_logit` | +| Uncertainty | 200-resample percentile bootstrap **over regions** | `percentile_bootstrap_ci` | `test_bootstrap_ci_brackets_the_point_estimate` | + +## Spatial statistics + +| Statistic | Implementation | Test | +|---|---|---| +| Moran's I + permutation null + BH-FDR | `interpretability/tier2/autocorrelation.py` | `test_spatial_statistics.py::test_morans_i_*` | +| Ripley cross-L, centred `L(r) − r` | `interpretability/tier2/ripley.py` | `test_ripley_l_is_near_zero_for_a_poisson_pattern` | +| Neighbourhood enrichment | `interpretability/tier2/neighborhood.py` | `test_plugins_and_tier2.py` | +| Cross-space agreement (ARI/AMI/NMI) | `interpretability/core/agreement.py` | `test_niche_agreement_*` | + +## Known divergences from `methods.tex` + +The manuscript is stale on these points. Each is deliberate. + +1. **Edge features.** `methods.tex` describes a learned edge-**type** embedding + propagated through the conv, with the distance channel supplied but unread. + The implemented models do the opposite: GINE on the scalar edge **length**, + no edge-type notion. *The manuscript needs updating.* +2. **Degree normalisation.** `methods.tex` says min-max; the run code uses + mean-max (`(v − mean)/(max − mean)`, clamped at 0), which maps every + below-average-degree node to zero importance. Implemented as in the code. +3. **Signed margin share bounds.** `methods.tex` states this is bounded in + `[-1, 1]`. It is not: it is a fraction of the bag margin, so one niche can + exceed 1 when another opposes it. The real invariant is that a full + partition sums to 1. +4. **Moran's I scope.** 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 — a defensible and + arguably cleaner object, but not numerically the same statistic. + +## Not implemented + +Claimed in `methods.tex` but absent here, and absent from the legacy code too. +Implement or remove before submission: + +- Archetypal analysis and the three simplex tests that depend on it. +- The 10×10 ensemble, ICC(2,1)/ICC(2,k), and the two-way ANOVA variance + decomposition. All training uses a single seed; there is no ensemble harness. +- Graphlet/motif analysis and its permutation null. +- Paired bootstrap of the cross-L condition difference (the contrast is two + separately aggregated mean curves). diff --git a/docs/model_component_contracts.md b/docs/model_component_contracts.md index 168821c..8d666d1 100644 --- a/docs/model_component_contracts.md +++ b/docs/model_component_contracts.md @@ -27,7 +27,7 @@ Optional fields: - `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_index`: `torch.LongTensor` for categorical annotations. +- `categorical_codes`: `torch.LongTensor` for categorical annotations. Non-requirements: @@ -114,3 +114,35 @@ Native-first policy: - 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). + +## Encoder Contract (updated) + +- `forward(x, edge_index, edge_attr=None, batch=None, categorical_codes=None, + return_graph_embedding=False)`. +- `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 returns `[n_items, n_classes]`. Callers apply softmax over + `dim=0` (instances) **per class column**. +- Only two widths are legal: 1 (one shared channel) or one per class. Enforced + by `grass_mil.contracts.validate_attention_width`, at build time and at use. +- `head_spaces()` exposes intermediate activations (`a_pre_tanh`, + `gate_product`) so analysis code does not reach into module internals. + +## 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 `cluster_attribution_summary` reports `identity_residual` so a +violation surfaces rather than propagating into the cluster summaries. diff --git a/docs/repository_map.md b/docs/repository_map.md index a8b19d5..cb14edb 100644 --- a/docs/repository_map.md +++ b/docs/repository_map.md @@ -17,7 +17,7 @@ Scope rules used here: - `LICENSE`: Project license text. - `README.md`: Primary entrypoint documentation and navigation index. - `Makefile`: Convenience targets for clean, format, tests, and default training. -- `requirements.txt`: Pip dependency list, including optional data/runtime libraries. +- `pyproject.toml`: Package metadata, dependencies and optional extras. - `environment.yaml`: Conda environment alternative with pinned major versions. - `setup.py`: Package metadata and console entry points (`train_command`, `eval_command`). - `pyproject.toml`: Pytest and coverage tool configuration. @@ -125,7 +125,6 @@ Attention presets: Head presets: - `configs/model/heads/graph.yaml`: Graph-level prediction MLP head preset. -- `configs/model/heads/node.yaml`: Node-level prediction MLP head preset. Loss presets: @@ -160,71 +159,73 @@ SSL preset: - `configs/trainer/ddp.yaml`: Multi-GPU DDP preset. - `configs/trainer/ddp_sim.yaml`: CPU `ddp_spawn` simulation preset. -## Source Code (`src/`) +## Source Code (`src/grass_mil/`) ### Package Marker -- `src/__init__.py`: Source package marker. +- `src/grass_mil/__init__.py`: Package marker. +- `src/grass_mil/_paths.py`: Best-effort `PROJECT_ROOT` resolution for the CLI entrypoints. +- `src/grass_mil/configs/`: Hydra config tree, shipped as package data. ### Entrypoints -- `src/train.py`: Hydra training entrypoint (train + optional test). -- `src/eval.py`: Hydra evaluation entrypoint (test from checkpoint). -- `src/loocv.py`: Hydra LOOCV entrypoint with fold discovery, per-fold runs, and aggregate summaries. +- `src/grass_mil/train.py`: Hydra training entrypoint (train + optional test). +- `src/grass_mil/eval.py`: Hydra evaluation entrypoint (test from checkpoint). +- `src/grass_mil/loocv.py`: Hydra LOOCV entrypoint with fold discovery, per-fold runs, and aggregate summaries. -### Data Pipeline (`src/data/`) +### Data Pipeline (`src/grass_mil/data/`) -- `src/data/__init__.py`: Data package marker. -- `src/data/spatial_omics_datamodule.py`: LightningDataModule wiring preprocessing, splitting, and loaders. +- `src/grass_mil/data/__init__.py`: Data package marker. +- `src/grass_mil/data/spatial_omics_datamodule.py`: LightningDataModule wiring preprocessing, splitting, and loaders. Data components: -- `src/data/components/__init__.py`: Data components package marker. -- `src/data/components/spatial_types.py`: Canonical in-memory spatial table dataclass. -- `src/data/components/loaders.py`: CSV/H5AD/SCE loaders and polygon parsing utilities. -- `src/data/components/patching.py`: Polygon/tile patch extraction helpers. -- `src/data/components/graph_builders.py`: Graph builder interfaces and Delaunay implementation. -- `src/data/components/feature_reducers.py`: Feature reducer interfaces and implementations (`identity`, `pca`). -- `src/data/components/precompute.py`: Precompute orchestration into persisted PyG graph artifacts/indexes. -- `src/data/components/datasets.py`: Dataset wrappers for persisted graph indices and transform application. -- `src/data/components/samplers.py`: Batch/sampler strategies (identity, shadow_native, shadow_custom). -- `src/data/components/transforms.py`: Transform instantiation and composition-vector transform. +- `src/grass_mil/data/components/__init__.py`: Data components package marker. +- `src/grass_mil/data/components/spatial_types.py`: Canonical in-memory spatial table dataclass. +- `src/grass_mil/data/components/loaders.py`: CSV/H5AD/SCE loaders and polygon parsing utilities. +- `src/grass_mil/data/components/patching.py`: Polygon/tile patch extraction helpers. +- `src/grass_mil/data/components/graph_builders.py`: Graph builder interfaces and Delaunay implementation. +- `src/grass_mil/data/components/feature_reducers.py`: Feature reducer interfaces and implementations (`identity`, `pca`). +- `src/grass_mil/data/components/precompute.py`: Precompute orchestration into persisted PyG graph artifacts/indexes. +- `src/grass_mil/data/components/datasets.py`: Dataset wrappers for persisted graph indices and transform application. +- `src/grass_mil/data/components/samplers.py`: Batch/sampler strategies (identity, shadow_native, shadow_custom). +- `src/grass_mil/data/components/transforms.py`: Transform instantiation and composition-vector transform. -### Model Runtime (`src/models/`) +### Model Runtime (`src/grass_mil/models/`) -- `src/models/__init__.py`: Model package marker. -- `src/models/supervised_module.py`: Unified supervised training module (mean + MIL behavior). -- `src/models/bgrl_module.py`: BGRL self-supervised training module. +- `src/grass_mil/models/__init__.py`: Model package marker. +- `src/grass_mil/models/supervised_module.py`: Unified supervised training module (mean + MIL behavior). +- `src/grass_mil/models/bgrl_module.py`: BGRL self-supervised training module. Model components: -- `src/models/components/__init__.py`: Component package exports. -- `src/models/components/backbones.py`: GNN encoder definitions and configuration contract. -- `src/models/components/pooling.py`: Graph pooling wrapper. -- `src/models/components/attention.py`: MIL attention blocks. -- `src/models/components/heads.py`: Graph/node MLP heads. -- `src/models/components/losses.py`: Weighted loss implementations (classification/regression/survival). -- `src/models/components/ssl.py`: BGRL and predictor modules. -- `src/models/components/factory.py`: Factory builders for encoder/attention/loss/head/SSL modules. +- `src/grass_mil/models/components/__init__.py`: Component package exports. +- `src/grass_mil/models/components/backbones.py`: GNN encoder definitions and configuration contract. +- `src/grass_mil/models/components/pooling.py`: Graph pooling wrapper. +- `src/grass_mil/models/components/attention.py`: MIL attention blocks. +- `src/grass_mil/models/components/heads.py`: Graph/node MLP heads. +- `src/grass_mil/models/components/losses.py`: Weighted loss implementations (classification/regression/survival). +- `src/grass_mil/models/components/ssl.py`: BGRL and predictor modules. +- `src/grass_mil/models/components/factory.py`: Factory builders for encoder/attention/loss/head/SSL modules. Training helpers: -- `src/models/training/__init__.py`: Training-helper exports. -- `src/models/training/builders.py`: Component build/inference helpers and task validation. -- `src/models/training/bagging.py`: Bag ID resolution, instance grouping/sampling, bag aggregation logic. -- `src/models/training/loss_utils.py`: Loss routing and MIL auxiliary/regularization utilities. -- `src/models/training/optimization.py`: Optimizer/scheduler instantiation and warmup wrapper. -- `src/models/training/checkpoint_init.py`: Checkpoint loading and encoder-key remapping. -- `src/models/training/ssl_runtime.py`: SSL augmentations and momentum schedule helpers. - -### Utilities (`src/utils/`) - -- `src/utils/__init__.py`: Utility export surface. -- `src/utils/pylogger.py`: Rank-aware logger adapter. -- `src/utils/instantiators.py`: Hydra callback/logger instantiation helpers. -- `src/utils/logging_utils.py`: Hyperparameter logging packaging for logger backends. -- `src/utils/rich_utils.py`: Rich config-tree and tag prompting utilities. -- `src/utils/utils.py`: Misc execution helpers (`extras`, `task_wrapper`, metric extraction). +- `src/grass_mil/models/training/__init__.py`: Training-helper exports. +- `src/grass_mil/models/training/builders.py`: Component build/inference helpers and task validation. +- `src/grass_mil/models/training/bagging.py`: Bag ID resolution, instance grouping/sampling, bag aggregation logic. +- `src/grass_mil/models/training/loss_utils.py`: Loss routing and MIL auxiliary/regularization utilities. +- `src/grass_mil/models/training/optimization.py`: Optimizer/scheduler instantiation and warmup wrapper. +- `src/grass_mil/models/training/checkpoint_init.py`: Checkpoint loading and encoder-key remapping. +- `src/grass_mil/models/training/ssl_runtime.py`: SSL augmentations and momentum schedule helpers. + +### Utilities (`src/grass_mil/utils/`) + +- `src/grass_mil/utils/__init__.py`: Utility export surface. +- `src/grass_mil/utils/pylogger.py`: Rank-aware logger adapter. +- `src/grass_mil/utils/instantiators.py`: Hydra callback/logger instantiation helpers. +- `src/grass_mil/utils/logging_utils.py`: Hyperparameter logging packaging for logger backends. +- `src/grass_mil/utils/rich_utils.py`: Rich config-tree and tag prompting utilities. +- `src/grass_mil/utils/utils.py`: Misc execution helpers (`extras`, `task_wrapper`, metric extraction). ## Tests (`tests/`) diff --git a/docs/workflows_train_infer.md b/docs/workflows_train_infer.md index 8c8acb1..d2744a0 100644 --- a/docs/workflows_train_infer.md +++ b/docs/workflows_train_infer.md @@ -14,7 +14,7 @@ From `/Users/lovrorabuzin/Projects/grass-mil_unification/grass-mil`: ```bash python -m venv .venv source .venv/bin/activate -pip install -r requirements.txt +pip install -e ".[interpretability,report,spatial-stats,io]" export PROJECT_ROOT=/Users/lovrorabuzin/Projects/grass-mil_unification/grass-mil ``` @@ -82,7 +82,7 @@ Key options (`data.csv.*`): Example override: ```bash -python src/train.py \ +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]" @@ -168,7 +168,7 @@ Important split persistence behavior: Typical fresh-data command: ```bash -python src/train.py \ +grass-mil-train \ task=finetune_mean \ model=supervised_module \ data=spatial_omics \ @@ -179,13 +179,13 @@ python src/train.py \ ### Step 4b: Run LOOCV (Single Fold Or Full Loop) -`src/loocv.py` reuses the train stack and runs fold-specific jobs by overriding +`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 -python src/loocv.py \ +grass-mil-loocv \ loocv.mode=single \ data.split.loocv.enabled=true \ data.split.loocv.fold_unit=region \ @@ -197,7 +197,7 @@ python src/loocv.py \ Single-fold by fold index (useful for array jobs): ```bash -python src/loocv.py \ +grass-mil-loocv \ loocv.mode=single \ loocv.fold_index=0 \ data.split.loocv.enabled=true \ @@ -210,7 +210,7 @@ mode, `loocv.fold_index` takes priority and a warning is emitted. Full loop across all discovered folds: ```bash -python src/loocv.py \ +grass-mil-loocv \ loocv.mode=all \ data.split.loocv.enabled=true \ data.split.loocv.fold_unit=region \ @@ -233,13 +233,13 @@ Safety note: Default training (mean aggregation): ```bash -python src/train.py task=finetune_mean model=supervised_module data=spatial_omics +grass-mil-train task=finetune_mean model=supervised_module data=spatial_omics ``` Evaluation of a selected checkpoint: ```bash -python src/eval.py \ +grass-mil-eval \ ckpt_path=/absolute/path/to/checkpoint.ckpt \ task=finetune_mean \ model=supervised_module \ @@ -265,7 +265,7 @@ Training artifacts (`paths.output_dir`): ### Step 1: Run SSL Pretraining (BGRL) ```bash -python src/train.py task=pretrain_bgrl model=bgrl_module data=spatial_omics +grass-mil-train task=pretrain_bgrl model=bgrl_module data=spatial_omics ``` Default pretraining task includes: @@ -286,7 +286,7 @@ If your callback tracks best metric, use best checkpoint path from training logs ### Step 3: Fine-Tune Using The Pretrained Encoder ```bash -python src/train.py \ +grass-mil-train \ task=finetune_mil \ model=supervised_module \ data=spatial_omics \ @@ -298,7 +298,7 @@ python src/train.py \ Optional encoder freezing: ```bash -python src/train.py \ +grass-mil-train \ task=finetune_mil \ model=supervised_module \ model.init_from_ckpt=/absolute/path/to/pretrain.ckpt \ @@ -326,13 +326,13 @@ Common failures: Mean aggregation: ```bash -python src/train.py task=finetune_mean model=supervised_module +grass-mil-train task=finetune_mean model=supervised_module ``` MIL attention aggregation: ```bash -python src/train.py task=finetune_mil model=supervised_module +grass-mil-train task=finetune_mil model=supervised_module ``` Notes: @@ -345,11 +345,11 @@ Notes: Switch backbone type: ```bash -python src/train.py model.encoder.conv_type=gin -python src/train.py model.encoder.conv_type=gcn -python src/train.py model.encoder.conv_type=gat -python src/train.py model.encoder.conv_type=graphsage -python src/train.py model.encoder.conv_type=gine +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: @@ -364,61 +364,42 @@ Important constraints: Common overrides: ```bash -python src/train.py \ +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=set2set \ - model.encoder.set2set_steps=4 + model.encoder.pooling=max ``` ### D) Attention Variants ```bash -python src/train.py task=finetune_mil \ +grass-mil-train task=finetune_mil \ model.attention.attention_type=gated -python src/train.py task=finetune_mil \ +grass-mil-train task=finetune_mil \ model.attention.attention_type=gated_projected \ model.attention.projection_dim=128 ``` -### E) MIL-Only Optional Terms +### E) MIL Region Accumulation -Region accumulation (manual optimizer stepping over region hyperbatches): +Manual optimizer stepping over region hyperbatches. This is a throughput +mechanism, not an objective: region cross-entropy remains the sole loss. ```bash -python src/train.py task=finetune_mil \ +grass-mil-train task=finetune_mil \ task.region_accumulation.enabled=true \ task.region_accumulation.hyperbatch_size=8 ``` -Node auxiliary loss: - -```bash -python src/train.py task=finetune_mil \ - task.node_aux.enabled=true \ - task.node_aux.target_mode=attention_shaped_ti \ - task.node_aux.loss_mode=weighted_bce \ - task.node_aux.weight=0.2 -``` - -Entropy regularization: - -```bash -python src/train.py task=finetune_mil \ - task.entropy_reg.enabled=true \ - task.entropy_reg.mode=attention \ - task.entropy_reg.weight=0.01 -``` - Dual LR groups for MIL backbone/attention: ```bash -python src/train.py task=finetune_mil \ +grass-mil-train task=finetune_mil \ task.optimization.backbone_lr=1e-3 \ task.optimization.attention_lr=4e-3 ``` @@ -428,26 +409,26 @@ python src/train.py task=finetune_mil \ Binary classification (default): ```bash -python src/train.py task.target_type=binary task.loss=categorical_bce +grass-mil-train task.target_type=binary task.loss=categorical_bce ``` Regression: ```bash -python src/train.py task.target_type=regression task.loss=regression_mse -python src/train.py task.target_type=regression task.loss=regression_huber +grass-mil-train task.target_type=regression task.loss=regression_mse +grass-mil-train task.target_type=regression task.loss=regression_huber ``` Survival: ```bash -python src/train.py task.target_type=survival task.loss=survival_coxsgd +grass-mil-train task.target_type=survival task.loss=survival_coxsgd ``` Target column selection: ```bash -python src/train.py task.target_columns="[outcome]" +grass-mil-train task.target_columns="[outcome]" ``` Target format requirements: @@ -459,10 +440,10 @@ Target format requirements: ### A) Supported CLI Inference Path (Checkpoint Evaluation) -`src/eval.py` is the official CLI inference/evaluation entrypoint. +`src/grass_mil/eval.py` is the official CLI inference/evaluation entrypoint. ```bash -python src/eval.py \ +grass-mil-eval \ ckpt_path=/absolute/path/to/checkpoint.ckpt \ task=finetune_mil \ model=supervised_module \ @@ -473,10 +454,10 @@ This runs `trainer.test(...)` and logs test metrics. ### B) Inference Utility Suite (Prediction + Metrics + Embeddings) -Use `src/inference/predict.py` for artifact-centric inference runs: +Use `src/grass_mil/inference/predict.py` for artifact-centric inference runs: ```bash -python src/inference/predict.py \ +grass-mil-predict \ ckpt_path=/absolute/path/to/checkpoint.ckpt \ task=finetune_mil \ model=supervised_module \ @@ -578,7 +559,7 @@ sampling but validation should approximate tissue-level uniform sampling. Example: ```bash -python src/train.py \ +grass-mil-train \ data.sampler.runtime.proportional_root_sampling=true \ data.val_sampler.name=shadow_custom \ data.val_sampler.kwargs="{}" \ @@ -627,7 +608,7 @@ Example: keep training weighted root sampling, but validate with deterministic, non-proportional sampling. ```bash -python src/train.py \ +grass-mil-train \ data.sampler.runtime.proportional_root_sampling=true \ data.val_sampler.name=shadow_custom \ data.val_sampler.kwargs="{}" \ diff --git a/environment.ci-spatial.yaml b/environment.ci-spatial.yaml index 457f164..2373596 100644 --- a/environment.ci-spatial.yaml +++ b/environment.ci-spatial.yaml @@ -19,6 +19,12 @@ dependencies: - pytorch=2.4.1 - pyg=2.6.1 - pytorch-scatter=2.1.2 + # Prebuilt from the pyg channel. Do NOT move this to the pip section: + # pip has no macOS-arm64 wheel, so it builds from source, and PyTorch's + # own c10/util/strong_type.h specializes std::is_arithmetic, which newer + # libc++ (Xcode 26) rejects outright. Any pip-built torch C++ extension + # fails on the macOS runner for that reason. + - pytorch-sparse=0.6.18 - pip: - lightning==2.5.6 - pytorch-lightning==2.5.6 @@ -29,5 +35,4 @@ dependencies: - plotly==5.24.1 - rootutils==1.0.7 - shapely==2.0.6 - - torch-sparse==0.6.18 - sh==2.2.2 diff --git a/pyproject.toml b/pyproject.toml index 36a180f..95b20ff 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,4 +1,48 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "grass-mil" +version = "0.1.0" +description = "Geometric deep learning framework for spatial omics data" +readme = "README.md" +requires-python = ">=3.11" +license = { file = "LICENSE" } +dependencies = [ + "torch>=2.0", + "lightning>=2.0", + "torchmetrics>=0.11.4", + "torch-geometric>=2.5", + "hydra-core==1.3.2", + "hydra-colorlog==1.2.0", + "omegaconf", + "numpy", + "pandas", + "scipy>=1.11", + "scikit-learn", + "rich", + "rootutils", +] + +[project.optional-dependencies] +interpretability = ["plotly", "umap-learn", "jinja2"] +report = ["kaleido", "playwright"] +spatial-stats = ["squidpy", "scanpy", "anndata"] +io = ["anndata", "shapely"] + +[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-predict = "grass_mil.inference.predict:main" +grass-mil-report = "grass_mil.interpretability.report_cli:main" + +[tool.hatch.build.targets.wheel] +packages = ["src/grass_mil"] + [tool.pytest.ini_options] +pythonpath = ["src", "."] addopts = [ "--color=yes", "--durations=0", @@ -10,9 +54,6 @@ filterwarnings = [ "ignore::UserWarning", ] log_cli = "True" -markers = [ - "slow: slow tests", -] minversion = "6.0" testpaths = "tests/" diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 4aad74a..0000000 --- a/requirements.txt +++ /dev/null @@ -1,41 +0,0 @@ -# --------- pytorch --------- # -torch>=2.0.0 -torchvision>=0.15.0 -lightning>=2.0.0 -torchmetrics>=0.11.4 - -# --------- hydra --------- # -hydra-core==1.3.2 -hydra-colorlog==1.2.0 -hydra-optuna-sweeper==1.2.0 - -# --------- 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 - -# --------- others --------- # -rootutils # standardizing the project root setup -pre-commit # hooks for applying linters on commit -rich # beautiful text formatting in terminal -pytest # tests - -# --------- data --------- # -numpy -pandas -scipy -scikit-learn -shapely -anndata -plotly -kaleido -playwright -jinja2 -scanpy -squidpy -torch-geometric -# optional: required only for SingleCellExperiment/rds loading -# rpy2 -# sh # for running bash commands in some tests (linux/macos only) diff --git a/setup.py b/setup.py deleted file mode 100644 index af121a8..0000000 --- a/setup.py +++ /dev/null @@ -1,19 +0,0 @@ -#!/usr/bin/env python - -from setuptools import find_packages, setup - -setup( - name="grass-mil", - version="0.0.1", - description="Geometric deep learning framework for spatial omics data", - install_requires=["lightning", "hydra-core"], - packages=find_packages(), - entry_points={ - "console_scripts": [ - "train_command = src.train:main", - "eval_command = src.eval:main", - "loocv_command = src.loocv:main", - "interpretability_report_command = src.interpretability.report_cli:main", - ] - }, -) diff --git a/src/__init__.py b/src/grass_mil/__init__.py similarity index 100% rename from src/__init__.py rename to src/grass_mil/__init__.py diff --git a/src/grass_mil/_paths.py b/src/grass_mil/_paths.py new file mode 100644 index 0000000..51ba02a --- /dev/null +++ b/src/grass_mil/_paths.py @@ -0,0 +1,29 @@ +"""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/configs/__init__.py b/src/grass_mil/configs/__init__.py similarity index 100% rename from configs/__init__.py rename to src/grass_mil/configs/__init__.py diff --git a/src/grass_mil/configs/augmentation/bgrl_paper.yaml b/src/grass_mil/configs/augmentation/bgrl_paper.yaml new file mode 100644 index 0000000..5874df8 --- /dev/null +++ b/src/grass_mil/configs/augmentation/bgrl_paper.yaml @@ -0,0 +1,17 @@ +# 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 new file mode 100644 index 0000000..2645df7 --- /dev/null +++ b/src/grass_mil/configs/augmentation/uniform.yaml @@ -0,0 +1,6 @@ +# 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/configs/callbacks/default.yaml b/src/grass_mil/configs/callbacks/default.yaml similarity index 100% rename from configs/callbacks/default.yaml rename to src/grass_mil/configs/callbacks/default.yaml diff --git a/configs/callbacks/early_stopping.yaml b/src/grass_mil/configs/callbacks/early_stopping.yaml similarity index 100% rename from configs/callbacks/early_stopping.yaml rename to src/grass_mil/configs/callbacks/early_stopping.yaml diff --git a/configs/callbacks/model_checkpoint.yaml b/src/grass_mil/configs/callbacks/model_checkpoint.yaml similarity index 100% rename from configs/callbacks/model_checkpoint.yaml rename to src/grass_mil/configs/callbacks/model_checkpoint.yaml diff --git a/configs/callbacks/model_summary.yaml b/src/grass_mil/configs/callbacks/model_summary.yaml similarity index 100% rename from configs/callbacks/model_summary.yaml rename to src/grass_mil/configs/callbacks/model_summary.yaml diff --git a/configs/callbacks/none.yaml b/src/grass_mil/configs/callbacks/none.yaml similarity index 100% rename from configs/callbacks/none.yaml rename to src/grass_mil/configs/callbacks/none.yaml diff --git a/configs/callbacks/rich_progress_bar.yaml b/src/grass_mil/configs/callbacks/rich_progress_bar.yaml similarity index 100% rename from configs/callbacks/rich_progress_bar.yaml rename to src/grass_mil/configs/callbacks/rich_progress_bar.yaml diff --git a/configs/data/sampler/identity.yaml b/src/grass_mil/configs/data/sampler/identity.yaml similarity index 100% rename from configs/data/sampler/identity.yaml rename to src/grass_mil/configs/data/sampler/identity.yaml diff --git a/src/grass_mil/configs/data/sampler/shadow_custom.yaml b/src/grass_mil/configs/data/sampler/shadow_custom.yaml new file mode 100644 index 0000000..8176031 --- /dev/null +++ b/src/grass_mil/configs/data/sampler/shadow_custom.yaml @@ -0,0 +1,33 @@ +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 new file mode 100644 index 0000000..97971fb --- /dev/null +++ b/src/grass_mil/configs/data/sampler/shadow_native.yaml @@ -0,0 +1,33 @@ +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/configs/data/spatial_omics.yaml b/src/grass_mil/configs/data/spatial_omics.yaml similarity index 84% rename from configs/data/spatial_omics.yaml rename to src/grass_mil/configs/data/spatial_omics.yaml index a7ae84c..1ad5f44 100644 --- a/configs/data/spatial_omics.yaml +++ b/src/grass_mil/configs/data/spatial_omics.yaml @@ -1,4 +1,4 @@ -_target_: src.data.spatial_omics_datamodule.SpatialOmicsDataModule +_target_: grass_mil.data.spatial_omics_datamodule.SpatialOmicsDataModule data_dir: ${paths.data_dir} raw_manifest_path: ${paths.data_dir}/raw/manifest.csv @@ -88,6 +88,18 @@ sampler: property_name: cell_type weight_mode: inverse min_weight: 1.0e-6 + # See configs/data/sampler/shadow_custom.yaml for the full documentation. + 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 # Optional validation-only sampler override. # When null, validation reuses `sampler`. diff --git a/configs/debug/default.yaml b/src/grass_mil/configs/debug/default.yaml similarity index 100% rename from configs/debug/default.yaml rename to src/grass_mil/configs/debug/default.yaml diff --git a/configs/debug/fdr.yaml b/src/grass_mil/configs/debug/fdr.yaml similarity index 100% rename from configs/debug/fdr.yaml rename to src/grass_mil/configs/debug/fdr.yaml diff --git a/configs/debug/limit.yaml b/src/grass_mil/configs/debug/limit.yaml similarity index 100% rename from configs/debug/limit.yaml rename to src/grass_mil/configs/debug/limit.yaml diff --git a/configs/debug/overfit.yaml b/src/grass_mil/configs/debug/overfit.yaml similarity index 100% rename from configs/debug/overfit.yaml rename to src/grass_mil/configs/debug/overfit.yaml diff --git a/configs/debug/profiler.yaml b/src/grass_mil/configs/debug/profiler.yaml similarity index 100% rename from configs/debug/profiler.yaml rename to src/grass_mil/configs/debug/profiler.yaml diff --git a/configs/eval.yaml b/src/grass_mil/configs/eval.yaml similarity index 100% rename from configs/eval.yaml rename to src/grass_mil/configs/eval.yaml diff --git a/configs/experiment/debug.yaml b/src/grass_mil/configs/experiment/debug.yaml similarity index 100% rename from configs/experiment/debug.yaml rename to src/grass_mil/configs/experiment/debug.yaml diff --git a/configs/experiment/notebook_data_shadow_custom.yaml b/src/grass_mil/configs/experiment/notebook_data_shadow_custom.yaml similarity index 100% rename from configs/experiment/notebook_data_shadow_custom.yaml rename to src/grass_mil/configs/experiment/notebook_data_shadow_custom.yaml diff --git a/configs/experiment/notebook_model_components.yaml b/src/grass_mil/configs/experiment/notebook_model_components.yaml similarity index 100% rename from configs/experiment/notebook_model_components.yaml rename to src/grass_mil/configs/experiment/notebook_model_components.yaml diff --git a/configs/experiment/notebook_training_logic.yaml b/src/grass_mil/configs/experiment/notebook_training_logic.yaml similarity index 100% rename from configs/experiment/notebook_training_logic.yaml rename to src/grass_mil/configs/experiment/notebook_training_logic.yaml diff --git a/configs/extras/default.yaml b/src/grass_mil/configs/extras/default.yaml similarity index 100% rename from configs/extras/default.yaml rename to src/grass_mil/configs/extras/default.yaml diff --git a/configs/hydra/default.yaml b/src/grass_mil/configs/hydra/default.yaml similarity index 100% rename from configs/hydra/default.yaml rename to src/grass_mil/configs/hydra/default.yaml diff --git a/configs/inference/aggregation/default.yaml b/src/grass_mil/configs/inference/aggregation/default.yaml similarity index 100% rename from configs/inference/aggregation/default.yaml rename to src/grass_mil/configs/inference/aggregation/default.yaml diff --git a/configs/inference/embeddings/default.yaml b/src/grass_mil/configs/inference/embeddings/default.yaml similarity index 100% rename from configs/inference/embeddings/default.yaml rename to src/grass_mil/configs/inference/embeddings/default.yaml diff --git a/configs/inference/interpretability/default.yaml b/src/grass_mil/configs/inference/interpretability/default.yaml similarity index 69% rename from configs/inference/interpretability/default.yaml rename to src/grass_mil/configs/inference/interpretability/default.yaml index 61c55f8..f34e7a6 100644 --- a/configs/inference/interpretability/default.yaml +++ b/src/grass_mil/configs/inference/interpretability/default.yaml @@ -15,6 +15,11 @@ 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: diff --git a/configs/inference/metrics/default.yaml b/src/grass_mil/configs/inference/metrics/default.yaml similarity index 100% rename from configs/inference/metrics/default.yaml rename to src/grass_mil/configs/inference/metrics/default.yaml diff --git a/configs/inference/predict.yaml b/src/grass_mil/configs/inference/predict.yaml similarity index 100% rename from configs/inference/predict.yaml rename to src/grass_mil/configs/inference/predict.yaml diff --git a/configs/interpretability/clustering/agglomerative.yaml b/src/grass_mil/configs/interpretability/clustering/agglomerative.yaml similarity index 100% rename from configs/interpretability/clustering/agglomerative.yaml rename to src/grass_mil/configs/interpretability/clustering/agglomerative.yaml diff --git a/configs/interpretability/clustering/dbscan.yaml b/src/grass_mil/configs/interpretability/clustering/dbscan.yaml similarity index 100% rename from configs/interpretability/clustering/dbscan.yaml rename to src/grass_mil/configs/interpretability/clustering/dbscan.yaml diff --git a/configs/interpretability/clustering/hdbscan.yaml b/src/grass_mil/configs/interpretability/clustering/hdbscan.yaml similarity index 100% rename from configs/interpretability/clustering/hdbscan.yaml rename to src/grass_mil/configs/interpretability/clustering/hdbscan.yaml diff --git a/configs/interpretability/clustering/kmeans.yaml b/src/grass_mil/configs/interpretability/clustering/kmeans.yaml similarity index 100% rename from configs/interpretability/clustering/kmeans.yaml rename to src/grass_mil/configs/interpretability/clustering/kmeans.yaml diff --git a/configs/interpretability/clustering/optics.yaml b/src/grass_mil/configs/interpretability/clustering/optics.yaml similarity index 100% rename from configs/interpretability/clustering/optics.yaml rename to src/grass_mil/configs/interpretability/clustering/optics.yaml diff --git a/configs/interpretability/clustering/spectral.yaml b/src/grass_mil/configs/interpretability/clustering/spectral.yaml similarity index 100% rename from configs/interpretability/clustering/spectral.yaml rename to src/grass_mil/configs/interpretability/clustering/spectral.yaml diff --git a/src/grass_mil/configs/interpretability/plugins/default.yaml b/src/grass_mil/configs/interpretability/plugins/default.yaml new file mode 100644 index 0000000..979321a --- /dev/null +++ b/src/grass_mil/configs/interpretability/plugins/default.yaml @@ -0,0 +1,89 @@ +enabled: + - niche_profiles +params: + niche_profiles: + variance_estimator: unbiased + neighborhood_enrichment: + label_column: niche_label + # One of: zscore, obs-exp, log2fc. + enrichment_mode: zscore + # n_perms=0 uses analytical expected/std (runtime warning emitted). + # Set n_perms>0 for Squidpy-style permutation expected/std + p-values. + n_perms: 0 + diff_neighborhood_enrichment: + label_column: niche_label + condition_column: condition + permutation_group_column: sample_id + # One of: zscore, obs-exp, log2fc. + enrichment_mode: zscore + # n_perms=0 computes pairwise differentials without permutation p-values. + n_perms: 0 + filtration_curves: + # Distances and thresholds are interpreted in micrometers. + threshold_start: 0.0 + threshold_stop: 55.0 + threshold_count: 500 + per_niche_cell_type_enrichment: + cell_type_column: cell_type + niche_column: niche_label + # Niches with fewer than 2 instances, and Background (-1) when + # skip_background is true, are omitted. + skip_background: true + # One of: zscore, obs-exp, log2fc. + enrichment_mode: obs-exp + # n_perms=0 uses analytical expected/std (runtime warning emitted). + n_perms: 0 + random_state: 42 + undirected: false + 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 + morans_i: + feature_prefix: comp_ + n_perms: 100 + random_state: 0 + min_nodes: 3 + # Contrast every niche against this one; -1 is the Background label. + reference_group: -1 + ripley: + label_column: niche_label + group_column: null + x_column: center_x + y_column: center_y + n_radii: 50 + max_fraction: 0.25 + min_count: 5 + # median | max | first. `first` reproduces the legacy order-dependent grid. + radius_source: median + 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 diff --git a/src/grass_mil/configs/interpretability/plugins/full.yaml b/src/grass_mil/configs/interpretability/plugins/full.yaml new file mode 100644 index 0000000..fd1e66c --- /dev/null +++ b/src/grass_mil/configs/interpretability/plugins/full.yaml @@ -0,0 +1,14 @@ +# Every builtin analysis enabled. Slower than the default preset: +# Moran's I and Ripley scale with niche size and permutation count. +defaults: + - default + - _self_ + +enabled: + - niche_profiles + - margin_attribution + - neighborhood_enrichment + - per_niche_cell_type_enrichment + - filtration_curves + - morans_i + - ripley diff --git a/configs/interpretability/reduction/pca.yaml b/src/grass_mil/configs/interpretability/reduction/pca.yaml similarity index 100% rename from configs/interpretability/reduction/pca.yaml rename to src/grass_mil/configs/interpretability/reduction/pca.yaml diff --git a/configs/interpretability/reduction/tsne.yaml b/src/grass_mil/configs/interpretability/reduction/tsne.yaml similarity index 100% rename from configs/interpretability/reduction/tsne.yaml rename to src/grass_mil/configs/interpretability/reduction/tsne.yaml diff --git a/configs/interpretability/reduction/umap.yaml b/src/grass_mil/configs/interpretability/reduction/umap.yaml similarity index 100% rename from configs/interpretability/reduction/umap.yaml rename to src/grass_mil/configs/interpretability/reduction/umap.yaml diff --git a/configs/interpretability/report.yaml b/src/grass_mil/configs/interpretability/report.yaml similarity index 100% rename from configs/interpretability/report.yaml rename to src/grass_mil/configs/interpretability/report.yaml diff --git a/configs/local/.gitkeep b/src/grass_mil/configs/local/.gitkeep similarity index 100% rename from configs/local/.gitkeep rename to src/grass_mil/configs/local/.gitkeep diff --git a/configs/logger/aim.yaml b/src/grass_mil/configs/logger/aim.yaml similarity index 100% rename from configs/logger/aim.yaml rename to src/grass_mil/configs/logger/aim.yaml diff --git a/configs/logger/comet.yaml b/src/grass_mil/configs/logger/comet.yaml similarity index 100% rename from configs/logger/comet.yaml rename to src/grass_mil/configs/logger/comet.yaml diff --git a/configs/logger/csv.yaml b/src/grass_mil/configs/logger/csv.yaml similarity index 100% rename from configs/logger/csv.yaml rename to src/grass_mil/configs/logger/csv.yaml diff --git a/configs/logger/many_loggers.yaml b/src/grass_mil/configs/logger/many_loggers.yaml similarity index 100% rename from configs/logger/many_loggers.yaml rename to src/grass_mil/configs/logger/many_loggers.yaml diff --git a/configs/logger/mlflow.yaml b/src/grass_mil/configs/logger/mlflow.yaml similarity index 100% rename from configs/logger/mlflow.yaml rename to src/grass_mil/configs/logger/mlflow.yaml diff --git a/configs/logger/neptune.yaml b/src/grass_mil/configs/logger/neptune.yaml similarity index 100% rename from configs/logger/neptune.yaml rename to src/grass_mil/configs/logger/neptune.yaml diff --git a/configs/logger/tensorboard.yaml b/src/grass_mil/configs/logger/tensorboard.yaml similarity index 100% rename from configs/logger/tensorboard.yaml rename to src/grass_mil/configs/logger/tensorboard.yaml diff --git a/configs/logger/wandb.yaml b/src/grass_mil/configs/logger/wandb.yaml similarity index 100% rename from configs/logger/wandb.yaml rename to src/grass_mil/configs/logger/wandb.yaml diff --git a/configs/loocv.yaml b/src/grass_mil/configs/loocv.yaml similarity index 100% rename from configs/loocv.yaml rename to src/grass_mil/configs/loocv.yaml diff --git a/src/grass_mil/configs/model/attention/gated.yaml b/src/grass_mil/configs/model/attention/gated.yaml new file mode 100644 index 0000000..db41112 --- /dev/null +++ b/src/grass_mil/configs/model/attention/gated.yaml @@ -0,0 +1,7 @@ +_target_: grass_mil.models.components.attention.AttnNetGated +input_dim: 128 +hidden_dim: 64 +dropout: false +# Attention is per-class: n_classes must equal the graph head's output_dim. +# The task preset sets both together (finetune_mil uses 2 for binary). +n_classes: 1 diff --git a/src/grass_mil/configs/model/attention/gated_projected.yaml b/src/grass_mil/configs/model/attention/gated_projected.yaml new file mode 100644 index 0000000..85478c6 --- /dev/null +++ b/src/grass_mil/configs/model/attention/gated_projected.yaml @@ -0,0 +1,8 @@ +_target_: grass_mil.models.components.attention.AttnNetGatedProjected +input_dim: 128 +projection_dim: 64 +hidden_dim: 16 +dropout: false +# Attention is per-class: n_classes must equal the graph head's output_dim. +# The task preset sets both together (finetune_mil uses 2 for binary). +n_classes: 1 diff --git a/configs/model/bgrl_module.yaml b/src/grass_mil/configs/model/bgrl_module.yaml similarity index 54% rename from configs/model/bgrl_module.yaml rename to src/grass_mil/configs/model/bgrl_module.yaml index 4c60dbf..1f4c1a8 100644 --- a/configs/model/bgrl_module.yaml +++ b/src/grass_mil/configs/model/bgrl_module.yaml @@ -1,4 +1,4 @@ -_target_: src.models.bgrl_module.BGRLModule +_target_: grass_mil.models.bgrl_module.BGRLModule _recursive_: false encoder: @@ -12,11 +12,16 @@ encoder: jk: last act: relu pooling: mean - set2set_steps: 3 gat_heads: 4 use_edge_attr: false edge_weight_index: 0 edge_attr_dim: null + # Selects a single scalar edge column (the edge length) for GINE message + # passing; null passes the whole edge_attr vector through. + edge_feature_index: null + # Learned embedding of a categorical node label, summed with the projection + # of the continuous features. null disables it (features come from `x` only). + categorical_embedding: null ssl: method: bgrl diff --git a/src/grass_mil/configs/model/components.yaml b/src/grass_mil/configs/model/components.yaml new file mode 100644 index 0000000..2116e22 --- /dev/null +++ b/src/grass_mil/configs/model/components.yaml @@ -0,0 +1,54 @@ +# @package _global_ + +model: + flags: + use_attention: true + use_ssl: false + + encoder: + _target_: grass_mil.models.components.backbones.EncoderConfig + input_dim: 64 + 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 + # Selects a single scalar edge column (the edge length) for GINE message + # passing; null passes the whole edge_attr vector through. + edge_feature_index: null + # Learned embedding of a categorical node label, summed with the projection + # of the continuous features. null disables it (features come from `x` only). + categorical_embedding: null + + attention: + _target_: grass_mil.models.components.attention.AttnNetGatedProjected + input_dim: 128 + projection_dim: 64 + hidden_dim: 16 + dropout: false + # Must equal graph_head.output_dim: attention is per-class. + n_classes: 1 + + ssl: + method: bgrl + predictor: + _target_: grass_mil.models.components.ssl.MLPPredictor + input_size: 128 + output_size: 128 + hidden_size: 512 + + graph_head: + _target_: grass_mil.models.components.heads.GraphPredictionHead + input_dim: 128 + output_dim: 1 + hidden_dim: 128 + num_layers: 2 + dropout: 0.0 diff --git a/configs/model/encoder/gat.yaml b/src/grass_mil/configs/model/encoder/gat.yaml similarity index 74% rename from configs/model/encoder/gat.yaml rename to src/grass_mil/configs/model/encoder/gat.yaml index 5c99954..dd080d4 100644 --- a/configs/model/encoder/gat.yaml +++ b/src/grass_mil/configs/model/encoder/gat.yaml @@ -1,4 +1,4 @@ -_target_: src.models.components.backbones.EncoderConfig +_target_: grass_mil.models.components.backbones.EncoderConfig input_dim: 64 hidden_dim: 128 out_dim: 128 @@ -9,7 +9,6 @@ norm: batchnorm jk: last act: relu pooling: mean -set2set_steps: 3 gat_heads: 4 use_edge_attr: false edge_weight_index: 0 diff --git a/configs/model/encoder/gcn.yaml b/src/grass_mil/configs/model/encoder/gcn.yaml similarity index 74% rename from configs/model/encoder/gcn.yaml rename to src/grass_mil/configs/model/encoder/gcn.yaml index 9cf83e8..75d9b27 100644 --- a/configs/model/encoder/gcn.yaml +++ b/src/grass_mil/configs/model/encoder/gcn.yaml @@ -1,4 +1,4 @@ -_target_: src.models.components.backbones.EncoderConfig +_target_: grass_mil.models.components.backbones.EncoderConfig input_dim: 64 hidden_dim: 128 out_dim: 128 @@ -9,7 +9,6 @@ norm: batchnorm jk: last act: relu pooling: mean -set2set_steps: 3 gat_heads: 4 use_edge_attr: true edge_weight_index: 0 diff --git a/configs/model/encoder/gin.yaml b/src/grass_mil/configs/model/encoder/gin.yaml similarity index 74% rename from configs/model/encoder/gin.yaml rename to src/grass_mil/configs/model/encoder/gin.yaml index 4b2dab7..ec63d74 100644 --- a/configs/model/encoder/gin.yaml +++ b/src/grass_mil/configs/model/encoder/gin.yaml @@ -1,4 +1,4 @@ -_target_: src.models.components.backbones.EncoderConfig +_target_: grass_mil.models.components.backbones.EncoderConfig input_dim: 64 hidden_dim: 128 out_dim: 128 @@ -9,7 +9,6 @@ norm: batchnorm jk: last act: relu pooling: mean -set2set_steps: 3 gat_heads: 4 use_edge_attr: false edge_weight_index: 0 diff --git a/configs/model/encoder/gine.yaml b/src/grass_mil/configs/model/encoder/gine.yaml similarity index 73% rename from configs/model/encoder/gine.yaml rename to src/grass_mil/configs/model/encoder/gine.yaml index 932c5c3..6d5546e 100644 --- a/configs/model/encoder/gine.yaml +++ b/src/grass_mil/configs/model/encoder/gine.yaml @@ -1,4 +1,4 @@ -_target_: src.models.components.backbones.EncoderConfig +_target_: grass_mil.models.components.backbones.EncoderConfig input_dim: 64 hidden_dim: 128 out_dim: 128 @@ -9,7 +9,6 @@ norm: batchnorm jk: last act: relu pooling: mean -set2set_steps: 3 gat_heads: 4 use_edge_attr: true edge_weight_index: 0 diff --git a/src/grass_mil/configs/model/encoder/gine_length.yaml b/src/grass_mil/configs/model/encoder/gine_length.yaml new file mode 100644 index 0000000..02ec00b --- /dev/null +++ b/src/grass_mil/configs/model/encoder/gine_length.yaml @@ -0,0 +1,36 @@ +# Production GRASS-MIL encoder: GINE conditioned on the scalar edge length, +# with cell identity supplied through a learned embedding rather than as a +# one-hot column of `x`. +# +# `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. +_target_: grass_mil.models.components.backbones.EncoderConfig +input_dim: 0 +hidden_dim: 256 +out_dim: 256 +num_layers: 4 +dropout: 0.2 +conv_type: gine +norm: layernorm +jk: last +act: relu +pooling: max +gat_heads: 4 + +# Message passing reads exactly one edge column: the edge length. Index 0 of +# `edge_attr` is `distance` for the delaunay graph builder's default +# `edge_features: ["distance", "neighbor"]`. +use_edge_attr: true +edge_weight_index: 0 +edge_attr_dim: 1 +edge_feature_index: 0 + +categorical_embedding: + label: cell_type + num_embeddings: null + column_index: null + # Reserve a trailing row so BGRL augmentations can mark a dropped node + # "unassigned" without resizing the table. + reserve_unassigned: true diff --git a/configs/model/encoder/graphsage.yaml b/src/grass_mil/configs/model/encoder/graphsage.yaml similarity index 74% rename from configs/model/encoder/graphsage.yaml rename to src/grass_mil/configs/model/encoder/graphsage.yaml index d5c3110..88df378 100644 --- a/configs/model/encoder/graphsage.yaml +++ b/src/grass_mil/configs/model/encoder/graphsage.yaml @@ -1,4 +1,4 @@ -_target_: src.models.components.backbones.EncoderConfig +_target_: grass_mil.models.components.backbones.EncoderConfig input_dim: 64 hidden_dim: 128 out_dim: 128 @@ -9,7 +9,6 @@ norm: batchnorm jk: last act: relu pooling: mean -set2set_steps: 3 gat_heads: 4 use_edge_attr: false edge_weight_index: 0 diff --git a/configs/model/heads/graph.yaml b/src/grass_mil/configs/model/heads/graph.yaml similarity index 52% rename from configs/model/heads/graph.yaml rename to src/grass_mil/configs/model/heads/graph.yaml index b3ecdb9..26752d9 100644 --- a/configs/model/heads/graph.yaml +++ b/src/grass_mil/configs/model/heads/graph.yaml @@ -1,4 +1,4 @@ -_target_: src.models.components.heads.GraphPredictionHead +_target_: grass_mil.models.components.heads.GraphPredictionHead input_dim: 128 output_dim: 1 hidden_dim: 128 diff --git a/src/grass_mil/configs/model/loss/categorical_bce.yaml b/src/grass_mil/configs/model/loss/categorical_bce.yaml new file mode 100644 index 0000000..3ebaef1 --- /dev/null +++ b/src/grass_mil/configs/model/loss/categorical_bce.yaml @@ -0,0 +1,2 @@ +_target_: grass_mil.models.components.losses.WeightedBCEWithLogitsLoss +pos_weight: null diff --git a/src/grass_mil/configs/model/loss/categorical_ce.yaml b/src/grass_mil/configs/model/loss/categorical_ce.yaml new file mode 100644 index 0000000..28218f2 --- /dev/null +++ b/src/grass_mil/configs/model/loss/categorical_ce.yaml @@ -0,0 +1,2 @@ +_target_: grass_mil.models.components.losses.WeightedCrossEntropyLoss +class_weight: null diff --git a/src/grass_mil/configs/model/loss/regression_huber.yaml b/src/grass_mil/configs/model/loss/regression_huber.yaml new file mode 100644 index 0000000..da58ab7 --- /dev/null +++ b/src/grass_mil/configs/model/loss/regression_huber.yaml @@ -0,0 +1,2 @@ +_target_: grass_mil.models.components.losses.WeightedHuberLoss +delta: 1.0 diff --git a/src/grass_mil/configs/model/loss/regression_mse.yaml b/src/grass_mil/configs/model/loss/regression_mse.yaml new file mode 100644 index 0000000..c459c7e --- /dev/null +++ b/src/grass_mil/configs/model/loss/regression_mse.yaml @@ -0,0 +1 @@ +_target_: grass_mil.models.components.losses.WeightedMSELoss diff --git a/src/grass_mil/configs/model/loss/survival_coxsgd.yaml b/src/grass_mil/configs/model/loss/survival_coxsgd.yaml new file mode 100644 index 0000000..0ac72a6 --- /dev/null +++ b/src/grass_mil/configs/model/loss/survival_coxsgd.yaml @@ -0,0 +1,3 @@ +_target_: grass_mil.models.components.losses.CoxSGDLoss +top_n: 5 +regularizer_weight: 0.05 diff --git a/configs/model/ssl/bgrl.yaml b/src/grass_mil/configs/model/ssl/bgrl.yaml similarity index 58% rename from configs/model/ssl/bgrl.yaml rename to src/grass_mil/configs/model/ssl/bgrl.yaml index 8671d4f..fc77634 100644 --- a/configs/model/ssl/bgrl.yaml +++ b/src/grass_mil/configs/model/ssl/bgrl.yaml @@ -1,6 +1,6 @@ method: bgrl predictor: - _target_: src.models.components.ssl.MLPPredictor + _target_: grass_mil.models.components.ssl.MLPPredictor input_size: 128 output_size: 128 hidden_size: 512 diff --git a/configs/model/supervised_module.yaml b/src/grass_mil/configs/model/supervised_module.yaml similarity index 64% rename from configs/model/supervised_module.yaml rename to src/grass_mil/configs/model/supervised_module.yaml index 10aa191..e74132f 100644 --- a/configs/model/supervised_module.yaml +++ b/src/grass_mil/configs/model/supervised_module.yaml @@ -1,4 +1,4 @@ -_target_: src.models.supervised_module.SupervisedModule +_target_: grass_mil.models.supervised_module.SupervisedModule _recursive_: false flags: @@ -15,11 +15,16 @@ encoder: jk: last act: relu pooling: mean - set2set_steps: 3 gat_heads: 4 use_edge_attr: false edge_weight_index: 0 edge_attr_dim: null + # Selects a single scalar edge column (the edge length) for GINE message + # passing; null passes the whole edge_attr vector through. + edge_feature_index: null + # Learned embedding of a categorical node label, summed with the projection + # of the continuous features. null disables it (features come from `x` only). + categorical_embedding: null # Attention block -- only used when task.aggregation == mil_attention. attention: @@ -28,6 +33,7 @@ attention: projection_dim: 64 hidden_dim: 16 dropout: false + # Must equal graph_head.output_dim: attention is per-class. n_classes: 1 graph_head: diff --git a/configs/optim/adamw.yaml b/src/grass_mil/configs/optim/adamw.yaml similarity index 100% rename from configs/optim/adamw.yaml rename to src/grass_mil/configs/optim/adamw.yaml diff --git a/configs/paths/default.yaml b/src/grass_mil/configs/paths/default.yaml similarity index 65% rename from configs/paths/default.yaml rename to src/grass_mil/configs/paths/default.yaml index ec81db2..45ddc6d 100644 --- a/configs/paths/default.yaml +++ b/src/grass_mil/configs/paths/default.yaml @@ -1,7 +1,9 @@ # path to root directory -# this requires PROJECT_ROOT environment variable to exist +# 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} +root_dir: ${oc.env:PROJECT_ROOT,${hydra:runtime.cwd}} # path to data directory data_dir: ${paths.root_dir}/data/ diff --git a/configs/scheduler/cosine_epoch.yaml b/src/grass_mil/configs/scheduler/cosine_epoch.yaml similarity index 100% rename from configs/scheduler/cosine_epoch.yaml rename to src/grass_mil/configs/scheduler/cosine_epoch.yaml diff --git a/configs/scheduler/cosine_step.yaml b/src/grass_mil/configs/scheduler/cosine_step.yaml similarity index 100% rename from configs/scheduler/cosine_step.yaml rename to src/grass_mil/configs/scheduler/cosine_step.yaml diff --git a/configs/task/finetune_mean.yaml b/src/grass_mil/configs/task/finetune_mean.yaml similarity index 100% rename from configs/task/finetune_mean.yaml rename to src/grass_mil/configs/task/finetune_mean.yaml diff --git a/configs/task/finetune_mil.yaml b/src/grass_mil/configs/task/finetune_mil.yaml similarity index 51% rename from configs/task/finetune_mil.yaml rename to src/grass_mil/configs/task/finetune_mil.yaml index 63c8857..e4001c1 100644 --- a/configs/task/finetune_mil.yaml +++ b/src/grass_mil/configs/task/finetune_mil.yaml @@ -7,9 +7,13 @@ task: bag_fallback_key: sample_id max_instances_per_bag: 0 instance_sampling: all - target_type: binary + # The published 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_bce + loss: categorical_ce optimization: backbone_lr: null attention_lr: null @@ -24,19 +28,11 @@ task: mode: manual_region_buffer hyperbatch_size: 8 flush_on_epoch_end: true - node_aux: - enabled: false - # graph_label | attention_shaped_ti - target_mode: attention_shaped_ti - # bce | weighted_bce - loss_mode: bce - weight: 1.0 - entropy_reg: - enabled: false - # attention | attention_shaped_target - mode: attention_shaped_target - weight: 0.01 - loss_weights: - region: 1.0 - node_aux: 1.0 - entropy: 0.0 + +# 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 new file mode 100644 index 0000000..cc9f46b --- /dev/null +++ b/src/grass_mil/configs/task/finetune_survival.yaml @@ -0,0 +1,43 @@ +# @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/configs/task/pretrain_bgrl.yaml b/src/grass_mil/configs/task/pretrain_bgrl.yaml similarity index 80% rename from configs/task/pretrain_bgrl.yaml rename to src/grass_mil/configs/task/pretrain_bgrl.yaml index e199f0d..fbbc2c3 100644 --- a/configs/task/pretrain_bgrl.yaml +++ b/src/grass_mil/configs/task/pretrain_bgrl.yaml @@ -21,6 +21,9 @@ task: 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 diff --git a/configs/train.yaml b/src/grass_mil/configs/train.yaml similarity index 100% rename from configs/train.yaml rename to src/grass_mil/configs/train.yaml diff --git a/configs/trainer/cpu.yaml b/src/grass_mil/configs/trainer/cpu.yaml similarity index 100% rename from configs/trainer/cpu.yaml rename to src/grass_mil/configs/trainer/cpu.yaml diff --git a/configs/trainer/ddp.yaml b/src/grass_mil/configs/trainer/ddp.yaml similarity index 100% rename from configs/trainer/ddp.yaml rename to src/grass_mil/configs/trainer/ddp.yaml diff --git a/configs/trainer/ddp_sim.yaml b/src/grass_mil/configs/trainer/ddp_sim.yaml similarity index 100% rename from configs/trainer/ddp_sim.yaml rename to src/grass_mil/configs/trainer/ddp_sim.yaml diff --git a/configs/trainer/default.yaml b/src/grass_mil/configs/trainer/default.yaml similarity index 100% rename from configs/trainer/default.yaml rename to src/grass_mil/configs/trainer/default.yaml diff --git a/configs/trainer/gpu.yaml b/src/grass_mil/configs/trainer/gpu.yaml similarity index 100% rename from configs/trainer/gpu.yaml rename to src/grass_mil/configs/trainer/gpu.yaml diff --git a/configs/trainer/mps.yaml b/src/grass_mil/configs/trainer/mps.yaml similarity index 100% rename from configs/trainer/mps.yaml rename to src/grass_mil/configs/trainer/mps.yaml diff --git a/src/interpretability/contracts.py b/src/grass_mil/contracts.py similarity index 62% rename from src/interpretability/contracts.py rename to src/grass_mil/contracts.py index e133083..f4b4063 100644 --- a/src/interpretability/contracts.py +++ b/src/grass_mil/contracts.py @@ -1,3 +1,12 @@ +"""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 @@ -7,6 +16,26 @@ 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: @@ -43,15 +72,11 @@ class ClusteringResult: @dataclass(frozen=True) -class ClusterSummary: - cluster_labels: np.ndarray +class NicheSummary: + niche_labels: np.ndarray composition: pd.DataFrame enrichment: pd.DataFrame - cluster_counts: pd.Series - weighted_scores: Optional[pd.Series] = None - mean_scores: Optional[pd.Series] = None - attention_present: Optional[pd.Series] = None - attention_lift_present: Optional[pd.Series] = None + niche_counts: pd.Series @dataclass(frozen=True) @@ -85,7 +110,7 @@ class PluginResult: @dataclass(frozen=True) -class ClusterTransferBundle: +class NicheTransferBundle: embedding_columns: List[str] id_column: str cluster_on_pca: bool @@ -99,9 +124,9 @@ class ClusterTransferBundle: class ReportBundle: dataset: InterpretabilityDataset reduction: Optional[ReductionResult] - cluster_feature_reduction: Optional[ReductionResult] + niche_feature_reduction: Optional[ReductionResult] clustering: Optional[ClusteringResult] - cluster_summary: Optional[ClusterSummary] + niche_summary: Optional[NicheSummary] plugin_results: Dict[str, PluginResult] artifacts_dir: Path metadata: Dict[str, Any] = field(default_factory=dict) diff --git a/src/data/__init__.py b/src/grass_mil/data/__init__.py similarity index 100% rename from src/data/__init__.py rename to src/grass_mil/data/__init__.py diff --git a/src/data/components/__init__.py b/src/grass_mil/data/components/__init__.py similarity index 100% rename from src/data/components/__init__.py rename to src/grass_mil/data/components/__init__.py diff --git a/src/data/components/datasets.py b/src/grass_mil/data/components/datasets.py similarity index 70% rename from src/data/components/datasets.py rename to src/grass_mil/data/components/datasets.py index b5827c6..ced8651 100644 --- a/src/data/components/datasets.py +++ b/src/grass_mil/data/components/datasets.py @@ -18,6 +18,23 @@ class ProcessedIndexEntry: split: Optional[str] = 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) @@ -44,7 +61,7 @@ def __len__(self) -> int: def __getitem__(self, idx: int) -> torch.nn.Module: entry = self.entries[idx] data = torch.load(entry.path) - return data + return _migrate_legacy_attributes(data) class TransformDataset(Dataset): diff --git a/src/grass_mil/data/components/ego_radius.py b/src/grass_mil/data/components/ego_radius.py new file mode 100644 index 0000000..6cffb75 --- /dev/null +++ b/src/grass_mil/data/components/ego_radius.py @@ -0,0 +1,180 @@ +"""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/data/components/feature_reducers.py b/src/grass_mil/data/components/feature_reducers.py similarity index 100% rename from src/data/components/feature_reducers.py rename to src/grass_mil/data/components/feature_reducers.py diff --git a/src/data/components/graph_builders.py b/src/grass_mil/data/components/graph_builders.py similarity index 100% rename from src/data/components/graph_builders.py rename to src/grass_mil/data/components/graph_builders.py diff --git a/src/data/components/loaders.py b/src/grass_mil/data/components/loaders.py similarity index 100% rename from src/data/components/loaders.py rename to src/grass_mil/data/components/loaders.py diff --git a/src/data/components/patching.py b/src/grass_mil/data/components/patching.py similarity index 100% rename from src/data/components/patching.py rename to src/grass_mil/data/components/patching.py diff --git a/src/data/components/precompute.py b/src/grass_mil/data/components/precompute.py similarity index 99% rename from src/data/components/precompute.py rename to src/grass_mil/data/components/precompute.py index f6b7108..41285ac 100644 --- a/src/data/components/precompute.py +++ b/src/grass_mil/data/components/precompute.py @@ -685,8 +685,8 @@ def _attach_categorical_indices( encoded_cols.append(encoded) slices[label_name] = col_idx - categorical_index = np.stack(encoded_cols, axis=1) - data.categorical_index = torch.from_numpy(categorical_index).long() + 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 diff --git a/src/data/components/samplers.py b/src/grass_mil/data/components/samplers.py similarity index 86% rename from src/data/components/samplers.py rename to src/grass_mil/data/components/samplers.py index b5ce44b..5780700 100644 --- a/src/data/components/samplers.py +++ b/src/grass_mil/data/components/samplers.py @@ -1,13 +1,25 @@ from __future__ import annotations import copy -from dataclasses import dataclass +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 torch_geometric.loader import DataLoader as PyGDataLoader @@ -54,6 +66,8 @@ class RuntimeShadowConfig: min_weight: float = 1e-6 subsample_fraction: float = 1.0 subsample_seed: Optional[int] = None + 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": @@ -97,38 +111,54 @@ def from_dict(cls, cfg: Optional[Dict[str, Any]]) -> "RuntimeShadowConfig": min_weight=min_weight, subsample_fraction=subsample_fraction, subsample_seed=subsample_seed, + interior_seeds=InteriorSeedConfig.from_dict(cfg.get("interior_seeds")), + ego_radius=EgoRadiusConfig.from_dict(cfg.get("ego_radius")), ) def _build_weighted_node_idx( - data: Data, *, property_name: str, weight_mode: str, min_weight: float + 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_index") or not hasattr( + if not hasattr(data, "categorical_codes") or not hasattr( data, "categorical_slices" ): return None - categorical_index = getattr(data, "categorical_index") + 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_index, torch.Tensor) - or categorical_index.numel() == 0 + not isinstance(categorical_codes, torch.Tensor) + or categorical_codes.numel() == 0 ): return None - if categorical_index.dim() != 2: + if categorical_codes.dim() != 2: return None prop_col = int(categorical_slices[property_name]) - if prop_col < 0 or prop_col >= int(categorical_index.size(1)): + if prop_col < 0 or prop_col >= int(categorical_codes.size(1)): return None - node_labels = categorical_index[:, prop_col].long().view(-1) + 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 ) @@ -155,11 +185,66 @@ def _build_weighted_node_idx( 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 - return torch.multinomial(probs, num_samples=num_samples, replacement=True).long() + 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: @@ -245,6 +330,11 @@ def __init__( 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: @@ -257,12 +347,18 @@ def __len__(self) -> int: 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 @@ -292,12 +388,18 @@ def __iter__(self): 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 @@ -326,7 +428,9 @@ def __iter__(self): node_idx=effective_root_candidates, replace=bool(self.runtime.replace), shuffle=bool(self.runtime.shuffle_subgraphs and self.shuffle), - transform=self.kwargs.get("transform"), + transform=_compose_ego_radius_transform( + self.kwargs.get("transform"), self.runtime + ), **unit_loader_kwargs, ) for sub_batch in loader: @@ -813,8 +917,12 @@ def _apply_transform_to_each_subgraph(self, batch: Batch) -> Batch: elif val.dim() == 0: setattr(sub_data, key, val) elif val.dim() > 0 and val.size(0) == batch.x.size(0): - # Node-level feature tensors are already represented by `x`. - continue + # 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 "edge_attr" in batch and val.size(0) == batch.edge_index.size(1): continue else: diff --git a/src/grass_mil/data/components/seed_sampling.py b/src/grass_mil/data/components/seed_sampling.py new file mode 100644 index 0000000..ec0cca7 --- /dev/null +++ b/src/grass_mil/data/components/seed_sampling.py @@ -0,0 +1,277 @@ +"""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/data/components/spatial_types.py b/src/grass_mil/data/components/spatial_types.py similarity index 100% rename from src/data/components/spatial_types.py rename to src/grass_mil/data/components/spatial_types.py diff --git a/src/data/components/transforms.py b/src/grass_mil/data/components/transforms.py similarity index 93% rename from src/data/components/transforms.py rename to src/grass_mil/data/components/transforms.py index 558e262..b4faf55 100644 --- a/src/data/components/transforms.py +++ b/src/grass_mil/data/components/transforms.py @@ -50,17 +50,17 @@ def __call__(self, data): if isinstance(labels, torch.Tensor) else np.asarray(labels) ) - elif hasattr(data, "categorical_index") and hasattr(data, "categorical_slices"): + 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_index[:, col_idx].detach().cpu().numpy() + labels_np = data.categorical_codes[:, col_idx].detach().cpu().numpy() else: raise AttributeError( - f"Missing label attribute '{self.label_attr}' or categorical_index for composition vector" + f"Missing label attribute '{self.label_attr}' or categorical_codes for composition vector" ) labels_np = labels_np.reshape(-1) diff --git a/src/data/spatial_omics_datamodule.py b/src/grass_mil/data/spatial_omics_datamodule.py similarity index 100% rename from src/data/spatial_omics_datamodule.py rename to src/grass_mil/data/spatial_omics_datamodule.py diff --git a/src/eval.py b/src/grass_mil/eval.py similarity index 90% rename from src/eval.py rename to src/grass_mil/eval.py index b7e09f9..288d5f6 100644 --- a/src/eval.py +++ b/src/grass_mil/eval.py @@ -1,14 +1,15 @@ from typing import Any, Dict, List, Tuple import hydra -import rootutils from lightning import LightningDataModule, LightningModule, Trainer from lightning.pytorch.loggers import Logger from omegaconf import DictConfig -rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True) +from grass_mil._paths import setup_project_root -from src.utils import ( +setup_project_root() + +from grass_mil.utils import ( RankedLogger, extras, instantiate_loggers, @@ -60,7 +61,7 @@ def evaluate(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: return metric_dict, object_dict -@hydra.main(version_base="1.3", config_path="../configs", config_name="eval.yaml") +@hydra.main(version_base="1.3", config_path="configs", config_name="eval.yaml") def main(cfg: DictConfig) -> None: """Entry point for evaluation.""" extras(cfg) diff --git a/src/grass_mil/inference/__init__.py b/src/grass_mil/inference/__init__.py new file mode 100644 index 0000000..b5db7e3 --- /dev/null +++ b/src/grass_mil/inference/__init__.py @@ -0,0 +1,11 @@ +from grass_mil.inference.aggregation import aggregate_group_logits +from grass_mil.inference.collectors import collect_inference_payload +from grass_mil.inference.embeddings import embeddings_to_dataframe +from grass_mil.inference.metrics import compute_task_metrics + +__all__ = [ + "aggregate_group_logits", + "collect_inference_payload", + "compute_task_metrics", + "embeddings_to_dataframe", +] diff --git a/src/inference/aggregation.py b/src/grass_mil/inference/aggregation.py similarity index 95% rename from src/inference/aggregation.py rename to src/grass_mil/inference/aggregation.py index 19e7f99..7e74ec0 100644 --- a/src/inference/aggregation.py +++ b/src/grass_mil/inference/aggregation.py @@ -5,7 +5,8 @@ import torch -from src.inference.schemas import AggregatedPredictionPayload, BatchPredictionPayload +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: @@ -254,9 +255,17 @@ def aggregate_group_logits( ) 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).reshape(-1) + 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.unsqueeze(-1)).sum(dim=0, keepdim=True) + 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()) diff --git a/src/inference/collectors.py b/src/grass_mil/inference/collectors.py similarity index 99% rename from src/inference/collectors.py rename to src/grass_mil/inference/collectors.py index 0ec8b66..32941dc 100644 --- a/src/inference/collectors.py +++ b/src/grass_mil/inference/collectors.py @@ -6,7 +6,7 @@ import torch from lightning import LightningDataModule, LightningModule, Trainer -from src.inference.schemas import ( +from grass_mil.inference.schemas import ( BatchPredictionPayload, CollectedInferencePayload, EmbeddingPayload, diff --git a/src/inference/embeddings.py b/src/grass_mil/inference/embeddings.py similarity index 87% rename from src/inference/embeddings.py rename to src/grass_mil/inference/embeddings.py index f7e0291..6910e34 100644 --- a/src/inference/embeddings.py +++ b/src/grass_mil/inference/embeddings.py @@ -2,11 +2,11 @@ import pandas as pd -from src.inference.io import ( +from grass_mil.inference.io import ( embeddings_to_dataframe as _embeddings_to_dataframe, node_embeddings_to_dataframe as _node_embeddings_to_dataframe, ) -from src.inference.schemas import EmbeddingPayload +from grass_mil.inference.schemas import EmbeddingPayload def embedding_payload_to_table(payload: EmbeddingPayload) -> pd.DataFrame: diff --git a/src/inference/interpretability_export.py b/src/grass_mil/inference/interpretability_export.py similarity index 89% rename from src/inference/interpretability_export.py rename to src/grass_mil/inference/interpretability_export.py index 0cf6792..b2dfd14 100644 --- a/src/inference/interpretability_export.py +++ b/src/grass_mil/inference/interpretability_export.py @@ -8,7 +8,7 @@ import pandas as pd import torch -from src.inference.schemas import BatchPredictionPayload +from grass_mil.inference.schemas import BatchPredictionPayload def _tensor_to_2d_cpu(values: Optional[torch.Tensor]) -> Optional[torch.Tensor]: @@ -75,26 +75,32 @@ def _make_unique_ids(values: Sequence[str]) -> List[str]: def _softmax_by_group(logits: torch.Tensor, groups: Sequence[str]) -> np.ndarray: - if logits.ndim != 2 or logits.shape[1] != 1: + """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 with exactly one column. " - "Multi-column attention logits are ambiguous for interpretability export; " - "emit a single attention logit per instance." + "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),), dtype=np.float64) + 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, 0] + chunk = logits[indices, :] scores = torch.softmax(chunk, dim=0).cpu().numpy().astype(np.float64) - for local, global_idx in enumerate(indices): - attention[global_idx] = float(scores[local]) + attention[indices, :] = scores return attention @@ -183,6 +189,7 @@ def build_instance_table( score_mode: str = "sigmoid", score_logit_index: int = 0, attention_column: str = "attention", + attention_class_index: Optional[int] = None, ) -> pd.DataFrame: logits = _tensor_to_2d_cpu(payload.instance_logits) if logits is None: @@ -232,7 +239,24 @@ def build_instance_table( attention_logits = _tensor_to_2d_cpu(payload.instance_attention_logits) if attention_logits is not None: - data[attention_column] = _softmax_by_group(attention_logits, bag_ids).tolist() + 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() diff --git a/src/inference/io.py b/src/grass_mil/inference/io.py similarity index 97% rename from src/inference/io.py rename to src/grass_mil/inference/io.py index e3de6d7..7c73dde 100644 --- a/src/inference/io.py +++ b/src/grass_mil/inference/io.py @@ -8,7 +8,7 @@ import pandas as pd import torch -from src.inference.schemas import BatchPredictionPayload, EmbeddingPayload +from grass_mil.inference.schemas import BatchPredictionPayload, EmbeddingPayload def ensure_output_dir(base_output_dir: Path, output_subdir: str) -> Path: diff --git a/src/inference/metrics.py b/src/grass_mil/inference/metrics.py similarity index 86% rename from src/inference/metrics.py rename to src/grass_mil/inference/metrics.py index 95f6444..eaee6c2 100644 --- a/src/inference/metrics.py +++ b/src/grass_mil/inference/metrics.py @@ -7,7 +7,7 @@ from omegaconf import DictConfig from sklearn.metrics import balanced_accuracy_score, roc_auc_score -from src.inference.schemas import BatchPredictionPayload +from grass_mil.inference.schemas import BatchPredictionPayload def _safe_div(num: float, den: float) -> float: @@ -163,28 +163,48 @@ def _regression_metrics(logits: torch.Tensor, targets: torch.Tensor) -> Dict[str def _survival_concordance_index( - risk_scores: torch.Tensor, times: torch.Tensor, events: torch.Tensor + 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 - n = int(r.shape[0]) - for i in range(n): - if e[i].item() <= 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 - for j in range(n): - if i == j: - continue - if t[i].item() < t[j].item(): - comparable += 1.0 - if r[i].item() > r[j].item(): - concordant += 1.0 - elif r[i].item() == r[j].item(): - ties += 1.0 + 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) diff --git a/src/inference/predict.py b/src/grass_mil/inference/predict.py similarity index 96% rename from src/inference/predict.py rename to src/grass_mil/inference/predict.py index 3b20041..88f51bb 100644 --- a/src/inference/predict.py +++ b/src/grass_mil/inference/predict.py @@ -4,21 +4,22 @@ from typing import Any, Dict, List, Tuple import hydra -import rootutils from lightning import LightningDataModule, LightningModule, Trainer from lightning.pytorch.loggers import Logger from omegaconf import DictConfig -rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True) +from grass_mil._paths import setup_project_root -from src.inference.aggregation import aggregate_group_logits # noqa: E402 -from src.inference.collectors import collect_inference_payload # noqa: E402 -from src.inference.interpretability_export import ( # noqa: E402 +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_spatial_table, resolve_composition_column_names, ) -from src.inference.io import ( # noqa: E402 +from grass_mil.inference.io import ( # noqa: E402 build_summary_payload, embeddings_to_dataframe, ensure_output_dir, @@ -27,12 +28,12 @@ write_dataframe, write_json, ) -from src.inference.metrics import compute_task_metrics # noqa: E402 -from src.inference.schemas import ( # noqa: E402 +from grass_mil.inference.metrics import compute_task_metrics # noqa: E402 +from grass_mil.inference.schemas import ( # noqa: E402 BatchPredictionPayload, PreforwardSubsamplingDecision, ) -from src.utils import ( # noqa: E402 +from grass_mil.utils import ( # noqa: E402 RankedLogger, extras, instantiate_loggers, @@ -421,6 +422,7 @@ def predict(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: 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"), ) instance_path = output_dir / str( (interpret_cfg or {}).get("instance_filename", "instance_table.csv") @@ -465,7 +467,7 @@ def predict(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: @hydra.main( version_base="1.3", - config_path="../../configs", + config_path="../configs", config_name="inference/predict.yaml", ) def main(cfg: DictConfig) -> None: diff --git a/src/inference/schemas.py b/src/grass_mil/inference/schemas.py similarity index 100% rename from src/inference/schemas.py rename to src/grass_mil/inference/schemas.py diff --git a/src/grass_mil/interpretability/__init__.py b/src/grass_mil/interpretability/__init__.py new file mode 100644 index 0000000..9a59eef --- /dev/null +++ b/src/grass_mil/interpretability/__init__.py @@ -0,0 +1,39 @@ +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 +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", + "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 new file mode 100644 index 0000000..2bf512d --- /dev/null +++ b/src/grass_mil/interpretability/core/__init__.py @@ -0,0 +1,27 @@ +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 new file mode 100644 index 0000000..bbb030b --- /dev/null +++ b/src/grass_mil/interpretability/core/agreement.py @@ -0,0 +1,103 @@ +"""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 new file mode 100644 index 0000000..7566099 --- /dev/null +++ b/src/grass_mil/interpretability/core/attribution.py @@ -0,0 +1,347 @@ +"""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 new file mode 100644 index 0000000..11e7a25 --- /dev/null +++ b/src/grass_mil/interpretability/core/biomarkers.py @@ -0,0 +1,66 @@ +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/interpretability/core/clustering.py b/src/grass_mil/interpretability/core/clustering.py similarity index 97% rename from src/interpretability/core/clustering.py rename to src/grass_mil/interpretability/core/clustering.py index 732ba3a..ac10e78 100644 --- a/src/interpretability/core/clustering.py +++ b/src/grass_mil/interpretability/core/clustering.py @@ -11,7 +11,7 @@ SpectralClustering, ) -from src.interpretability.contracts import ClusteringResult +from grass_mil.contracts import ClusteringResult def run_clustering( diff --git a/src/interpretability/core/data.py b/src/grass_mil/interpretability/core/data.py similarity index 96% rename from src/interpretability/core/data.py rename to src/grass_mil/interpretability/core/data.py index 978faaa..a39278d 100644 --- a/src/interpretability/core/data.py +++ b/src/grass_mil/interpretability/core/data.py @@ -6,7 +6,7 @@ import numpy as np import pandas as pd -from src.interpretability.contracts import EmbeddingSet, InterpretabilityDataset +from grass_mil.contracts import EmbeddingSet, InterpretabilityDataset def load_table(path: Path) -> pd.DataFrame: diff --git a/src/grass_mil/interpretability/core/niches.py b/src/grass_mil/interpretability/core/niches.py new file mode 100644 index 0000000..9933ed8 --- /dev/null +++ b/src/grass_mil/interpretability/core/niches.py @@ -0,0 +1,167 @@ +"""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, Optional, Sequence + +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/interpretability/core/reduction.py b/src/grass_mil/interpretability/core/reduction.py similarity index 96% rename from src/interpretability/core/reduction.py rename to src/grass_mil/interpretability/core/reduction.py index 88fca8a..c0e889a 100644 --- a/src/interpretability/core/reduction.py +++ b/src/grass_mil/interpretability/core/reduction.py @@ -6,7 +6,7 @@ from sklearn.decomposition import PCA from sklearn.manifold import TSNE -from src.interpretability.contracts import ReductionResult +from grass_mil.contracts import ReductionResult def run_reduction( diff --git a/src/interpretability/core/transfer.py b/src/grass_mil/interpretability/core/transfer.py similarity index 75% rename from src/interpretability/core/transfer.py rename to src/grass_mil/interpretability/core/transfer.py index 0edf4b3..cdc8ebe 100644 --- a/src/interpretability/core/transfer.py +++ b/src/grass_mil/interpretability/core/transfer.py @@ -8,17 +8,17 @@ import pandas as pd from sklearn.neighbors import KNeighborsClassifier -from src.interpretability.contracts import ( - ClusterTransferBundle, +from grass_mil.contracts import ( + NicheTransferBundle, InterpretabilityDataset, ReportBundle, ) -from src.interpretability.core.data import extract_embedding_set +from grass_mil.interpretability.core.data import extract_embedding_set _TRANSFER_BUNDLE_SCHEMA_VERSION = 1 -def fit_cluster_transfer_from_report_bundle( +def fit_niche_transfer_from_report_bundle( dataset: InterpretabilityDataset, report_bundle: ReportBundle, *, @@ -26,55 +26,53 @@ def fit_cluster_transfer_from_report_bundle( n_neighbors: int = 15, weights: str = "uniform", metric: str = "minkowski", -) -> ClusterTransferBundle: +) -> NicheTransferBundle: if report_bundle.clustering is None: - raise ValueError( - "Cannot fit cluster transfer: report bundle is missing clustering result." - ) + 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 cluster transfer: clustering labels must be 1D.") + raise ValueError("Cannot fit niche transfer: clustering labels must be 1D.") if emb_set.matrix.shape[0] != labels.shape[0]: raise ValueError( - "Cannot fit cluster transfer: embedding rows and clustering labels length mismatch." + "Cannot fit niche transfer: embedding rows and clustering labels length mismatch." ) if emb_set.matrix.shape[0] == 0: - raise ValueError("Cannot fit cluster transfer on an empty instance table.") + raise ValueError("Cannot fit niche transfer on an empty instance table.") n_neighbors = int(n_neighbors) if n_neighbors < 1: - raise ValueError("Cannot fit cluster transfer: n_neighbors must be >= 1.") + raise ValueError("Cannot fit niche transfer: n_neighbors must be >= 1.") if n_neighbors > int(labels.shape[0]): raise ValueError( - "Cannot fit cluster transfer: n_neighbors cannot exceed number of training rows." + "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.cluster_feature_reduction is None: + if report_bundle.niche_feature_reduction is None: raise ValueError( - "Cannot fit cluster transfer: cluster_on_pca=true but " - "report bundle is missing cluster_feature_reduction." + "Cannot fit niche transfer: cluster_on_pca=true but " + "report bundle is missing niche_feature_reduction." ) - pca_result = report_bundle.cluster_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 cluster transfer: cluster_feature_reduction fitted object " + "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 cluster transfer: cluster_feature_reduction embedding must be 2D." + "Cannot fit niche transfer: niche_feature_reduction embedding must be 2D." ) if train_features.shape[0] != labels.shape[0]: raise ValueError( - "Cannot fit cluster transfer: cluster_feature_reduction rows and clustering " + "Cannot fit niche transfer: niche_feature_reduction rows and clustering " "labels length mismatch." ) @@ -90,7 +88,7 @@ def fit_cluster_transfer_from_report_bundle( "cluster_on_pca": bool(cluster_on_pca), "cluster_method": report_bundle.clustering.method, } - return ClusterTransferBundle( + return NicheTransferBundle( embedding_columns=list(emb_set.feature_columns), id_column=str(dataset.id_column), cluster_on_pca=cluster_on_pca, @@ -101,8 +99,8 @@ def fit_cluster_transfer_from_report_bundle( ) -def apply_cluster_transfer( - bundle: ClusterTransferBundle, +def apply_niche_transfer( + bundle: NicheTransferBundle, instance_table: pd.DataFrame, *, id_column: str | None = None, @@ -110,13 +108,11 @@ def apply_cluster_transfer( output_confidence_column: str = "transfer_confidence", ) -> pd.DataFrame: if len(instance_table) == 0: - raise ValueError("Cannot apply cluster transfer to an empty instance table.") + 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 cluster transfer: missing id column {resolved_id_column!r}." - ) + 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 @@ -124,13 +120,13 @@ def apply_cluster_transfer( if missing_embeddings: preview = ", ".join(sorted(missing_embeddings)[:8]) raise ValueError( - "Cannot apply cluster transfer: instance table is missing embedding columns: " + "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 cluster transfer: transfer bundle is missing a fitted kNN model." + "Cannot apply niche transfer: transfer bundle is missing a fitted kNN model." ) query_matrix = instance_table[bundle.embedding_columns].to_numpy(dtype=float) @@ -138,7 +134,7 @@ def apply_cluster_transfer( if bundle.pca_model is not None: if not hasattr(bundle.pca_model, "transform"): raise ValueError( - "Cannot apply cluster transfer: transfer bundle PCA model is missing " + "Cannot apply niche transfer: transfer bundle PCA model is missing " "transform(...)." ) query_features = np.asarray(bundle.pca_model.transform(query_features)) @@ -159,7 +155,7 @@ def apply_cluster_transfer( ) -def save_cluster_transfer_bundle(bundle: ClusterTransferBundle, path: str | Path) -> None: +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 = { @@ -175,7 +171,7 @@ def save_cluster_transfer_bundle(bundle: ClusterTransferBundle, path: str | Path joblib.dump(payload, output_path) -def load_cluster_transfer_bundle(path: str | Path) -> ClusterTransferBundle: +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.") @@ -210,7 +206,7 @@ def load_cluster_transfer_bundle(path: str | Path) -> ClusterTransferBundle: if knn_model is None or not hasattr(knn_model, "predict"): raise ValueError("Invalid transfer bundle format: missing or invalid knn_model.") - return ClusterTransferBundle( + 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"]), diff --git a/src/interpretability/pipeline.py b/src/grass_mil/interpretability/pipeline.py similarity index 78% rename from src/interpretability/pipeline.py rename to src/grass_mil/interpretability/pipeline.py index 2e16b93..0e6e31e 100644 --- a/src/interpretability/pipeline.py +++ b/src/grass_mil/interpretability/pipeline.py @@ -1,5 +1,7 @@ from __future__ import annotations +from dataclasses import replace + import json from pathlib import Path from typing import Any, Dict, Optional @@ -7,19 +9,20 @@ import numpy as np import pandas as pd -from src.interpretability.contracts import ( - ClusterSummary, +from grass_mil.contracts import ( + NicheSummary, InterpretabilityDataset, PluginResult, ReportBundle, ) -from src.interpretability.core.biomarkers import cluster_biomarker_summary -from src.interpretability.core.clustering import run_clustering -from src.interpretability.core.data import extract_embedding_set -from src.interpretability.core.reduction import run_reduction -from src.interpretability.plugins.base import PluginContext -from src.interpretability.plugins.builtin import create_builtin_registry -from src.interpretability.plugins.registry import PluginRegistry +from grass_mil.interpretability.core.biomarkers import niche_composition_summary +from grass_mil.interpretability.core.niches import relabel_niches_by_composition +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( @@ -131,7 +134,7 @@ def run_interpretability_pipeline( if clustering_enabled: method = str(clustering_cfg.get("method", "kmeans")) params = dict(clustering_cfg.get("params", {})) - # Contract: when cluster_on_pca is disabled, cluster on raw embeddings. + # 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 @@ -147,9 +150,21 @@ def run_interpretability_pipeline( clustering_result = run_clustering(method, clustering_input, **params) labels = clustering_result.labels - cluster_summary: Optional[ClusterSummary] = None + 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: - cluster_summary = cluster_biomarker_summary( + niche_summary = niche_composition_summary( dataset.instance_table, labels, composition_prefix=str(config.get("composition_prefix", "comp_")), @@ -160,8 +175,8 @@ def run_interpretability_pipeline( state={ "reduction": reduction_result, "clustering": clustering_result, - "cluster_labels": labels, - "cluster_summary": cluster_summary, + "niche_labels": labels, + "niche_summary": niche_summary, } ) plugin_results: Dict[str, PluginResult] = {} @@ -183,8 +198,8 @@ def run_interpretability_pipeline( # Persist core artifacts base_instance = dataset.instance_table.copy() if labels is not None: - base_instance["cluster_label"] = labels - _write_table(base_instance, artifacts_dir / "instance_with_clusters.csv") + 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])] @@ -193,13 +208,13 @@ def run_interpretability_pipeline( ) _write_table(red_df, artifacts_dir / "reduction.csv") if clustering_result is not None: - cluster_df = pd.DataFrame( + niche_df = pd.DataFrame( { dataset.id_column: dataset.instance_table[dataset.id_column].astype(str).tolist(), - "cluster_label": clustering_result.labels, + "niche_label": clustering_result.labels, } ) - _write_table(cluster_df, artifacts_dir / "cluster_labels.csv") + _write_table(niche_df, artifacts_dir / "niche_labels.csv") summary_payload: Dict[str, Any] = { "rows": int(len(dataset.instance_table)), @@ -208,15 +223,16 @@ def run_interpretability_pipeline( "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, - cluster_feature_reduction=combo_pca_result if cluster_on_pca else None, + niche_feature_reduction=combo_pca_result if cluster_on_pca else None, clustering=clustering_result, - cluster_summary=cluster_summary, + 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 new file mode 100644 index 0000000..30bab8c --- /dev/null +++ b/src/grass_mil/interpretability/plugins/__init__.py @@ -0,0 +1,15 @@ +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/interpretability/plugins/base.py b/src/grass_mil/interpretability/plugins/base.py similarity index 84% rename from src/interpretability/plugins/base.py rename to src/grass_mil/interpretability/plugins/base.py index 525524f..d778118 100644 --- a/src/interpretability/plugins/base.py +++ b/src/grass_mil/interpretability/plugins/base.py @@ -3,7 +3,7 @@ from dataclasses import dataclass, field from typing import Any, Dict, List, Protocol -from src.interpretability.contracts import InterpretabilityDataset, PluginResult +from grass_mil.contracts import InterpretabilityDataset, PluginResult @dataclass diff --git a/src/grass_mil/interpretability/plugins/builtin.py b/src/grass_mil/interpretability/plugins/builtin.py new file mode 100644 index 0000000..335cc78 --- /dev/null +++ b/src/grass_mil/interpretability/plugins/builtin.py @@ -0,0 +1,658 @@ +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.autocorrelation import ( + diff_morans_i_vs_reference, + run_morans_i, +) +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.ripley import aggregate_ripley +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 NeighborhoodEnrichmentPlugin(InterpretabilityPlugin): + name: str = "neighborhood_enrichment" + + 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 neighborhood enrichment plugin.") + table = dataset.instance_table.copy() + table["niche_label"] = np.asarray(context.state["niche_labels"]) + result = run_neighborhood_enrichment( + table, + dataset.spatial_table, + label_column=str(params.get("label_column", "niche_label")), + id_column=str(params.get("id_column", dataset.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")), + ) + payload = { + "enrichment": result.enrichment, + "observed": result.observed, + "expected": result.expected, + "pvalues": result.pvalues, + } + section_tables = { + "enrichment": result.enrichment, + "observed": result.observed, + "expected": result.expected, + } + if result.pvalues is not None: + section_tables["pvalues"] = result.pvalues + sections = [ + ReportSection( + title="Neighborhood Enrichment", + description="Niche adjacency enrichment relative to expected connectivity.", + tables=section_tables, + ) + ] + return PluginResult(name=self.name, payload=payload, sections=sections) + + +@dataclass +class DiffNeighborhoodEnrichmentPlugin(InterpretabilityPlugin): + name: str = "diff_neighborhood_enrichment" + + 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 differential neighborhood enrichment.") + table = dataset.instance_table.copy() + table["niche_label"] = np.asarray(context.state["niche_labels"]) + condition_column = str( + params.get("condition_column", dataset.condition_column or "condition") + ) + pairwise = run_diff_neighborhood_enrichment( + table, + dataset.spatial_table, + label_column=str(params.get("label_column", "niche_label")), + condition_column=condition_column, + permutation_group_column=str(params.get("permutation_group_column", "sample_id")), + id_column=str(params.get("id_column", dataset.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")), + ) + payload = { + "enrichment_by_pair": {pair: out.enrichment for pair, out in pairwise.items()}, + "observed_by_pair": {pair: out.observed for pair, out in pairwise.items()}, + "expected_by_pair": {pair: out.expected for pair, out in pairwise.items()}, + "pvalues_by_pair": { + pair: out.pvalues for pair, out in pairwise.items() if out.pvalues is not None + }, + } + sections: List[ReportSection] = [] + for pair, out in pairwise.items(): + section_tables = { + "enrichment": out.enrichment, + "observed": out.observed, + "expected": out.expected, + } + if out.pvalues is not None: + section_tables["pvalues"] = out.pvalues + sections.append( + ReportSection( + title=f"Differential Neighborhood Enrichment ({pair})", + description=( + "Pairwise condition differential neighborhood enrichment " + "(left condition minus right condition)." + ), + tables=section_tables, + ) + ) + return PluginResult(name=self.name, payload=payload, sections=sections) + + +@dataclass +class FiltrationCurvesPlugin(InterpretabilityPlugin): + name: str = "filtration_curves" + + 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 filtration curves plugin.") + table = dataset.instance_table.copy() + table["niche_label"] = np.asarray(context.state["niche_labels"]) + thresholds = _resolve_filtration_thresholds(dict(params)) + result = compute_filtration_curves( + table, + dataset.spatial_table, + thresholds=thresholds, + niche_column=str(params.get("niche_column", "niche_label")), + cell_type_column=str(params.get("cell_type_column", "cell_type")), + id_column=str(params.get("id_column", dataset.id_column)), + distance_column=str(params.get("distance_column", "distance")), + scale_within_cluster=bool(params.get("scale_within_cluster", True)), + ) + payload: Dict[str, Any] = {"thresholds": result.thresholds, "curves": result.curves} + sections = [ + ReportSection( + title="Filtration Curves", + description="Distance-threshold accumulation by niche and cell type.", + metadata={"thresholds": result.thresholds.tolist()}, + ) + ] + 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"), + "cluster_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) + + +@dataclass +class PerNicheCellTypeEnrichmentPlugin(InterpretabilityPlugin): + """Compute cell-type neighborhood enrichment within each niche separately. + + For each niche, subsets the node table and spatial table to nodes belonging + to that niche, then runs neighborhood enrichment with label_column=cell_type. + This answers: "within niche C, which cell types are spatially co-located?" + """ + + name: str = "per_niche_cell_type_enrichment" + + 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 per-niche cell-type enrichment.") + cell_type_column = str( + params.get("cell_type_column", dataset.cell_type_column or "cell_type") + ) + niche_column = str(params.get("niche_column", "niche_label")) + id_column = str(params.get("id_column", dataset.id_column)) + + table = dataset.instance_table.copy() + table[niche_column] = np.asarray(context.state["niche_labels"]) + + unique_niches = sorted(set(table[niche_column])) + skip_background = bool(params.get("skip_background", True)) + + enrichment_by_niche: Dict[str, pd.DataFrame] = {} + sections: List[ReportSection] = [] + + for niche_id in unique_niches: + if skip_background and niche_id == -1: + continue + niche_mask = table[niche_column] == niche_id + niche_table = table[niche_mask].copy() + if len(niche_table) < 2: + continue + + niche_node_ids = set(niche_table[id_column].astype(str)) + spatial = dataset.spatial_table.copy() + niche_spatial = spatial[ + spatial["source_id"].astype(str).isin(niche_node_ids) + & spatial["target_id"].astype(str).isin(niche_node_ids) + ] + if niche_spatial.empty: + continue + + try: + result = run_neighborhood_enrichment( + niche_table, + niche_spatial, + label_column=cell_type_column, + id_column=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", "obs-exp")), + ) + enrichment_by_niche[str(niche_id)] = result.enrichment + sections.append( + ReportSection( + title=f"Cell-Type Enrichment — Niche {niche_id}", + description=( + f"Cell-type neighborhood enrichment within niche {niche_id}." + ), + tables={"enrichment": result.enrichment}, + ) + ) + except (ValueError, KeyError): + continue + + payload: Dict[str, Any] = {"enrichment_by_niche": enrichment_by_niche} + 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 + + +@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." + ) + + result = niche_attribution_summary( + table, + labels, + attention_columns=attention_columns, + logit_columns=logit_columns, + bag_id_column=str(params.get("bag_id_column", dataset.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)), + 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 MoransIPlugin(InterpretabilityPlugin): + """Spatial autocorrelation per niche, with a permutation null.""" + + name: str = "morans_i" + + 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("Moran's I requires a spatial table.") + table = dataset.instance_table.copy() + niche_column = str(params.get("niche_column", "niche_label")) + table[niche_column] = np.asarray(context.state["niche_labels"]) + + feature_columns = _resolve_columns( + table, + str(params.get("feature_prefix", "comp_")), + params.get("feature_columns"), + ) + if not feature_columns: + raise ValueError("Moran's I found no feature columns to evaluate.") + + result = run_morans_i( + table, + dataset.spatial_table, + feature_columns=feature_columns, + group_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)), + min_nodes=int(params.get("min_nodes", 3)), + ) + payload: Dict[str, Any] = { + "statistic": result.statistic, + "pvalue": result.pvalue, + "qvalue": result.qvalue, + "n_nodes": result.n_nodes, + } + tables = {"morans_i": result.statistic, "qvalue": result.qvalue} + + reference = params.get("reference_group", -1) + if reference in result.statistic.index: + differential = diff_morans_i_vs_reference(result, reference_group=reference) + payload["differential_z"] = differential + tables["differential_z"] = differential + + sections = [ + ReportSection( + title="Spatial Autocorrelation (Moran's I)", + description=( + "Global Moran's I per niche over the instance graph, with a " + "permutation null and BH-FDR adjustment." + ), + tables=tables, + ) + ] + return PluginResult(name=self.name, payload=payload, sections=sections) + + +@dataclass +class RipleyPlugin(InterpretabilityPlugin): + """Centred cross-L curves over instance centroids.""" + + name: str = "ripley" + + def required_inputs(self) -> List[str]: + return ["niche_labels"] + + def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: + table = dataset.instance_table.copy() + label_column = str(params.get("label_column", "niche_label")) + if label_column == "niche_label": + table[label_column] = np.asarray(context.state["niche_labels"]) + + pairs = params.get("pairs") + if pairs is not None: + pairs = [(str(a), str(b)) for a, b in pairs] + + result = aggregate_ripley( + table, + label_column=label_column, + group_column=params.get("group_column"), + x_column=str(params.get("x_column", "center_x")), + y_column=str(params.get("y_column", "center_y")), + pairs=pairs, + 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")), + ) + curves = result.to_frame() + payload = { + "radii": result.radii, + "curves": curves, + "n_groups": result.n_groups, + } + sections = [ + ReportSection( + title="Ripley Cross-L", + description=( + "Centred cross-L, L(r) - r: zero under complete spatial randomness, " + "positive under clustering, negative under regularity." + ), + tables={"cross_l": curves}, + ) + ] + 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(NeighborhoodEnrichmentPlugin()) + registry.register(DiffNeighborhoodEnrichmentPlugin()) + registry.register(FiltrationCurvesPlugin()) + registry.register(TissueGraphPlugin()) + registry.register(PerNicheCellTypeEnrichmentPlugin()) + registry.register(MarginAttributionPlugin()) + registry.register(MoransIPlugin()) + registry.register(RipleyPlugin()) + registry.register(NicheAgreementPlugin()) + + +def create_builtin_registry() -> PluginRegistry: + registry = create_plugin_registry() + register_builtin_plugins(registry) + return registry diff --git a/src/interpretability/plugins/registry.py b/src/grass_mil/interpretability/plugins/registry.py similarity index 95% rename from src/interpretability/plugins/registry.py rename to src/grass_mil/interpretability/plugins/registry.py index 870236f..656cc09 100644 --- a/src/interpretability/plugins/registry.py +++ b/src/grass_mil/interpretability/plugins/registry.py @@ -2,7 +2,7 @@ from typing import Dict, Iterable, List -from src.interpretability.plugins.base import InterpretabilityPlugin +from grass_mil.interpretability.plugins.base import InterpretabilityPlugin class PluginRegistry: diff --git a/src/interpretability/report_cli.py b/src/grass_mil/interpretability/report_cli.py similarity index 85% rename from src/interpretability/report_cli.py rename to src/grass_mil/interpretability/report_cli.py index 1c89839..9fd8ecb 100644 --- a/src/interpretability/report_cli.py +++ b/src/grass_mil/interpretability/report_cli.py @@ -6,14 +6,15 @@ from typing import Any, Dict, Tuple import hydra -import rootutils from omegaconf import DictConfig, OmegaConf -rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True) +from grass_mil._paths import setup_project_root -from src.interpretability.core.data import load_interpretability_dataset # noqa: E402 -from src.interpretability.pipeline import run_interpretability_pipeline # noqa: E402 -from src.interpretability.reporting.render import render_interpretability_report # noqa: E402 +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__) @@ -65,7 +66,7 @@ def run_report(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: @hydra.main( version_base="1.3", - config_path="../../configs", + config_path="../configs", config_name="interpretability/report.yaml", ) def main(cfg: DictConfig) -> None: diff --git a/src/grass_mil/interpretability/reporting/__init__.py b/src/grass_mil/interpretability/reporting/__init__.py new file mode 100644 index 0000000..14e9546 --- /dev/null +++ b/src/grass_mil/interpretability/reporting/__init__.py @@ -0,0 +1,3 @@ +from grass_mil.interpretability.reporting.render import render_interpretability_report + +__all__ = ["render_interpretability_report"] diff --git a/src/interpretability/reporting/html.py b/src/grass_mil/interpretability/reporting/html.py similarity index 97% rename from src/interpretability/reporting/html.py rename to src/grass_mil/interpretability/reporting/html.py index c26bf7c..4613a91 100644 --- a/src/interpretability/reporting/html.py +++ b/src/grass_mil/interpretability/reporting/html.py @@ -6,7 +6,7 @@ import pandas as pd -from src.interpretability.contracts import ReportSection +from grass_mil.contracts import ReportSection def _render_table(title: str, table: pd.DataFrame) -> str: diff --git a/src/interpretability/reporting/pdf.py b/src/grass_mil/interpretability/reporting/pdf.py similarity index 100% rename from src/interpretability/reporting/pdf.py rename to src/grass_mil/interpretability/reporting/pdf.py diff --git a/src/grass_mil/interpretability/reporting/plotly_builders.py b/src/grass_mil/interpretability/reporting/plotly_builders.py new file mode 100644 index 0000000..4079e3a --- /dev/null +++ b/src/grass_mil/interpretability/reporting/plotly_builders.py @@ -0,0 +1,492 @@ +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/interpretability/reporting/render.py b/src/grass_mil/interpretability/reporting/render.py similarity index 79% rename from src/interpretability/reporting/render.py rename to src/grass_mil/interpretability/reporting/render.py index 2e0d1be..065a207 100644 --- a/src/interpretability/reporting/render.py +++ b/src/grass_mil/interpretability/reporting/render.py @@ -4,12 +4,12 @@ from pathlib import Path from typing import Any, Dict -from src.interpretability.contracts import ReportBundle -from src.interpretability.reporting.html import render_html_report -from src.interpretability.reporting.pdf import render_pdf_via_playwright -from src.interpretability.reporting.plotly_builders import bundle_figures -from src.interpretability.reporting.sections import build_report_sections -from src.interpretability.reporting.snapshot import export_plotly_snapshots +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__) diff --git a/src/interpretability/reporting/sections.py b/src/grass_mil/interpretability/reporting/sections.py similarity index 61% rename from src/interpretability/reporting/sections.py rename to src/grass_mil/interpretability/reporting/sections.py index f83eb93..9797347 100644 --- a/src/interpretability/reporting/sections.py +++ b/src/grass_mil/interpretability/reporting/sections.py @@ -2,7 +2,7 @@ from typing import List -from src.interpretability.contracts import ReportBundle, ReportSection +from grass_mil.contracts import ReportBundle, ReportSection def build_report_sections(bundle: ReportBundle) -> List[ReportSection]: @@ -15,14 +15,14 @@ def build_report_sections(bundle: ReportBundle) -> List[ReportSection]: metadata=bundle.metadata, ) ) - if bundle.cluster_summary is not None: + if bundle.niche_summary is not None: sections.append( ReportSection( - title="Cluster Summary", - description="Cluster counts and enrichment tables.", + title="Niche Summary", + description="Niche counts and enrichment tables.", tables={ - "cluster_counts": bundle.cluster_summary.cluster_counts.to_frame(name="count"), - "enrichment": bundle.cluster_summary.enrichment, + "niche_counts": bundle.niche_summary.niche_counts.to_frame(name="count"), + "enrichment": bundle.niche_summary.enrichment, }, ) ) diff --git a/src/interpretability/reporting/snapshot.py b/src/grass_mil/interpretability/reporting/snapshot.py similarity index 100% rename from src/interpretability/reporting/snapshot.py rename to src/grass_mil/interpretability/reporting/snapshot.py diff --git a/src/interpretability/tier2/__init__.py b/src/grass_mil/interpretability/tier2/__init__.py similarity index 61% rename from src/interpretability/tier2/__init__.py rename to src/grass_mil/interpretability/tier2/__init__.py index 5b261f0..dafdaff 100644 --- a/src/interpretability/tier2/__init__.py +++ b/src/grass_mil/interpretability/tier2/__init__.py @@ -1,9 +1,9 @@ -from src.interpretability.tier2.filtration import compute_filtration_curves -from src.interpretability.tier2.neighborhood import ( +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 src.interpretability.tier2.tissue_graph import ( +from grass_mil.interpretability.tier2.tissue_graph import ( build_tissue_graph_figure, prepare_tissue_graph_view, ) diff --git a/src/grass_mil/interpretability/tier2/autocorrelation.py b/src/grass_mil/interpretability/tier2/autocorrelation.py new file mode 100644 index 0000000..296dd7c --- /dev/null +++ b/src/grass_mil/interpretability/tier2/autocorrelation.py @@ -0,0 +1,245 @@ +"""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, List, Optional, Sequence + +import numpy as np +import pandas as pd + +__all__ = ["MoranResult", "morans_i", "run_morans_i", "diff_morans_i_vs_reference"] + + +@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), + ) + + +def diff_morans_i_vs_reference( + result: MoranResult, + *, + reference_group: object = -1, +) -> pd.DataFrame: + """Differential autocorrelation against a reference group. + + ``z = (I_k - I_ref) / sqrt(var_k + var_ref)``, using the permutation-null + variances. The reference is typically the Background niche. + """ + if reference_group not in result.statistic.index: + raise ValueError( + f"reference_group {reference_group!r} is not present; available: " + f"{list(result.statistic.index)}." + ) + ref_stat = result.statistic.loc[reference_group] + ref_var = result.variance.loc[reference_group] + + rows: List[pd.Series] = [] + index: List[object] = [] + for group in result.statistic.index: + if group == reference_group: + continue + delta = result.statistic.loc[group] - ref_stat + pooled = np.sqrt(result.variance.loc[group].astype(float) + ref_var.astype(float)) + rows.append(delta / pooled.replace(0.0, np.nan)) + index.append(group) + if not rows: + return pd.DataFrame(columns=result.statistic.columns) + return pd.DataFrame(rows, index=index) diff --git a/src/interpretability/tier2/filtration.py b/src/grass_mil/interpretability/tier2/filtration.py similarity index 83% rename from src/interpretability/tier2/filtration.py rename to src/grass_mil/interpretability/tier2/filtration.py index a2d2761..65f8132 100644 --- a/src/interpretability/tier2/filtration.py +++ b/src/grass_mil/interpretability/tier2/filtration.py @@ -3,7 +3,7 @@ import numpy as np import pandas as pd -from src.interpretability.contracts import FiltrationCurvesResult +from grass_mil.contracts import FiltrationCurvesResult def compute_filtration_curves( @@ -11,14 +11,14 @@ def compute_filtration_curves( spatial_table: pd.DataFrame, *, thresholds: np.ndarray, - cluster_column: str = "cluster_label", + 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 cluster_column not in node_table.columns: - raise ValueError(f"Missing cluster column {cluster_column!r}.") + 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: @@ -28,17 +28,17 @@ def compute_filtration_curves( 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, cluster_column, cell_type_column]].copy() + 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[cluster_column] = nodes[cluster_column].astype(str) + 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)[cluster_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) @@ -47,7 +47,7 @@ def compute_filtration_curves( edges[distance_column] = edges[distance_column].astype(float) edges["source_cluster"] = edges["source_id"].map(node_to_cluster) - clusters = sorted(nodes[cluster_column].astype(str).unique()) + 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) @@ -58,11 +58,11 @@ def compute_filtration_curves( 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 clusters + c: {ct: np.zeros_like(thresholds, dtype=float) for ct in cell_types} for c in niches } - for cluster in clusters: - c_edges = edges[edges["source_cluster"] == str(cluster)] + 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) @@ -100,11 +100,11 @@ def compute_filtration_curves( 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[cluster][ct] = per_ct_original[ct_idx].astype(float, copy=False) + curves[niche][ct] = per_ct_original[ct_idx].astype(float, copy=False) if scale_within_cluster: for ct in cell_types: - vmax = float(curves[cluster][ct].max()) + vmax = float(curves[niche][ct].max()) if vmax > 0: - curves[cluster][ct] = curves[cluster][ct] / vmax + curves[niche][ct] = curves[niche][ct] / vmax return FiltrationCurvesResult(thresholds=thresholds, curves=curves) diff --git a/src/interpretability/tier2/neighborhood.py b/src/grass_mil/interpretability/tier2/neighborhood.py similarity index 98% rename from src/interpretability/tier2/neighborhood.py rename to src/grass_mil/interpretability/tier2/neighborhood.py index c9fd9f5..c8cb6f2 100644 --- a/src/interpretability/tier2/neighborhood.py +++ b/src/grass_mil/interpretability/tier2/neighborhood.py @@ -7,7 +7,7 @@ import numpy as np import pandas as pd -from src.interpretability.contracts import NeighborhoodEnrichmentResult +from grass_mil.contracts import NeighborhoodEnrichmentResult def _coerce_edges(spatial_table: pd.DataFrame) -> pd.DataFrame: @@ -77,8 +77,8 @@ def _analytical_expected_and_std( if n_nodes == 0 or n_edges == 0: zeros = pd.DataFrame(0.0, index=categories, columns=categories) return zeros, zeros - cluster_counts = node_labels.value_counts().reindex(categories, fill_value=0.0).astype(float) - probs = cluster_counts.values / float(n_nodes) + 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) diff --git a/src/grass_mil/interpretability/tier2/ripley.py b/src/grass_mil/interpretability/tier2/ripley.py new file mode 100644 index 0000000..14a1e79 --- /dev/null +++ b/src/grass_mil/interpretability/tier2/ripley.py @@ -0,0 +1,227 @@ +"""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/interpretability/tier2/tissue_graph.py b/src/grass_mil/interpretability/tier2/tissue_graph.py similarity index 96% rename from src/interpretability/tier2/tissue_graph.py rename to src/grass_mil/interpretability/tier2/tissue_graph.py index 1f6abdd..c29659e 100644 --- a/src/interpretability/tier2/tissue_graph.py +++ b/src/grass_mil/interpretability/tier2/tissue_graph.py @@ -33,7 +33,7 @@ def prepare_tissue_graph_view( id_column: str, x_column: str = "center_x", y_column: str = "center_y", - label_column: str = "cluster_label", + label_column: str = "niche_label", coerce_ids_to_str: bool = True, include_edge_distances: bool = True, ) -> TissueGraphView: @@ -91,7 +91,7 @@ def prepare_tissue_graph_view( edge_cols.append("distance") edges = edges[edge_cols].reset_index(drop=True) - cluster_counts = ( + niche_counts = ( nodes[str(label_column)].astype(str).value_counts().sort_index().astype(int).to_dict() ) metadata: Dict[str, Any] = { @@ -103,7 +103,7 @@ def prepare_tissue_graph_view( "label_column": str(label_column), "node_count": int(len(nodes)), "edge_count": int(len(edges)), - "cluster_counts": cluster_counts, + "niche_counts": niche_counts, } return TissueGraphView(nodes=nodes.reset_index(drop=True), edges=edges, metadata=metadata) @@ -130,7 +130,7 @@ def build_tissue_graph_figure( nodes = view.nodes.copy() metadata = dict(view.metadata) - label_column = str(metadata.get("label_column", "cluster_label")) + 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")) @@ -178,7 +178,7 @@ def build_tissue_graph_figure( "color": label_codes, "colorscale": colorscale, "showscale": True, - "colorbar": {"title": "cluster"}, + "colorbar": {"title": "niche"}, }, customdata=np.stack( [ diff --git a/src/loocv.py b/src/grass_mil/loocv.py similarity index 96% rename from src/loocv.py rename to src/grass_mil/loocv.py index 3b1dcf0..e2fbac4 100644 --- a/src/loocv.py +++ b/src/grass_mil/loocv.py @@ -12,15 +12,16 @@ import hydra import lightning as L import pandas as pd -import rootutils from hydra.utils import instantiate from omegaconf import DictConfig, open_dict -rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True) +from grass_mil._paths import setup_project_root -from src.data.components.precompute import discover_viable_loocv_fold_ids # noqa: E402 -from src.train import train # noqa: E402 -from src.utils import RankedLogger, extras # noqa: E402 +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) @@ -274,7 +275,7 @@ def run_loocv(cfg: DictConfig) -> Dict[str, Any]: return summary -@hydra.main(version_base="1.3", config_path="../configs", config_name="loocv.yaml") +@hydra.main(version_base="1.3", config_path="configs", config_name="loocv.yaml") def main(cfg: DictConfig) -> None: extras(cfg) run_loocv(cfg) diff --git a/src/models/__init__.py b/src/grass_mil/models/__init__.py similarity index 91% rename from src/models/__init__.py rename to src/grass_mil/models/__init__.py index 697d481..f3fab0b 100644 --- a/src/models/__init__.py +++ b/src/grass_mil/models/__init__.py @@ -8,7 +8,6 @@ GraphPooling, GraphPredictionHead, MLPPredictor, - NodePredictionHead, WeightedBCEWithLogitsLoss, WeightedCrossEntropyLoss, WeightedHuberLoss, @@ -17,7 +16,6 @@ build_encoder, build_graph_head, build_loss, - build_node_head, build_pooling, build_ssl, ) @@ -36,7 +34,6 @@ "GraphPooling", "GraphPredictionHead", "MLPPredictor", - "NodePredictionHead", "WeightedBCEWithLogitsLoss", "WeightedCrossEntropyLoss", "WeightedHuberLoss", @@ -45,7 +42,6 @@ "build_encoder", "build_graph_head", "build_loss", - "build_node_head", "build_pooling", "build_ssl", ] diff --git a/src/models/bgrl_module.py b/src/grass_mil/models/bgrl_module.py similarity index 62% rename from src/models/bgrl_module.py rename to src/grass_mil/models/bgrl_module.py index 6c936e8..a2f1d51 100644 --- a/src/models/bgrl_module.py +++ b/src/grass_mil/models/bgrl_module.py @@ -9,14 +9,29 @@ from .components import EncoderConfig, build_encoder, build_ssl from .training import ( CosineWarmup, - augment_graph, + 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) + return list(names) if names else None + + class BGRLModule(L.LightningModule): def __init__( self, @@ -42,27 +57,52 @@ def __init__( self.init_strict = init_strict self.encoder_init_map = encoder_init_map - self._drop_edge_p1 = float(self.task_cfg.get("drop_edge_p1", 0.0)) - self._drop_edge_p2 = float(self.task_cfg.get("drop_edge_p2", 0.0)) - self._drop_feat_p1 = float(self.task_cfg.get("drop_feat_p1", 0.0)) - self._drop_feat_p2 = float(self.task_cfg.get("drop_feat_p2", 0.0)) + # 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 self._momentum_base = float(self.task_cfg.get("momentum", 0.99)) self._momentum_min = float(self.task_cfg.get("momentum_min", 0.99)) self._warmup_steps = int(self.task_cfg.get("warmup_steps", 0)) self._total_steps = int(self.task_cfg.get("total_steps", 0)) self._momentum_scheduler: Optional[CosineWarmup] = None + # 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 = dict(self.hparams.encoder) - if encoder_cfg.get("input_dim", 0) in (None, 0): - inferred = infer_encoder_input_dim(self) - if inferred <= 0: - raise ValueError("BGRLModule could not infer encoder.input_dim from datamodule.") - encoder_cfg["input_dim"] = inferred + 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.") @@ -97,19 +137,31 @@ def _encode_tuple(self, batch): 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 = augment_graph( - batch, - drop_edge_p=self._drop_edge_p1, - drop_feat_p=self._drop_feat_p1, - ) - view2 = augment_graph( - batch, - drop_edge_p=self._drop_edge_p2, - drop_feat_p=self._drop_feat_p2, - ) + 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 = ( diff --git a/src/models/components/__init__.py b/src/grass_mil/models/components/__init__.py similarity index 84% rename from src/models/components/__init__.py rename to src/grass_mil/models/components/__init__.py index 134b8cb..2d89b69 100644 --- a/src/models/components/__init__.py +++ b/src/grass_mil/models/components/__init__.py @@ -1,15 +1,15 @@ 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_node_head, build_pooling, build_ssl, ) -from .heads import GraphPredictionHead, NodePredictionHead +from .heads import GraphPredictionHead from .losses import ( CoxSGDLoss, WeightedBCEWithLogitsLoss, @@ -25,12 +25,13 @@ "AttnNetGatedProjected", "BGRL", "CoxSGDLoss", + "CategoricalEmbeddingConfig", "EncoderConfig", + "NodeInputEmbedding", "GNNEncoder", "GraphPooling", "GraphPredictionHead", "MLPPredictor", - "NodePredictionHead", "WeightedBCEWithLogitsLoss", "WeightedCrossEntropyLoss", "WeightedHuberLoss", @@ -39,7 +40,6 @@ "build_encoder", "build_graph_head", "build_loss", - "build_node_head", "build_pooling", "build_ssl", ] diff --git a/src/models/components/attention.py b/src/grass_mil/models/components/attention.py similarity index 62% rename from src/models/components/attention.py rename to src/grass_mil/models/components/attention.py index f586965..8266146 100644 --- a/src/models/components/attention.py +++ b/src/grass_mil/models/components/attention.py @@ -29,6 +29,19 @@ def forward(self, x): logits = self.attention_c(a * b) return logits, 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 the primary clustering space for the interpretability suite because + it is low-dimensional and directly upstream of the attention weights; + clustering the saturated output collapses the extremes together. + """ + pre_tanh = self.attention_a[0](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 variant of gated MIL attention network.""" @@ -58,3 +71,18 @@ def forward(self, x): h = self.projection(x) logits = self.attention_c(self.attention_a(h) * self.attention_b(h)) return logits, x + + 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. It is the primary clustering space for the interpretability + suite: 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/models/components/backbones.py b/src/grass_mil/models/components/backbones.py similarity index 73% rename from src/models/components/backbones.py rename to src/grass_mil/models/components/backbones.py index d7ee95e..05178b1 100644 --- a/src/models/components/backbones.py +++ b/src/grass_mil/models/components/backbones.py @@ -2,17 +2,20 @@ import warnings from dataclasses import dataclass -from typing import Literal, Optional +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 src.models.components.backbones") from exc + raise ImportError( + "torch_geometric is required for grass_mil.models.components.backbones" + ) from exc NormType = Literal["batchnorm", "layernorm", "none"] @@ -32,11 +35,15 @@ class EncoderConfig: jk: JKType = "last" act: Literal["relu", "gelu", "leaky_relu"] = "relu" pooling: Optional[str] = None - set2set_steps: int = 3 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 def _get_activation(name: str) -> nn.Module: @@ -71,10 +78,10 @@ def __init__(self, cfg: EncoderConfig): self.act = _get_activation(cfg.act) self.dropout = nn.Dropout(cfg.dropout) - self.input_proj = ( - nn.Identity() - if cfg.input_dim == cfg.hidden_dim - else nn.Linear(cfg.input_dim, cfg.hidden_dim) + self.input_proj = NodeInputEmbedding( + input_dim=cfg.input_dim, + hidden_dim=cfg.hidden_dim, + categorical=CategoricalEmbeddingConfig.from_dict(cfg.categorical_embedding), ) self.layers = nn.ModuleList([self._build_conv_layer() for _ in range(cfg.num_layers)]) @@ -92,7 +99,6 @@ def __init__(self, cfg: EncoderConfig): GraphPooling( name=cfg.pooling, input_dim=final_dim, - set2set_steps=cfg.set2set_steps, ) if cfg.pooling is not None else None @@ -141,6 +147,22 @@ def _supports_edge_attr(self) -> bool: 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}'. " @@ -174,9 +196,29 @@ def _apply_conv( 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, edge_attr) + 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": @@ -195,9 +237,10 @@ def forward( 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) + 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) diff --git a/src/grass_mil/models/components/embeddings.py b/src/grass_mil/models/components/embeddings.py new file mode 100644 index 0000000..82bc6ed --- /dev/null +++ b/src/grass_mil/models/components/embeddings.py @@ -0,0 +1,213 @@ +"""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. + """ + + label: str = "cell_type" + num_embeddings: Optional[int] = None + column_index: Optional[int] = None + reserve_unassigned: bool = True + + @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"} + 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, + ) -> None: + super().__init__() + if input_dim < 0: + raise ValueError(f"input_dim must be >= 0, got {input_dim}.") + self.input_dim = int(input_dim) + self.hidden_dim = int(hidden_dim) + self.categorical = categorical + + 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}." + ) + 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 + + @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, +) -> Optional[CategoricalEmbeddingConfig]: + """Fill in vocabulary size and column index from datamodule metadata. + + Explicit config values always win; ``None`` triggers inference, mirroring + the ``encoder.input_dim`` contract. + """ + 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." + ) + 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, + ) diff --git a/src/models/components/factory.py b/src/grass_mil/models/components/factory.py similarity index 94% rename from src/models/components/factory.py rename to src/grass_mil/models/components/factory.py index 8de3e66..21085d7 100644 --- a/src/models/components/factory.py +++ b/src/grass_mil/models/components/factory.py @@ -4,7 +4,7 @@ from .attention import AttnNetGated, AttnNetGatedProjected from .backbones import EncoderConfig, GNNEncoder -from .heads import GraphPredictionHead, NodePredictionHead +from .heads import GraphPredictionHead from .losses import ( CoxSGDLoss, WeightedBCEWithLogitsLoss, @@ -42,10 +42,6 @@ def build_graph_head(config: Dict[str, Any]) -> GraphPredictionHead: return GraphPredictionHead(**config) -def build_node_head(config: Dict[str, Any]) -> NodePredictionHead: - return NodePredictionHead(**config) - - def build_ssl(use_ssl: bool, encoder, ssl_config: Optional[Dict[str, Any]] = None): if not use_ssl: return None diff --git a/src/models/components/heads.py b/src/grass_mil/models/components/heads.py similarity index 95% rename from src/models/components/heads.py rename to src/grass_mil/models/components/heads.py index aad991f..8c5aaec 100644 --- a/src/models/components/heads.py +++ b/src/grass_mil/models/components/heads.py @@ -32,7 +32,3 @@ def forward(self, x): class GraphPredictionHead(_MLPHead): pass - - -class NodePredictionHead(_MLPHead): - pass diff --git a/src/models/components/losses.py b/src/grass_mil/models/components/losses.py similarity index 100% rename from src/models/components/losses.py rename to src/grass_mil/models/components/losses.py diff --git a/src/grass_mil/models/components/pooling.py b/src/grass_mil/models/components/pooling.py new file mode 100644 index 0000000..46b7e9e --- /dev/null +++ b/src/grass_mil/models/components/pooling.py @@ -0,0 +1,44 @@ +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/models/components/ssl.py b/src/grass_mil/models/components/ssl.py similarity index 100% rename from src/models/components/ssl.py rename to src/grass_mil/models/components/ssl.py diff --git a/src/models/supervised_module.py b/src/grass_mil/models/supervised_module.py similarity index 85% rename from src/models/supervised_module.py rename to src/grass_mil/models/supervised_module.py index 2b46678..be447c5 100644 --- a/src/models/supervised_module.py +++ b/src/grass_mil/models/supervised_module.py @@ -5,23 +5,19 @@ import lightning as L import torch import torch.distributed as dist -from src.inference.aggregation import aggregate_group_logits -from src.inference.metrics import compute_task_metrics_from_tensors -from src.inference.schemas import BatchPredictionPayload +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_mil_aux_targets, build_supervised_components, - compute_aux_node_loss, compute_binary_accuracy, compute_categorical_accuracy, - compute_entropy_regularization, compute_supervised_loss, extract_bag_ids, gather_bag_targets, - gather_instance_logits, group_instance_indices_by_bag, instantiate_optimizer, instantiate_scheduler_with_warmup, @@ -98,19 +94,13 @@ def __init__( "task.region_accumulation.hyperbatch_size must be >= 1 " "when region accumulation is enabled." ) - self.node_aux_cfg = dict(self.task_cfg.get("node_aux", {})) - self.entropy_reg_cfg = dict(self.task_cfg.get("entropy_reg", {})) - self.loss_weights_cfg = dict(self.task_cfg.get("loss_weights", {})) else: self.region_accum_cfg = {} self.region_accum_enabled = False self.region_accum_hyperbatch_size = 0 self.region_accum_flush_on_epoch_end = False - self.node_aux_cfg = {} - self.entropy_reg_cfg = {} - self.loss_weights_cfg = {} - self._region_total_loss_buffer: list[torch.Tensor] = [] + 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. @@ -183,6 +173,7 @@ def _forward_instances(self, batch): 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) @@ -249,17 +240,17 @@ def _compute_instance_composition( label_name: str, device: torch.device, ) -> Optional[torch.Tensor]: - categorical_index = getattr(batch, "categorical_index", None) - if not isinstance(categorical_index, torch.Tensor) or categorical_index.numel() == 0: + categorical_codes = getattr(batch, "categorical_codes", None) + if not isinstance(categorical_codes, torch.Tensor) or categorical_codes.numel() == 0: return None - categorical_index = categorical_index.to(device=device).long() - if categorical_index.ndim != 2: + 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_index.shape[1]): + if col_idx is None or col_idx < 0 or col_idx >= int(categorical_codes.shape[1]): return None - node_labels = categorical_index[:, int(col_idx)] + node_labels = categorical_codes[:, int(col_idx)] if node_labels.numel() == 0: return None batch_index = getattr(batch, "batch", None) @@ -322,6 +313,14 @@ def _compute_instance_centroids( 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, @@ -424,7 +423,9 @@ def _build_predict_group_metadata( if compute_attention_logits: sub_emb = graph_emb.index_select(0, idx) attn_logits, _ = self.attention(sub_emb) - instance_attention_logits_chunks.append(attn_logits.reshape(-1, 1)) + 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) @@ -459,7 +460,7 @@ def _build_predict_group_metadata( metadata["instance_attention_logits"] = ( torch.cat(instance_attention_logits_chunks, dim=0) if instance_attention_logits_chunks - else torch.empty((0, 1), device=patch_logits.device) + else torch.empty((0, self._attention_width()), device=patch_logits.device) ) if instance_graph_chunks: metadata["instance_graphs"] = instance_graph_chunks @@ -519,6 +520,7 @@ def collect_graph_embeddings( 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, ) bag_ids = extract_bag_ids( @@ -625,89 +627,6 @@ def _compute_region_loss( bag_weights=bag_weights, ) - def _compute_optional_terms( - self, - *, - patch_logits: torch.Tensor, - bag_targets: torch.Tensor, - bag_ids: list[str], - bag_indices: list[list[int]], - bag_attention: dict[str, torch.Tensor], - stage: str, - ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: - """Compute node auxiliary loss and entropy regularisation. - - Only active during training when MIL attention is enabled and the - relevant sub-configs are turned on. - """ - if not self.use_attention: - return None, None - if stage != "train" or self.task_cfg.get("target_type") != "binary": - return None, None - - node_aux_enabled = bool(self.node_aux_cfg.get("enabled", False)) - entropy_enabled = bool(self.entropy_reg_cfg.get("enabled", False)) - if not node_aux_enabled and not entropy_enabled: - return None, None - - aux_logits = gather_instance_logits(patch_logits, bag_indices) - aux_targets, aux_weights = build_mil_aux_targets( - bag_targets=bag_targets, - bag_indices=bag_indices, - bag_ids=bag_ids, - bag_attention=bag_attention, - target_mode=str(self.node_aux_cfg.get("target_mode", "attention_shaped_ti")), - ) - - node_aux_loss = None - if node_aux_enabled: - loss_mode = str(self.node_aux_cfg.get("loss_mode", "bce")) - node_aux_loss = compute_aux_node_loss( - aux_logits=aux_logits, - aux_targets=aux_targets, - aux_weights=aux_weights if loss_mode == "weighted_bce" else None, - loss_mode=loss_mode, - ) - - entropy_reg = None - if entropy_enabled: - entropy_mode = str(self.entropy_reg_cfg.get("mode", "attention_shaped_target")) - if entropy_mode == "attention": - entropy_values = torch.cat( - [bag_attention[bag_id].reshape(-1, 1) for bag_id in bag_ids], - dim=0, - ) - elif entropy_mode == "attention_shaped_target": - entropy_values = aux_targets - else: - raise ValueError( - "entropy_reg.mode must be one of " "['attention', 'attention_shaped_target']" - ) - entropy_reg = compute_entropy_regularization(entropy_values) - - return node_aux_loss, entropy_reg - - def _compose_total_loss( - self, - region_loss: torch.Tensor, - node_aux_loss: Optional[torch.Tensor], - entropy_reg: Optional[torch.Tensor], - ) -> torch.Tensor: - if not self.use_attention: - return region_loss - - region_w = float(self.loss_weights_cfg.get("region", 1.0)) - node_w = float(self.node_aux_cfg.get("weight", self.loss_weights_cfg.get("node_aux", 1.0))) - entropy_w = float( - self.entropy_reg_cfg.get("weight", self.loss_weights_cfg.get("entropy", 0.0)) - ) - total = region_w * region_loss - if node_aux_loss is not None: - total = total + node_w * node_aux_loss - if entropy_reg is not None: - total = total + entropy_w * entropy_reg - return total - def _compute_losses( self, *, @@ -719,22 +638,13 @@ def _compute_losses( bag_indices: list[list[int]], bag_attention: Optional[dict[str, torch.Tensor]], stage: str, - ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], torch.Tensor]: - region_loss = self._compute_region_loss( + ) -> 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, ) - node_aux_loss, entropy_reg = self._compute_optional_terms( - patch_logits=patch_logits, - bag_targets=bag_targets, - bag_ids=ordered_bag_ids, - bag_indices=bag_indices, - bag_attention=bag_attention or {}, - stage=stage, - ) - total_loss = self._compose_total_loss(region_loss, node_aux_loss, entropy_reg) - return region_loss, node_aux_loss, entropy_reg, total_loss def _log_stage_metrics( self, @@ -744,8 +654,6 @@ def _log_stage_metrics( bag_logits: torch.Tensor, bag_targets: torch.Tensor, region_loss: Optional[torch.Tensor] = None, - node_aux_loss: Optional[torch.Tensor] = None, - entropy_reg: Optional[torch.Tensor] = None, ) -> None: if stage not in {"train", "val", "test"}: return @@ -764,22 +672,6 @@ def _log_stage_metrics( on_epoch=True, prog_bar=False, ) - if node_aux_loss is not None: - self.log( - "train/node_aux_loss", - node_aux_loss, - on_step=True, - on_epoch=True, - prog_bar=False, - ) - if entropy_reg is not None: - self.log( - "train/entropy_reg", - entropy_reg, - 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( @@ -813,7 +705,7 @@ def _shared_step(self, batch, stage: str): bag_targets = payload["bag_targets"] bag_weights = payload["bag_weights"] - region_loss, node_aux_loss, entropy_reg, total_loss = self._compute_losses( + region_loss = self._compute_losses( patch_logits=patch_logits, bag_logits=bag_logits, bag_targets=bag_targets, @@ -825,16 +717,14 @@ def _shared_step(self, batch, stage: str): ) self._log_stage_metrics( stage=stage, - loss_value=total_loss if stage == "train" else region_loss, + loss_value=region_loss, bag_logits=bag_logits, bag_targets=bag_targets, region_loss=region_loss if self.use_attention else None, - node_aux_loss=node_aux_loss, - entropy_reg=entropy_reg, ) result = { - "loss": total_loss if stage == "train" else region_loss, + "loss": region_loss, "bag_ids": list(ordered_bag_ids), "bag_logits": bag_logits, "bag_targets": bag_targets, @@ -845,13 +735,7 @@ def _shared_step(self, batch, stage: str): "bag_attention": bag_attention, } if self.use_attention: - result.update( - { - "region_loss": region_loss, - "node_aux_loss": node_aux_loss, - "entropy_reg": entropy_reg, - } - ) + result["region_loss"] = region_loss return result # ------------------------------------------------------------------ @@ -873,12 +757,10 @@ def training_step(self, batch, batch_idx): patch_logits = payload["patch_logits"] region_losses: list[torch.Tensor] = [] - total_losses: list[torch.Tensor] = [] - node_losses: list[torch.Tensor] = [] - entropies: 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, node_aux_loss, entropy_reg, total_loss = self._compute_losses( + region_loss = self._compute_losses( patch_logits=patch_logits, bag_logits=bag_logits[idx : idx + 1], bag_targets=bag_targets[idx : idx + 1], @@ -889,19 +771,15 @@ def training_step(self, batch, batch_idx): stage="train", ) region_losses.append(region_loss.detach()) - total_losses.append(total_loss) - if node_aux_loss is not None: - node_losses.append(node_aux_loss.detach()) - if entropy_reg is not None: - entropies.append(entropy_reg.detach()) + buffered_region_losses.append(region_loss) - if not total_losses: + 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_total_loss = torch.stack(total_losses).mean() - self._region_total_loss_buffer.append(step_total_loss) + 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: @@ -914,26 +792,22 @@ def training_step(self, batch, batch_idx): ) self.log( "train/region_buffer_size", - float(len(self._region_total_loss_buffer)), + float(len(self._region_loss_buffer)), on_step=True, on_epoch=False, prog_bar=False, ) - mean_total_loss = step_total_loss.detach() + mean_region_loss_value = step_region_loss.detach() mean_region_loss = torch.stack(region_losses).mean() - mean_node_aux = torch.stack(node_losses).mean() if node_losses else None - mean_entropy = torch.stack(entropies).mean() if entropies else None self._log_stage_metrics( stage="train", - loss_value=mean_total_loss, + loss_value=mean_region_loss_value, bag_logits=bag_logits, bag_targets=bag_targets, region_loss=mean_region_loss, - node_aux_loss=mean_node_aux, - entropy_reg=mean_entropy, ) - return mean_total_loss + return mean_region_loss_value def validation_step(self, batch, batch_idx): out = self._shared_step(batch, "val") @@ -1026,9 +900,9 @@ def predict_step(self, batch, batch_idx, dataloader_idx=0): # ------------------------------------------------------------------ def _flush_region_buffer_if_needed(self, *, force: bool) -> list[torch.Tensor]: - if not self._region_total_loss_buffer: + if not self._region_loss_buffer: return [] - if not force and len(self._region_total_loss_buffer) < self.region_accum_hyperbatch_size: + if not force and len(self._region_loss_buffer) < self.region_accum_hyperbatch_size: return [] optimizer = self.optimizers() @@ -1036,16 +910,16 @@ def _flush_region_buffer_if_needed(self, *, force: bool) -> list[torch.Tensor]: optimizer = optimizer[0] flush_losses: list[torch.Tensor] = [] - while self._region_total_loss_buffer: + while self._region_loss_buffer: if force: - chunk = self._region_total_loss_buffer + chunk = self._region_loss_buffer else: - if len(self._region_total_loss_buffer) < self.region_accum_hyperbatch_size: + if len(self._region_loss_buffer) < self.region_accum_hyperbatch_size: break - chunk = self._region_total_loss_buffer[: self.region_accum_hyperbatch_size] + chunk = self._region_loss_buffer[: self.region_accum_hyperbatch_size] - total_loss = torch.stack(chunk, dim=0).mean() - self.manual_backward(total_loss) + flush_loss = torch.stack(chunk, dim=0).mean() + self.manual_backward(flush_loss) optimizer.step() optimizer.zero_grad(set_to_none=True) @@ -1064,9 +938,9 @@ def _flush_region_buffer_if_needed(self, *, force: bool) -> list[torch.Tensor]: on_epoch=force, prog_bar=False, ) - flush_losses.append(total_loss.detach()) + flush_losses.append(flush_loss.detach()) - del self._region_total_loss_buffer[: len(chunk)] + del self._region_loss_buffer[: len(chunk)] if force: break @@ -1074,7 +948,7 @@ def _flush_region_buffer_if_needed(self, *, force: bool) -> list[torch.Tensor]: def on_train_start(self) -> None: if self.region_accum_enabled: - self._region_total_loss_buffer.clear() + self._region_loss_buffer.clear() def on_validation_epoch_start(self) -> None: self._reset_val_epoch_buffers() diff --git a/src/models/training/__init__.py b/src/grass_mil/models/training/__init__.py similarity index 72% rename from src/models/training/__init__.py rename to src/grass_mil/models/training/__init__.py index 5343d9c..e847a6b 100644 --- a/src/models/training/__init__.py +++ b/src/grass_mil/models/training/__init__.py @@ -7,30 +7,37 @@ maybe_sample_indices, select_target_columns, ) +from .augmentations import ( + build_augmentation, + degree_importance, + drop_edges, + drop_features, + drop_importance, + feature_noise, +) 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 ( - build_mil_aux_targets, - compute_aux_node_loss, compute_binary_accuracy, compute_categorical_accuracy, - compute_entropy_regularization, compute_supervised_loss, - gather_instance_logits, ) from .optimization import ( instantiate_optimizer, instantiate_scheduler, instantiate_scheduler_with_warmup, ) -from .ssl_runtime import CosineWarmup, augment_graph +from .ssl_runtime import CosineWarmup __all__ = [ "aggregate_bag_logits_attention", @@ -41,20 +48,24 @@ "maybe_sample_indices", "select_target_columns", "build_supervised_components", + "build_augmentation", + "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", - "build_mil_aux_targets", - "compute_aux_node_loss", "compute_binary_accuracy", "compute_categorical_accuracy", - "compute_entropy_regularization", "compute_supervised_loss", - "gather_instance_logits", "instantiate_optimizer", "instantiate_scheduler", "instantiate_scheduler_with_warmup", "CosineWarmup", - "augment_graph", ] diff --git a/src/grass_mil/models/training/augmentations.py b/src/grass_mil/models/training/augmentations.py new file mode 100644 index 0000000..15b7650 --- /dev/null +++ b/src/grass_mil/models/training/augmentations.py @@ -0,0 +1,296 @@ +"""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, Dict, 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 + + draws = torch.rand((edge_index.size(1),), generator=generator, device=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): + """Zero whole feature columns with probability ``p``.""" + if p <= 0.0 or 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, device=x.device) + x[:, draws < float(p)] = 0 + out.x = x + return out + + +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/models/training/bagging.py b/src/grass_mil/models/training/bagging.py similarity index 87% rename from src/models/training/bagging.py rename to src/grass_mil/models/training/bagging.py index 691077f..d4135d8 100644 --- a/src/models/training/bagging.py +++ b/src/grass_mil/models/training/bagging.py @@ -4,6 +4,8 @@ import torch +from grass_mil.contracts import validate_attention_width + try: from torch_geometric.data import Data except Exception as exc: # pragma: no cover @@ -130,6 +132,23 @@ def aggregate_bag_logits_attention( 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 = {} @@ -143,8 +162,15 @@ def aggregate_bag_logits_attention( sub_emb = embeddings[selected] sub_logits = logits[selected] attn_logits, _ = attention(sub_emb) - attn_scores = torch.softmax(attn_logits.squeeze(-1), dim=0) - weighted_logits = (sub_logits * attn_scores.unsqueeze(-1)).sum(dim=0, keepdim=True) + 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) diff --git a/src/grass_mil/models/training/builders.py b/src/grass_mil/models/training/builders.py new file mode 100644 index 0000000..57b5158 --- /dev/null +++ b/src/grass_mil/models/training/builders.py @@ -0,0 +1,186 @@ +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) + ) + 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, + ) + cfg["categorical_embedding"] = asdict(resolved) + return cfg + + +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/models/training/checkpoint_init.py b/src/grass_mil/models/training/checkpoint_init.py similarity index 59% rename from src/models/training/checkpoint_init.py rename to src/grass_mil/models/training/checkpoint_init.py index 94b052b..e555018 100644 --- a/src/models/training/checkpoint_init.py +++ b/src/grass_mil/models/training/checkpoint_init.py @@ -45,3 +45,41 @@ def remap_encoder_keys( new_key = f"encoder.{key[len('model.online_encoder.') :]}" remapped[new_key] = value return remapped + + +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") + 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/loss_utils.py b/src/grass_mil/models/training/loss_utils.py new file mode 100644 index 0000000..b78882d --- /dev/null +++ b/src/grass_mil/models/training/loss_utils.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Sequence + +import torch +import torch.nn.functional as F + + +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/models/training/optimization.py b/src/grass_mil/models/training/optimization.py similarity index 100% rename from src/models/training/optimization.py rename to src/grass_mil/models/training/optimization.py diff --git a/src/grass_mil/models/training/ssl_runtime.py b/src/grass_mil/models/training/ssl_runtime.py new file mode 100644 index 0000000..f1bad8c --- /dev/null +++ b/src/grass_mil/models/training/ssl_runtime.py @@ -0,0 +1,22 @@ +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/train.py b/src/grass_mil/train.py similarity index 95% rename from src/train.py rename to src/grass_mil/train.py index 728616e..1f1eaf7 100644 --- a/src/train.py +++ b/src/grass_mil/train.py @@ -2,14 +2,15 @@ import hydra import lightning as L -import rootutils from lightning import Callback, LightningDataModule, LightningModule, Trainer from lightning.pytorch.loggers import Logger from omegaconf import DictConfig, OmegaConf -rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True) +from grass_mil._paths import setup_project_root -from src.utils import ( +setup_project_root() + +from grass_mil.utils import ( RankedLogger, extras, get_metric_value, @@ -134,7 +135,7 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: return metric_dict, object_dict -@hydra.main(version_base="1.3", config_path="../configs", config_name="train.yaml") +@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) diff --git a/src/grass_mil/utils/__init__.py b/src/grass_mil/utils/__init__.py new file mode 100644 index 0000000..78c7e75 --- /dev/null +++ b/src/grass_mil/utils/__init__.py @@ -0,0 +1,5 @@ +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 diff --git a/src/utils/instantiators.py b/src/grass_mil/utils/instantiators.py similarity index 97% rename from src/utils/instantiators.py rename to src/grass_mil/utils/instantiators.py index 82b9278..b823cfb 100644 --- a/src/utils/instantiators.py +++ b/src/grass_mil/utils/instantiators.py @@ -5,7 +5,7 @@ from lightning.pytorch.loggers import Logger from omegaconf import DictConfig -from src.utils import pylogger +from grass_mil.utils import pylogger log = pylogger.RankedLogger(__name__, rank_zero_only=True) diff --git a/src/utils/logging_utils.py b/src/grass_mil/utils/logging_utils.py similarity index 97% rename from src/utils/logging_utils.py rename to src/grass_mil/utils/logging_utils.py index 360abcd..01d6cdd 100644 --- a/src/utils/logging_utils.py +++ b/src/grass_mil/utils/logging_utils.py @@ -3,7 +3,7 @@ from lightning_utilities.core.rank_zero import rank_zero_only from omegaconf import OmegaConf -from src.utils import pylogger +from grass_mil.utils import pylogger log = pylogger.RankedLogger(__name__, rank_zero_only=True) diff --git a/src/utils/pylogger.py b/src/grass_mil/utils/pylogger.py similarity index 100% rename from src/utils/pylogger.py rename to src/grass_mil/utils/pylogger.py diff --git a/src/utils/rich_utils.py b/src/grass_mil/utils/rich_utils.py similarity index 98% rename from src/utils/rich_utils.py rename to src/grass_mil/utils/rich_utils.py index aeec680..ab2d158 100644 --- a/src/utils/rich_utils.py +++ b/src/grass_mil/utils/rich_utils.py @@ -9,7 +9,7 @@ from omegaconf import DictConfig, OmegaConf, open_dict from rich.prompt import Prompt -from src.utils import pylogger +from grass_mil.utils import pylogger log = pylogger.RankedLogger(__name__, rank_zero_only=True) diff --git a/src/utils/utils.py b/src/grass_mil/utils/utils.py similarity index 98% rename from src/utils/utils.py rename to src/grass_mil/utils/utils.py index 02b5576..5dbdbc1 100644 --- a/src/utils/utils.py +++ b/src/grass_mil/utils/utils.py @@ -4,7 +4,7 @@ from omegaconf import DictConfig -from src.utils import pylogger, rich_utils +from grass_mil.utils import pylogger, rich_utils log = pylogger.RankedLogger(__name__, rank_zero_only=True) diff --git a/src/inference/__init__.py b/src/inference/__init__.py deleted file mode 100644 index 4706399..0000000 --- a/src/inference/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -from src.inference.aggregation import aggregate_group_logits -from src.inference.collectors import collect_inference_payload -from src.inference.embeddings import embeddings_to_dataframe -from src.inference.metrics import compute_task_metrics - -__all__ = [ - "aggregate_group_logits", - "collect_inference_payload", - "compute_task_metrics", - "embeddings_to_dataframe", -] diff --git a/src/interpretability/__init__.py b/src/interpretability/__init__.py deleted file mode 100644 index a736047..0000000 --- a/src/interpretability/__init__.py +++ /dev/null @@ -1,43 +0,0 @@ -from src.interpretability.core.biomarkers import ( - cluster_attention_summary, - cluster_biomarker_summary, - cluster_survival_attention_summary, -) -from src.interpretability.core.clustering import run_clustering -from src.interpretability.core.reduction import run_reduction -from src.interpretability.core.transfer import ( - apply_cluster_transfer, - fit_cluster_transfer_from_report_bundle, - load_cluster_transfer_bundle, - save_cluster_transfer_bundle, -) -from src.interpretability.pipeline import run_interpretability_pipeline -from src.interpretability.reporting.render import render_interpretability_report -from src.interpretability.tier2.filtration import compute_filtration_curves -from src.interpretability.tier2.neighborhood import ( - run_diff_neighborhood_enrichment, - run_neighborhood_enrichment, -) -from src.interpretability.tier2.tissue_graph import ( - build_tissue_graph_figure, - prepare_tissue_graph_view, -) - -__all__ = [ - "run_reduction", - "run_clustering", - "cluster_biomarker_summary", - "cluster_attention_summary", - "cluster_survival_attention_summary", - "run_neighborhood_enrichment", - "run_diff_neighborhood_enrichment", - "compute_filtration_curves", - "prepare_tissue_graph_view", - "build_tissue_graph_figure", - "fit_cluster_transfer_from_report_bundle", - "apply_cluster_transfer", - "save_cluster_transfer_bundle", - "load_cluster_transfer_bundle", - "run_interpretability_pipeline", - "render_interpretability_report", -] diff --git a/src/interpretability/core/__init__.py b/src/interpretability/core/__init__.py deleted file mode 100644 index 8e92430..0000000 --- a/src/interpretability/core/__init__.py +++ /dev/null @@ -1,28 +0,0 @@ -from src.interpretability.core.biomarkers import ( - cluster_attention_summary, - cluster_biomarker_summary, - cluster_survival_attention_summary, -) -from src.interpretability.core.clustering import run_clustering -from src.interpretability.core.data import extract_embedding_set, load_interpretability_dataset -from src.interpretability.core.reduction import run_reduction -from src.interpretability.core.transfer import ( - apply_cluster_transfer, - fit_cluster_transfer_from_report_bundle, - load_cluster_transfer_bundle, - save_cluster_transfer_bundle, -) - -__all__ = [ - "run_reduction", - "run_clustering", - "cluster_biomarker_summary", - "cluster_attention_summary", - "cluster_survival_attention_summary", - "load_interpretability_dataset", - "extract_embedding_set", - "fit_cluster_transfer_from_report_bundle", - "apply_cluster_transfer", - "save_cluster_transfer_bundle", - "load_cluster_transfer_bundle", -] diff --git a/src/interpretability/core/biomarkers.py b/src/interpretability/core/biomarkers.py deleted file mode 100644 index daede4a..0000000 --- a/src/interpretability/core/biomarkers.py +++ /dev/null @@ -1,149 +0,0 @@ -from __future__ import annotations - -import numpy as np -import pandas as pd - -from src.interpretability.contracts import ClusterSummary - - -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 cluster_biomarker_summary( - table: pd.DataFrame, - cluster_labels: np.ndarray, - *, - composition_prefix: str = "comp_", - variance_estimator: str = "unbiased", -) -> ClusterSummary: - if len(table) != int(cluster_labels.shape[0]): - raise ValueError("table rows and cluster_labels length mismatch.") - - comp = _infer_composition_matrix(table, composition_prefix=composition_prefix) - labels = pd.Series(cluster_labels, name="cluster_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 ClusterSummary( - cluster_labels=np.asarray(cluster_labels), - composition=composition, - enrichment=enrichment, - cluster_counts=counts, - ) - - -def cluster_attention_summary( - table: pd.DataFrame, - cluster_labels: np.ndarray, - *, - attention_column: str = "attention", - score_column: str = "score", - bag_id_column: str = "bag_id", - composition_prefix: str = "comp_", - variance_estimator: str = "unbiased", - eps: float = 1e-8, -) -> ClusterSummary: - summary = cluster_biomarker_summary( - table, - cluster_labels, - composition_prefix=composition_prefix, - variance_estimator=variance_estimator, - ) - frame = table.copy() - frame["cluster_label"] = cluster_labels - for col in (attention_column, score_column, bag_id_column): - if col not in frame.columns: - raise ValueError(f"Missing required column {col!r} for attention summary.") - frame[attention_column] = frame[attention_column].astype(float) - frame[score_column] = frame[score_column].astype(float) - - weighted_frame = frame.assign(weighted=frame[attention_column] * frame[score_column]) - weighted = ( - weighted_frame.groupby("cluster_label") - .apply(lambda g: float(g["weighted"].sum() / max(g[attention_column].sum(), eps))) - .sort_index() - ) - mean_scores = frame.groupby("cluster_label")[score_column].mean().sort_index() - - cluster_attn_in_bag = ( - frame.groupby([bag_id_column, "cluster_label"])[attention_column] - .sum() - .unstack(fill_value=0.0) - ) - cluster_count_in_bag = ( - frame.groupby([bag_id_column, "cluster_label"]).size().unstack(fill_value=0) - ) - present_mask = cluster_count_in_bag > 0 - - attn_present = cluster_attn_in_bag.where(present_mask).mean(axis=0, skipna=True).sort_index() - - bag_totals = cluster_count_in_bag.sum(axis=1).replace(0, 1) - abundance = cluster_count_in_bag.div(bag_totals, axis=0) - lift = cluster_attn_in_bag / abundance.clip(lower=eps) - lift_present = lift.where(present_mask).mean(axis=0, skipna=True).sort_index() - - return ClusterSummary( - cluster_labels=summary.cluster_labels, - composition=summary.composition, - enrichment=summary.enrichment, - cluster_counts=summary.cluster_counts, - weighted_scores=weighted, - mean_scores=mean_scores, - attention_present=attn_present, - attention_lift_present=lift_present, - ) - - -def cluster_survival_attention_summary( - table: pd.DataFrame, - cluster_labels: np.ndarray, - *, - attention_column: str = "attention", - hazard_column: str = "hazard", - bag_id_column: str = "bag_id", - composition_prefix: str = "comp_", - variance_estimator: str = "unbiased", - eps: float = 1e-8, -) -> ClusterSummary: - renamed = table.copy() - if hazard_column not in renamed.columns: - raise ValueError(f"Missing hazard column {hazard_column!r}.") - renamed["score"] = renamed[hazard_column] - return cluster_attention_summary( - renamed, - cluster_labels, - attention_column=attention_column, - score_column="score", - bag_id_column=bag_id_column, - composition_prefix=composition_prefix, - variance_estimator=variance_estimator, - eps=eps, - ) diff --git a/src/interpretability/plugins/__init__.py b/src/interpretability/plugins/__init__.py deleted file mode 100644 index 2958f97..0000000 --- a/src/interpretability/plugins/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -from src.interpretability.plugins.base import InterpretabilityPlugin, PluginContext -from src.interpretability.plugins.builtin import create_builtin_registry, register_builtin_plugins -from src.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/interpretability/plugins/builtin.py b/src/interpretability/plugins/builtin.py deleted file mode 100644 index 2bff268..0000000 --- a/src/interpretability/plugins/builtin.py +++ /dev/null @@ -1,368 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Dict, List - -import numpy as np -import pandas as pd - -from src.interpretability.contracts import ( - ClusterSummary, - PluginResult, - ReportSection, -) -from src.interpretability.core.biomarkers import ( - cluster_attention_summary, - cluster_biomarker_summary, -) -from src.interpretability.plugins.base import InterpretabilityPlugin, PluginContext -from src.interpretability.plugins.registry import PluginRegistry, create_plugin_registry -from src.interpretability.tier2.filtration import compute_filtration_curves -from src.interpretability.tier2.neighborhood import ( - run_diff_neighborhood_enrichment, - run_neighborhood_enrichment, -) -from src.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, -) -> ClusterSummary | None: - summary = context.state.get("cluster_summary") - if not isinstance(summary, ClusterSummary): - return None - if summary.cluster_labels.shape != labels.shape: - return None - if not np.array_equal(summary.cluster_labels, labels): - return None - return summary - - -@dataclass -class ClusterProfilesPlugin(InterpretabilityPlugin): - name: str = "cluster_profiles" - - def required_inputs(self) -> List[str]: - return ["cluster_labels"] - - def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: - labels = np.asarray(context.state["cluster_labels"]) - summary = _context_cluster_summary(context, labels=labels) - if summary is None: - summary = cluster_biomarker_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.cluster_counts, - } - sections = [ - ReportSection( - title="Cluster Profiles", - description="Cluster-level composition and enrichment summary.", - tables={ - "composition": summary.composition, - "enrichment": summary.enrichment, - }, - ) - ] - return PluginResult(name=self.name, payload=payload, sections=sections) - - -@dataclass -class AttentionAttributionPlugin(InterpretabilityPlugin): - name: str = "attention_attribution" - - def required_inputs(self) -> List[str]: - return ["cluster_labels"] - - def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: - labels = np.asarray(context.state["cluster_labels"]) - attn_summary = cluster_attention_summary( - dataset.instance_table, - labels, - attention_column=str(params.get("attention_column", "attention")), - score_column=str(params.get("score_column", "score")), - bag_id_column=str(params.get("bag_id_column", dataset.bag_id_column)), - composition_prefix=str(params.get("composition_prefix", "comp_")), - variance_estimator=str(params.get("variance_estimator", "unbiased")), - ) - base_summary = _context_cluster_summary(context, labels=labels) - if base_summary is not None: - summary = ClusterSummary( - cluster_labels=base_summary.cluster_labels, - composition=base_summary.composition, - enrichment=base_summary.enrichment, - cluster_counts=base_summary.cluster_counts, - weighted_scores=attn_summary.weighted_scores, - mean_scores=attn_summary.mean_scores, - attention_present=attn_summary.attention_present, - attention_lift_present=attn_summary.attention_lift_present, - ) - else: - summary = attn_summary - payload = { - "weighted_scores": summary.weighted_scores, - "mean_scores": summary.mean_scores, - "attention_present": summary.attention_present, - "attention_lift_present": summary.attention_lift_present, - } - section_tables = {} - for key, value in payload.items(): - if value is not None: - section_tables[key] = value.to_frame(name=key) - sections = [ - ReportSection( - title="Attention Attribution", - description="Attention-weighted and abundance-corrected cluster scores.", - tables=section_tables, - ) - ] - return PluginResult(name=self.name, payload=payload, sections=sections) - - -@dataclass -class NeighborhoodEnrichmentPlugin(InterpretabilityPlugin): - name: str = "neighborhood_enrichment" - - def required_inputs(self) -> List[str]: - return ["spatial_table", "cluster_labels"] - - def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: - if dataset.spatial_table is None: - raise ValueError("Spatial table is required for neighborhood enrichment plugin.") - table = dataset.instance_table.copy() - table["cluster_label"] = np.asarray(context.state["cluster_labels"]) - result = run_neighborhood_enrichment( - table, - dataset.spatial_table, - label_column=str(params.get("label_column", "cluster_label")), - id_column=str(params.get("id_column", dataset.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")), - ) - payload = { - "enrichment": result.enrichment, - "observed": result.observed, - "expected": result.expected, - "pvalues": result.pvalues, - } - section_tables = { - "enrichment": result.enrichment, - "observed": result.observed, - "expected": result.expected, - } - if result.pvalues is not None: - section_tables["pvalues"] = result.pvalues - sections = [ - ReportSection( - title="Neighborhood Enrichment", - description="Cluster adjacency enrichment relative to expected connectivity.", - tables=section_tables, - ) - ] - return PluginResult(name=self.name, payload=payload, sections=sections) - - -@dataclass -class DiffNeighborhoodEnrichmentPlugin(InterpretabilityPlugin): - name: str = "diff_neighborhood_enrichment" - - def required_inputs(self) -> List[str]: - return ["spatial_table", "cluster_labels"] - - def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: - if dataset.spatial_table is None: - raise ValueError("Spatial table is required for differential neighborhood enrichment.") - table = dataset.instance_table.copy() - table["cluster_label"] = np.asarray(context.state["cluster_labels"]) - condition_column = str( - params.get("condition_column", dataset.condition_column or "condition") - ) - pairwise = run_diff_neighborhood_enrichment( - table, - dataset.spatial_table, - label_column=str(params.get("label_column", "cluster_label")), - condition_column=condition_column, - permutation_group_column=str(params.get("permutation_group_column", "sample_id")), - id_column=str(params.get("id_column", dataset.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")), - ) - payload = { - "enrichment_by_pair": {pair: out.enrichment for pair, out in pairwise.items()}, - "observed_by_pair": {pair: out.observed for pair, out in pairwise.items()}, - "expected_by_pair": {pair: out.expected for pair, out in pairwise.items()}, - "pvalues_by_pair": { - pair: out.pvalues for pair, out in pairwise.items() if out.pvalues is not None - }, - } - sections: List[ReportSection] = [] - for pair, out in pairwise.items(): - section_tables = { - "enrichment": out.enrichment, - "observed": out.observed, - "expected": out.expected, - } - if out.pvalues is not None: - section_tables["pvalues"] = out.pvalues - sections.append( - ReportSection( - title=f"Differential Neighborhood Enrichment ({pair})", - description=( - "Pairwise condition differential neighborhood enrichment " - "(left condition minus right condition)." - ), - tables=section_tables, - ) - ) - return PluginResult(name=self.name, payload=payload, sections=sections) - - -@dataclass -class FiltrationCurvesPlugin(InterpretabilityPlugin): - name: str = "filtration_curves" - - def required_inputs(self) -> List[str]: - return ["spatial_table", "cluster_labels"] - - def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: - if dataset.spatial_table is None: - raise ValueError("Spatial table is required for filtration curves plugin.") - table = dataset.instance_table.copy() - table["cluster_label"] = np.asarray(context.state["cluster_labels"]) - thresholds = _resolve_filtration_thresholds(dict(params)) - result = compute_filtration_curves( - table, - dataset.spatial_table, - thresholds=thresholds, - cluster_column=str(params.get("cluster_column", "cluster_label")), - cell_type_column=str(params.get("cell_type_column", "cell_type")), - id_column=str(params.get("id_column", dataset.id_column)), - distance_column=str(params.get("distance_column", "distance")), - scale_within_cluster=bool(params.get("scale_within_cluster", True)), - ) - payload: Dict[str, Any] = {"thresholds": result.thresholds, "curves": result.curves} - sections = [ - ReportSection( - title="Filtration Curves", - description="Distance-threshold accumulation by cluster and cell type.", - metadata={"thresholds": result.thresholds.tolist()}, - ) - ] - 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", "cluster_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["cluster_label"] = np.asarray(context.state["cluster_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", "cluster_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) - cluster_counts = meta.get("cluster_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"), - "cluster_count": len(cluster_counts) - if isinstance(cluster_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 cluster-colored cells.", - tables={"summary": section_summary}, - metadata=meta, - ) - ] - return PluginResult(name=self.name, payload=payload, sections=sections) - - -def register_builtin_plugins(registry: PluginRegistry) -> None: - registry.register(ClusterProfilesPlugin()) - registry.register(AttentionAttributionPlugin()) - registry.register(NeighborhoodEnrichmentPlugin()) - registry.register(DiffNeighborhoodEnrichmentPlugin()) - registry.register(FiltrationCurvesPlugin()) - registry.register(TissueGraphPlugin()) - - -def create_builtin_registry() -> PluginRegistry: - registry = create_plugin_registry() - register_builtin_plugins(registry) - return registry diff --git a/src/interpretability/reporting/__init__.py b/src/interpretability/reporting/__init__.py deleted file mode 100644 index 3174914..0000000 --- a/src/interpretability/reporting/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from src.interpretability.reporting.render import render_interpretability_report - -__all__ = ["render_interpretability_report"] diff --git a/src/interpretability/reporting/plotly_builders.py b/src/interpretability/reporting/plotly_builders.py deleted file mode 100644 index fc82448..0000000 --- a/src/interpretability/reporting/plotly_builders.py +++ /dev/null @@ -1,188 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, Dict, List - -import numpy as np -import pandas as pd - -from src.interpretability.contracts import ReportBundle -from src.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_cluster_enrichment_heatmap(bundle: ReportBundle) -> "go.Figure": - go = _require_plotly() - if bundle.cluster_summary is None: - raise ValueError("Cluster summary is missing.") - enr = bundle.cluster_summary.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="Cluster Enrichment Heatmap", - xaxis_title="Features", - yaxis_title="Cluster", - 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 cluster, 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"{cluster}:{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 "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 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.cluster_summary is not None: - figures.append(("cluster_enrichment", build_cluster_enrichment_heatmap(bundle))) - for name, fig in build_plugin_figures(bundle).items(): - figures.append((name, fig)) - return figures diff --git a/src/models/components/pooling.py b/src/models/components/pooling.py deleted file mode 100644 index 1fb8769..0000000 --- a/src/models/components/pooling.py +++ /dev/null @@ -1,45 +0,0 @@ -from __future__ import annotations - -from typing import Optional - -from torch import nn - -try: - from torch_geometric.nn import ( - GlobalAttention, - Set2Set, - global_add_pool, - global_max_pool, - global_mean_pool, - ) -except Exception as exc: # pragma: no cover - raise ImportError("torch_geometric is required for src.models.components.pooling") from exc - - -class GraphPooling(nn.Module): - """Thin pooling wrapper with explicit output_dim contract.""" - - def __init__(self, name: Optional[str], input_dim: int, set2set_steps: int = 3): - super().__init__() - self.name = (name or "mean").lower() - self.input_dim = input_dim - if self.name == "sum": - self.pool = global_add_pool - self.output_dim = input_dim - elif self.name == "mean": - self.pool = global_mean_pool - self.output_dim = input_dim - elif self.name == "max": - self.pool = global_max_pool - self.output_dim = input_dim - elif self.name == "attention": - self.pool = GlobalAttention(gate_nn=nn.Linear(input_dim, 1)) - self.output_dim = input_dim - elif self.name == "set2set": - self.pool = Set2Set(input_dim, processing_steps=set2set_steps) - self.output_dim = input_dim * 2 - else: - raise ValueError(f"Unsupported graph pooling: {name}") - - def forward(self, x, batch): - return self.pool(x, batch) diff --git a/src/models/training/builders.py b/src/models/training/builders.py deleted file mode 100644 index 1ed51ba..0000000 --- a/src/models/training/builders.py +++ /dev/null @@ -1,101 +0,0 @@ -from __future__ import annotations - -from typing import Any, Dict, Optional - -import lightning as L -import torch - -from src.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 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}'." - ) - - -def _resolve_encoder_cfg(encoder_cfg: Dict[str, Any], inferred_input_dim: int) -> Dict[str, Any]: - cfg = dict(encoder_cfg) - if cfg.get("input_dim", 0) in (None, 0): - if inferred_input_dim <= 0: - raise ValueError( - "Could not infer encoder input_dim from datamodule; " - "set model.encoder.input_dim explicitly." - ) - cfg["input_dim"] = inferred_input_dim - return cfg - - -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) - 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") - 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/models/training/loss_utils.py b/src/models/training/loss_utils.py deleted file mode 100644 index 12e8965..0000000 --- a/src/models/training/loss_utils.py +++ /dev/null @@ -1,109 +0,0 @@ -from __future__ import annotations - -from typing import Any, Dict, List, Optional, Sequence - -import torch -import torch.nn.functional as F - - -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 build_mil_aux_targets( - *, - bag_targets: torch.Tensor, - bag_indices: List[List[int]], - bag_ids: Sequence[str], - bag_attention: Dict[str, torch.Tensor], - target_mode: str, -) -> tuple[torch.Tensor, torch.Tensor]: - aux_targets: List[torch.Tensor] = [] - aux_weights: List[torch.Tensor] = [] - for bag_pos, (bag_id, idxs) in enumerate(zip(bag_ids, bag_indices)): - bag_target = bag_targets[bag_pos : bag_pos + 1] - repeated_target = bag_target.repeat(len(idxs), 1) - if target_mode == "graph_label": - aux_targets.append(repeated_target) - aux_weights.append(torch.ones_like(repeated_target)) - continue - if target_mode == "attention_shaped_ti": - attn = bag_attention[bag_id].detach().reshape(-1, 1).to(repeated_target.device) - target = 0.5 + (repeated_target - 0.5) * attn - aux_targets.append(target) - aux_weights.append(attn) - continue - raise ValueError( - "node_aux.target_mode must be one of ['graph_label', 'attention_shaped_ti']" - ) - return torch.cat(aux_targets, dim=0), torch.cat(aux_weights, dim=0) - - -def gather_instance_logits( - patch_logits: torch.Tensor, bag_indices: List[List[int]] -) -> torch.Tensor: - return torch.cat([patch_logits[idxs] for idxs in bag_indices], dim=0) - - -def compute_aux_node_loss( - *, - aux_logits: torch.Tensor, - aux_targets: torch.Tensor, - aux_weights: Optional[torch.Tensor], - loss_mode: str, -) -> torch.Tensor: - if loss_mode not in {"bce", "weighted_bce"}: - raise ValueError("node_aux.loss_mode must be one of ['bce', 'weighted_bce']") - loss = F.binary_cross_entropy_with_logits(aux_logits, aux_targets.float(), reduction="none") - if loss_mode == "weighted_bce": - if aux_weights is None: - raise ValueError("weighted_bce requires non-null aux_weights.") - loss = loss * aux_weights - return loss.mean() - - -def compute_entropy_regularization(values: torch.Tensor, eps: float = 1.0e-8) -> torch.Tensor: - probs = values.float().clamp(min=eps, max=1.0) - return -(probs * torch.log(probs + eps)).mean() - - -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/models/training/ssl_runtime.py b/src/models/training/ssl_runtime.py deleted file mode 100644 index 72ad460..0000000 --- a/src/models/training/ssl_runtime.py +++ /dev/null @@ -1,43 +0,0 @@ -from __future__ import annotations - -import math -from dataclasses import dataclass - -import torch - -try: - from torch_geometric.data import Data - from torch_geometric.utils import dropout_edge -except Exception as exc: # pragma: no cover - raise ImportError("torch_geometric is required for ssl runtime helpers") from exc - - -@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 - - -def augment_graph(batch: Data, *, drop_edge_p: float, drop_feat_p: float) -> Data: - aug = batch.clone() - if drop_feat_p > 0 and hasattr(aug, "x") and aug.x.numel() > 0: - feat_mask = torch.rand_like(aug.x) > drop_feat_p - aug.x = aug.x * feat_mask.float() - if drop_edge_p > 0 and hasattr(aug, "edge_index"): - edge_index, edge_mask = dropout_edge(aug.edge_index, p=drop_edge_p) - aug.edge_index = edge_index - if hasattr(aug, "edge_attr") and aug.edge_attr is not None: - aug.edge_attr = aug.edge_attr[edge_mask] - return aug diff --git a/src/utils/__init__.py b/src/utils/__init__.py deleted file mode 100644 index 5b0707c..0000000 --- a/src/utils/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from src.utils.instantiators import instantiate_callbacks, instantiate_loggers -from src.utils.logging_utils import log_hyperparameters -from src.utils.pylogger import RankedLogger -from src.utils.rich_utils import enforce_tags, print_config_tree -from src.utils.utils import extras, get_metric_value, task_wrapper diff --git a/tests/configs/test_component_instantiation.py b/tests/configs/test_component_instantiation.py index 8030315..d8682b6 100644 --- a/tests/configs/test_component_instantiation.py +++ b/tests/configs/test_component_instantiation.py @@ -1,5 +1,7 @@ from pathlib import Path +from tests.helpers.config_paths import CONFIGS_DIR + import hydra from omegaconf import OmegaConf @@ -10,7 +12,7 @@ def _instantiate_yaml(path: Path): def test_model_component_yaml_instantiation(): - base = Path("configs/model") + base = CONFIGS_DIR / "model" targets = [ base / "encoder" / "gin.yaml", base / "encoder" / "gcn.yaml", @@ -20,7 +22,6 @@ def test_model_component_yaml_instantiation(): base / "attention" / "gated.yaml", base / "attention" / "gated_projected.yaml", base / "heads" / "graph.yaml", - base / "heads" / "node.yaml", base / "loss" / "categorical_ce.yaml", base / "loss" / "categorical_bce.yaml", base / "loss" / "regression_mse.yaml", @@ -33,7 +34,7 @@ def test_model_component_yaml_instantiation(): def test_training_module_yaml_instantiation(): - base = Path("configs/model") + base = CONFIGS_DIR / "model" cases = [ ("bgrl_module.yaml", "pretrain_bgrl.yaml", "cosine_step.yaml"), ("supervised_module.yaml", "finetune_mean.yaml", "cosine_epoch.yaml"), @@ -42,17 +43,17 @@ def test_training_module_yaml_instantiation(): for model_name, task_name, scheduler_name in cases: path = base / model_name cfg = OmegaConf.load(path) - cfg.optim = OmegaConf.load(Path("configs/optim/adamw.yaml")) - cfg.scheduler = OmegaConf.load(Path("configs/scheduler") / scheduler_name) - cfg.task = OmegaConf.load(Path("configs/task") / task_name).task + cfg.optim = OmegaConf.load(CONFIGS_DIR / "optim" / "adamw.yaml") + cfg.scheduler = OmegaConf.load(CONFIGS_DIR / "scheduler" / scheduler_name) + cfg.task = OmegaConf.load(CONFIGS_DIR / "task" / task_name).task obj = hydra.utils.instantiate(cfg) assert obj is not None def test_sampler_strategy_resolution(): - from src.data.components.samplers import get_sampler_strategy + from grass_mil.data.components.samplers import get_sampler_strategy - base = Path("configs/data/sampler") + 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)) @@ -60,10 +61,10 @@ def test_sampler_strategy_resolution(): def test_default_component_and_data_sampler_configs(): - model_cfg = OmegaConf.load(Path("configs/model/components.yaml")) + model_cfg = OmegaConf.load(CONFIGS_DIR / "model" / "components.yaml") assert model_cfg.model.attention._target_ == ( - "src.models.components.attention.AttnNetGatedProjected" + "grass_mil.models.components.attention.AttnNetGatedProjected" ) - data_cfg = OmegaConf.load(Path("configs/data/spatial_omics.yaml")) + data_cfg = OmegaConf.load(CONFIGS_DIR / "data" / "spatial_omics.yaml") assert data_cfg.sampler.name == "shadow_custom" diff --git a/tests/conftest.py b/tests/conftest.py index 17a288d..634d545 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,7 +4,7 @@ import pytest import rootutils -from hydra import compose, initialize +from hydra import compose, initialize_config_module from hydra.core.global_hydra import GlobalHydra from omegaconf import DictConfig, OmegaConf, open_dict @@ -15,7 +15,7 @@ def cfg_train_global() -> DictConfig: :return: A DictConfig object containing a default Hydra configuration for training. """ - with initialize(version_base="1.3", config_path="../configs"): + 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 @@ -51,7 +51,7 @@ def cfg_eval_global() -> DictConfig: :return: A DictConfig containing a default Hydra configuration for evaluation. """ - with initialize(version_base="1.3", config_path="../configs"): + 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 @@ -82,7 +82,7 @@ def cfg_eval_global() -> DictConfig: @pytest.fixture(scope="package") def cfg_predict_global() -> DictConfig: """A pytest fixture for inference prediction configuration.""" - with initialize(version_base="1.3", config_path="../configs"): + with initialize_config_module(version_base="1.3", config_module="grass_mil.configs"): cfg = compose( config_name="inference/predict.yaml", return_hydra_config=True, @@ -114,7 +114,7 @@ def cfg_predict_global() -> DictConfig: @pytest.fixture(scope="package") def cfg_interpret_global() -> DictConfig: - with initialize(version_base="1.3", config_path="../configs"): + with initialize_config_module(version_base="1.3", config_module="grass_mil.configs"): cfg = compose( config_name="interpretability/report.yaml", return_hydra_config=True, diff --git a/tests/data/test_datamodule_sampler_exposure.py b/tests/data/test_datamodule_sampler_exposure.py index e156fb3..d642eca 100644 --- a/tests/data/test_datamodule_sampler_exposure.py +++ b/tests/data/test_datamodule_sampler_exposure.py @@ -1,7 +1,7 @@ from pathlib import Path -from src.data.components.samplers import SamplerConfig -from src.data.spatial_omics_datamodule import SpatialOmicsDataModule +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: diff --git a/tests/data/test_loader_flags.py b/tests/data/test_loader_flags.py index 0bea9ab..b2f6262 100644 --- a/tests/data/test_loader_flags.py +++ b/tests/data/test_loader_flags.py @@ -5,11 +5,11 @@ pytest.importorskip("lightning") -from src.data.components.feature_reducers import FeatureReducerConfig -from src.data.components.graph_builders import GraphBuilderConfig -from src.data.components.loaders import CsvConfig, H5adConfig, SceConfig -from src.data.components.patching import TileConfig -from src.data.components.precompute import ( +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, diff --git a/tests/data/test_precompute_manifest_regions.py b/tests/data/test_precompute_manifest_regions.py index ee68317..2c426d7 100644 --- a/tests/data/test_precompute_manifest_regions.py +++ b/tests/data/test_precompute_manifest_regions.py @@ -5,7 +5,7 @@ pytest.importorskip("lightning") -from src.data.components.precompute import ManifestConfig, SpatialOmicsPreprocessor +from grass_mil.data.components.precompute import ManifestConfig, SpatialOmicsPreprocessor def _validator() -> SimpleNamespace: diff --git a/tests/data/test_precompute_reducer_modes.py b/tests/data/test_precompute_reducer_modes.py index fbb73aa..443ee77 100644 --- a/tests/data/test_precompute_reducer_modes.py +++ b/tests/data/test_precompute_reducer_modes.py @@ -3,11 +3,11 @@ pytest.importorskip("lightning") -from src.data.components.feature_reducers import FeatureReducerConfig -from src.data.components.graph_builders import GraphBuilderConfig -from src.data.components.loaders import CsvConfig, H5adConfig, SceConfig -from src.data.components.patching import TileConfig -from src.data.components.precompute import ( +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, @@ -15,7 +15,7 @@ PreparedSample, SpatialOmicsPreprocessor, ) -from src.data.components.spatial_types import SpatialOmicsTable +from grass_mil.data.components.spatial_types import SpatialOmicsTable def _make_preprocessor(fit_mode: str) -> SpatialOmicsPreprocessor: diff --git a/tests/data/test_samplers.py b/tests/data/test_samplers.py index 845b4ce..45f370a 100644 --- a/tests/data/test_samplers.py +++ b/tests/data/test_samplers.py @@ -42,13 +42,13 @@ def _toy_dataset_with_graph_attrs(n_graphs: int = 2, n_nodes: int = 20): def _attach_cell_type_metadata(data, labels): - data.categorical_index = torch.tensor(labels, dtype=torch.long).view(-1, 1) + 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 src.data.components.samplers import _select_sequence_value_for_subgraph + 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"]] @@ -62,7 +62,7 @@ def test_sequence_assignment_preserves_schema_name_lists_without_nesting(): def test_sequence_assignment_does_not_slice_flat_schema_name_lists_by_length(): - from src.data.components.samplers import _select_sequence_value_for_subgraph + 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"] @@ -76,7 +76,7 @@ def test_sequence_assignment_does_not_slice_flat_schema_name_lists_by_length(): def test_sequence_assignment_still_slices_per_subgraph_non_schema_lists(): - from src.data.components.samplers import _select_sequence_value_for_subgraph + from grass_mil.data.components.samplers import _select_sequence_value_for_subgraph value = ["sample_0", "sample_1"] out = _select_sequence_value_for_subgraph( @@ -90,7 +90,7 @@ def test_sequence_assignment_still_slices_per_subgraph_non_schema_lists(): def test_identity_sampler_strategy(): pytest.importorskip("torch_geometric") - from src.data.components.samplers import get_sampler_strategy + from grass_mil.data.components.samplers import get_sampler_strategy dataset = _toy_dataset() strategy = get_sampler_strategy({"name": "identity", "kwargs": {}}) @@ -113,7 +113,7 @@ def test_shadow_native_requires_no_post_transform(): except Exception: pytest.skip("torch_sparse is not available") - from src.data.components.samplers import get_sampler_strategy + 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": {}}) @@ -134,7 +134,7 @@ def test_shadow_custom_unit_loader(): except Exception: pytest.skip("torch_sparse is not available") - from src.data.components.samplers import get_sampler_strategy + 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": {}}) @@ -152,8 +152,8 @@ def test_shadow_custom_unit_loader(): def test_shadow_custom_without_torch_sparse_rejects_transform(monkeypatch): pytest.importorskip("torch_geometric") - from src.data.components import samplers as samplers_mod - from src.data.components.samplers import ShadowCustomStrategy + 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) @@ -170,7 +170,7 @@ def test_shadow_custom_without_torch_sparse_rejects_transform(monkeypatch): def test_shadow_custom_runtime_dataset_loader_uses_unit_loader(monkeypatch): pytest.importorskip("torch_geometric") - from src.data.components.samplers import get_sampler_strategy + from grass_mil.data.components.samplers import get_sampler_strategy dataset = _toy_dataset_with_graph_attrs(n_graphs=2) strategy = get_sampler_strategy( @@ -212,7 +212,7 @@ def test_shadow_runtime_normalizes_graph_level_metadata(monkeypatch): pytest.importorskip("torch_geometric") from torch_geometric.data import Batch - from src.data.components.samplers import get_sampler_strategy + from grass_mil.data.components.samplers import get_sampler_strategy dataset = _toy_dataset_with_graph_attrs(n_graphs=1) strategy = get_sampler_strategy( @@ -251,7 +251,7 @@ def _fake_unit_loader(data, **kwargs): def test_runtime_loader_len_matches_yielded_batches(monkeypatch): pytest.importorskip("torch_geometric") - from src.data.components.samplers import get_sampler_strategy + 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( @@ -291,7 +291,7 @@ def _fake_unit_loader(data, **kwargs): def test_runtime_shadow_config_proportional_defaults(): - from src.data.components.samplers import RuntimeShadowConfig + from grass_mil.data.components.samplers import RuntimeShadowConfig cfg = RuntimeShadowConfig.from_dict({}) assert cfg.proportional_root_sampling is True @@ -302,7 +302,7 @@ def test_runtime_shadow_config_proportional_defaults(): def test_shadow_runtime_passes_weighted_node_idx(monkeypatch): pytest.importorskip("torch_geometric") - from src.data.components.samplers import get_sampler_strategy + 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]) @@ -346,7 +346,7 @@ def _fake_unit_loader(data, **kwargs): def test_shadow_runtime_falls_back_to_none_when_property_missing(monkeypatch): pytest.importorskip("torch_geometric") - from src.data.components.samplers import get_sampler_strategy + 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] @@ -393,7 +393,7 @@ def _fake_unit_loader(data, **kwargs): def test_shadow_runtime_can_disable_proportional_sampling(monkeypatch): pytest.importorskip("torch_geometric") - from src.data.components.samplers import get_sampler_strategy + 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]) @@ -436,14 +436,14 @@ def _fake_unit_loader(data, **kwargs): @pytest.mark.parametrize("weight_mode", ["inverse", "sqrt_inverse", "proportional"]) def test_shadow_runtime_accepts_weight_modes(weight_mode): - from src.data.components.samplers import RuntimeShadowConfig + 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 src.data.components.samplers import RuntimeShadowConfig + from grass_mil.data.components.samplers import RuntimeShadowConfig cfg = RuntimeShadowConfig.from_dict({}) assert cfg.subsample_fraction == 1.0 @@ -452,7 +452,7 @@ def test_runtime_shadow_config_accepts_subsample_defaults(): @pytest.mark.parametrize("fraction", [0.0, -0.1, 1.1]) def test_runtime_shadow_config_rejects_invalid_subsample_fraction(fraction): - from src.data.components.samplers import RuntimeShadowConfig + from grass_mil.data.components.samplers import RuntimeShadowConfig with pytest.raises(ValueError, match="subsample_fraction"): RuntimeShadowConfig.from_dict({"subsample_fraction": fraction}) @@ -460,7 +460,7 @@ def test_runtime_shadow_config_rejects_invalid_subsample_fraction(fraction): def test_shadow_runtime_subsamples_roots_with_seed(monkeypatch): pytest.importorskip("torch_geometric") - from src.data.components.samplers import get_sampler_strategy + 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] @@ -509,7 +509,7 @@ def _fake_unit_loader(data, **kwargs): def test_shadow_runtime_subsampled_len_matches_yielded_batches(monkeypatch): pytest.importorskip("torch_geometric") - from src.data.components.samplers import get_sampler_strategy + 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( @@ -545,3 +545,178 @@ def _fake_unit_loader(data, **kwargs): ) 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 diff --git a/tests/data/test_seed_sampling.py b/tests/data/test_seed_sampling.py new file mode 100644 index 0000000..8ae76b0 --- /dev/null +++ b/tests/data/test_seed_sampling.py @@ -0,0 +1,223 @@ +"""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 index e0f46c4..366ed5d 100644 --- a/tests/eval/test_aggregation.py +++ b/tests/eval/test_aggregation.py @@ -3,8 +3,8 @@ import torch import pytest -from src.inference.aggregation import aggregate_group_logits -from src.inference.schemas import BatchPredictionPayload +from grass_mil.inference.aggregation import aggregate_group_logits +from grass_mil.inference.schemas import BatchPredictionPayload def test_aggregate_group_logits_mean() -> None: @@ -51,15 +51,16 @@ def test_aggregate_group_logits_mean_drops_non_aggregable_attention_rows() -> No 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=[torch.tensor([1.0]), torch.tensor([9.0]), torch.tensor([1.0])], + 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], [9.0], [1.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"], @@ -71,7 +72,61 @@ def test_aggregate_group_logits_attention_weighted_uses_instance_softmax() -> No assert out.bag_attention is not None assert len(out.bag_attention) == 2 assert out.bag_attention[0] is not None - assert torch.allclose(out.bag_attention[0], torch.tensor([0.000335, 0.999665]), atol=1e-4) + # 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: diff --git a/tests/eval/test_collectors.py b/tests/eval/test_collectors.py index 58bdf44..f1911b3 100644 --- a/tests/eval/test_collectors.py +++ b/tests/eval/test_collectors.py @@ -4,7 +4,7 @@ import torch from types import SimpleNamespace -from src.inference.collectors import collect_inference_payload +from grass_mil.inference.collectors import collect_inference_payload class _FakeTrainer: diff --git a/tests/eval/test_embeddings.py b/tests/eval/test_embeddings.py index 3fd5e76..0c4838d 100644 --- a/tests/eval/test_embeddings.py +++ b/tests/eval/test_embeddings.py @@ -2,8 +2,8 @@ import torch -from src.inference.embeddings import embeddings_to_dataframe, node_embeddings_to_dataframe -from src.inference.schemas import EmbeddingPayload +from grass_mil.inference.embeddings import embeddings_to_dataframe, node_embeddings_to_dataframe +from grass_mil.inference.schemas import EmbeddingPayload def test_embedding_payload_to_dataframes() -> None: diff --git a/tests/eval/test_interpretability_export.py b/tests/eval/test_interpretability_export.py index d95e30a..46f6ef2 100644 --- a/tests/eval/test_interpretability_export.py +++ b/tests/eval/test_interpretability_export.py @@ -6,8 +6,8 @@ import pytest import torch -from src.inference.interpretability_export import build_instance_table, build_spatial_table -from src.inference.schemas import BatchPredictionPayload +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: @@ -184,13 +184,19 @@ def test_build_spatial_table_requires_payload_connectivity() -> None: build_spatial_table(instance_frame, payload) -def test_build_instance_table_rejects_multi_column_attention_logits() -> None: +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.5]]), + 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"], @@ -199,5 +205,38 @@ def test_build_instance_table_rejects_multi_column_attention_logits() -> None: 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="exactly one column"): - build_instance_table(payload, require_composition=True) + 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) diff --git a/tests/eval/test_io.py b/tests/eval/test_io.py index 56838e9..e9b00bc 100644 --- a/tests/eval/test_io.py +++ b/tests/eval/test_io.py @@ -4,7 +4,7 @@ import math from pathlib import Path -from src.inference.io import write_json +from grass_mil.inference.io import write_json def test_write_json_sanitizes_non_finite_floats(tmp_path: Path) -> None: diff --git a/tests/eval/test_metrics.py b/tests/eval/test_metrics.py index 1be9fa2..f8981bd 100644 --- a/tests/eval/test_metrics.py +++ b/tests/eval/test_metrics.py @@ -1,11 +1,14 @@ from __future__ import annotations +import math + +import pytest import torch from omegaconf import OmegaConf from sklearn.metrics import roc_auc_score -from src.inference.metrics import compute_task_metrics -from src.inference.schemas import BatchPredictionPayload +from grass_mil.inference.metrics import compute_task_metrics +from grass_mil.inference.schemas import BatchPredictionPayload def test_compute_binary_metrics() -> None: @@ -122,3 +125,107 @@ def test_compute_binary_multitask_metrics_are_per_task_not_flattened() -> None: 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 index 97c53ae..05f9384 100644 --- a/tests/eval/test_predict_cli.py +++ b/tests/eval/test_predict_cli.py @@ -10,21 +10,20 @@ from omegaconf import DictConfig, open_dict from omegaconf import OmegaConf -from src.inference.schemas import ( +from grass_mil.inference.schemas import ( AggregatedPredictionPayload, BatchPredictionPayload, CollectedInferencePayload, EmbeddingPayload, ) -from src.inference.predict import ( +from grass_mil.inference.predict import ( _needs_instance_payload, _plan_preforward_subsampling, predict, ) -from src.train import train +from grass_mil.train import train -@pytest.mark.slow 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 @@ -186,7 +185,7 @@ def test_needs_instance_payload(cfg_dict, expected) -> None: def test_predict_raises_when_preforward_subsampling_is_requested_but_ineffective( monkeypatch, tmp_path: Path ) -> None: - import src.inference.predict as predict_module + import grass_mil.inference.predict as predict_module datamodule = _DataModule(sampler_name="identity", runtime={"enabled": True}) @@ -243,7 +242,7 @@ def _fake_instantiate(cfg, *args, **kwargs): def test_predict_disables_post_subsampling_only_when_preforward_is_effective( monkeypatch, tmp_path: Path ) -> None: - import src.inference.predict as predict_module + import grass_mil.inference.predict as predict_module datamodule = _DataModule(sampler_name="shadow_native", runtime={"enabled": True}) captured = {} @@ -345,7 +344,7 @@ def _fake_aggregate(payload, **kwargs): def test_predict_writes_embeddings_from_single_pass_collected_payload( monkeypatch, tmp_path: Path ) -> None: - import src.inference.predict as predict_module + import grass_mil.inference.predict as predict_module datamodule = _DataModule(sampler_name="shadow_native", runtime={"enabled": True}) wrote_paths = [] @@ -465,7 +464,7 @@ def _fake_write_dataframe(frame, path): def test_predict_summary_includes_aggregation_metadata(monkeypatch, tmp_path: Path) -> None: - import src.inference.predict as predict_module + import grass_mil.inference.predict as predict_module class _FakeDataModule: pass @@ -611,7 +610,7 @@ def test_predict_raises_when_ckpt_path_missing(tmp_path: Path) -> None: def test_predict_can_export_interpretability_tables(monkeypatch, tmp_path: Path) -> None: - import src.inference.predict as predict_module + import grass_mil.inference.predict as predict_module class _FakeDataModule: class _Dataset: diff --git a/tests/helpers/config_paths.py b/tests/helpers/config_paths.py new file mode 100644 index 0000000..57f9b60 --- /dev/null +++ b/tests/helpers/config_paths.py @@ -0,0 +1,12 @@ +"""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/synthetic_datamodule.py b/tests/helpers/synthetic_datamodule.py index ce8b36d..6585981 100644 --- a/tests/helpers/synthetic_datamodule.py +++ b/tests/helpers/synthetic_datamodule.py @@ -16,19 +16,29 @@ def _make_graph( 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, ) - graph_y = torch.tensor([[label]], dtype=torch.float32) - graph_w = torch.ones_like(graph_y) + 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 = ["label"] + data.graph_label_names = ["time", "event"] if survival else ["label"] return data @@ -40,6 +50,7 @@ def __init__( pin_memory: bool = False, input_dim: int = 8, graphs_per_region: int = 20, + survival: bool = False, ) -> None: super().__init__() self.batch_size = batch_size @@ -47,6 +58,7 @@ def __init__( 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 @@ -54,9 +66,23 @@ def __init__( 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)) data.append( - _make_graph("region_1", 1.0, input_dim=self.input_dim, patch_idx=2 * i + 1) + _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 diff --git a/tests/interpretability/test_attribution.py b/tests/interpretability/test_attribution.py new file mode 100644 index 0000000..923cfe3 --- /dev/null +++ b/tests/interpretability/test_attribution.py @@ -0,0 +1,325 @@ +"""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_core.py b/tests/interpretability/test_core.py index ba65521..0eec984 100644 --- a/tests/interpretability/test_core.py +++ b/tests/interpretability/test_core.py @@ -4,12 +4,11 @@ import pandas as pd import pytest -from src.interpretability.core.biomarkers import ( - cluster_attention_summary, - cluster_biomarker_summary, +from grass_mil.interpretability.core.biomarkers import ( + niche_composition_summary, ) -from src.interpretability.core.clustering import run_clustering -from src.interpretability.core.reduction import run_reduction +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: @@ -37,19 +36,9 @@ def test_cluster_biomarker_and_attention_summary() -> None: } ) labels = np.array([0, 0, 1, 1, 1, 0]) - bio = cluster_biomarker_summary(table, labels) + bio = niche_composition_summary(table, labels) assert "comp_A" in bio.composition.columns - assert bio.cluster_counts.sum() == len(table) - - att = cluster_attention_summary( - table, - labels, - attention_column="attention", - score_column="score", - bag_id_column="bag_id", - ) - assert att.weighted_scores is not None - assert att.attention_lift_present is not None + assert bio.niche_counts.sum() == len(table) def test_cluster_biomarker_requires_composition_columns() -> None: @@ -61,7 +50,7 @@ def test_cluster_biomarker_requires_composition_columns() -> None: ) labels = np.array([0, 1]) with pytest.raises(ValueError, match="No composition columns found"): - cluster_biomarker_summary(table, labels) + niche_composition_summary(table, labels) def test_cluster_biomarker_variance_estimator_is_tunable() -> None: @@ -74,8 +63,8 @@ def test_cluster_biomarker_variance_estimator_is_tunable() -> None: ) labels = np.array([0, 1, 1]) - unbiased = cluster_biomarker_summary(table, labels, variance_estimator="unbiased") - biased = cluster_biomarker_summary(table, labels, variance_estimator="biased") + 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) @@ -92,7 +81,7 @@ def test_cluster_biomarker_variance_estimator_rejects_unknown_values() -> None: ) labels = np.array([0, 1]) with pytest.raises(ValueError, match="Unsupported variance_estimator"): - cluster_biomarker_summary(table, labels, variance_estimator="invalid_mode") + niche_composition_summary(table, labels, variance_estimator="invalid_mode") def test_hdbscan_clustering_forwards_selection_and_density_knobs(monkeypatch) -> None: diff --git a/tests/interpretability/test_niches.py b/tests/interpretability/test_niches.py new file mode 100644 index 0000000..9ef213e --- /dev/null +++ b/tests/interpretability/test_niches.py @@ -0,0 +1,236 @@ +"""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 index 08f3eae..f1149bc 100644 --- a/tests/interpretability/test_pipeline_and_report.py +++ b/tests/interpretability/test_pipeline_and_report.py @@ -8,15 +8,15 @@ from hydra.core.hydra_config import HydraConfig from omegaconf import DictConfig, open_dict -from src.interpretability.contracts import ClusteringResult -from src.interpretability.core.data import load_interpretability_dataset -from src.interpretability.pipeline import run_interpretability_pipeline -from src.interpretability.report_cli import run_report +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 src.interpretability.reporting.plotly_builders import bundle_figures -from src.interpretability.reporting.render import render_interpretability_report +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]: @@ -69,10 +69,7 @@ def test_pipeline_and_report_render(monkeypatch, tmp_path: Path) -> None: "params": {"n_clusters": 2, "linkage": "ward"}, }, "plugins": { - "enabled": ["cluster_profiles", "attention_attribution"], - "params": { - "attention_attribution": {"attention_column": "attention", "score_column": "score"} - }, + "enabled": ["niche_profiles"], }, } bundle = run_interpretability_pipeline(dataset, cfg, artifacts_dir=tmp_path / "artifacts") @@ -80,11 +77,11 @@ def test_pipeline_and_report_render(monkeypatch, tmp_path: Path) -> None: assert len(figs) >= 2 monkeypatch.setattr( - "src.interpretability.reporting.render.render_pdf_via_playwright", + "grass_mil.interpretability.reporting.render.render_pdf_via_playwright", lambda html_path, pdf_path: pdf_path, # type: ignore[lambda-assign] ) monkeypatch.setattr( - "src.interpretability.reporting.render.export_plotly_snapshots", + "grass_mil.interpretability.reporting.render.export_plotly_snapshots", lambda figures, output_dir, scale=2.0: {}, # type: ignore[lambda-assign] ) out = render_interpretability_report( @@ -116,7 +113,7 @@ def test_report_render_tolerates_snapshot_export_failure_when_pdf_disabled( "params": {"n_clusters": 2, "linkage": "ward"}, }, "plugins": { - "enabled": ["cluster_profiles"], + "enabled": ["niche_profiles"], "params": {}, }, } @@ -126,7 +123,7 @@ def _raise_snapshot_error(figures, output_dir, scale=2.0): # type: ignore[unuse raise RuntimeError("kaleido unavailable") monkeypatch.setattr( - "src.interpretability.reporting.render.export_plotly_snapshots", + "grass_mil.interpretability.reporting.render.export_plotly_snapshots", _raise_snapshot_error, ) out = render_interpretability_report( @@ -177,7 +174,7 @@ def test_pipeline_and_report_render_with_diff_neighborhood_plugin( assert "diff_neighborhood_enrichment_X_Y" in figs monkeypatch.setattr( - "src.interpretability.reporting.render.export_plotly_snapshots", + "grass_mil.interpretability.reporting.render.export_plotly_snapshots", lambda figures, output_dir, scale=2.0: {}, # type: ignore[lambda-assign] ) out = render_interpretability_report( @@ -222,7 +219,7 @@ def test_pipeline_and_report_render_with_tissue_graph_plugin(monkeypatch, tmp_pa assert "tissue_graph_sample_id_sx" in figs monkeypatch.setattr( - "src.interpretability.reporting.render.export_plotly_snapshots", + "grass_mil.interpretability.reporting.render.export_plotly_snapshots", lambda figures, output_dir, scale=2.0: {}, # type: ignore[lambda-assign] ) out = render_interpretability_report( @@ -248,9 +245,9 @@ def test_pipeline_fails_fast_when_plugin_required_inputs_missing_cluster_labels( cfg = { "reduction": {"enabled": True, "method": "pca", "params": {"n_components": 2}}, "clustering": {"enabled": False}, - "plugins": {"enabled": ["cluster_profiles"], "params": {}}, + "plugins": {"enabled": ["niche_profiles"], "params": {}}, } - with pytest.raises(ValueError, match="missing required inputs: cluster_labels"): + with pytest.raises(ValueError, match="missing required inputs: niche_labels"): run_interpretability_pipeline(dataset, cfg, artifacts_dir=tmp_path / "artifacts") @@ -299,7 +296,9 @@ def _capture_run_clustering(method, x, **params): # type: ignore[no-untyped-def fitted_object=None, ) - monkeypatch.setattr("src.interpretability.pipeline.run_clustering", _capture_run_clustering) + monkeypatch.setattr( + "grass_mil.interpretability.pipeline.run_clustering", _capture_run_clustering + ) cfg = { "reduction": {"enabled": True, "method": "pca", "params": {"n_components": 1}}, "clustering": { @@ -314,7 +313,7 @@ def _capture_run_clustering(method, x, **params): # type: ignore[no-untyped-def assert "x" in captured assert captured["x"].shape == expected_raw.shape assert np.allclose(captured["x"], expected_raw) - assert bundle.cluster_feature_reduction is None + assert bundle.niche_feature_reduction is None def test_pipeline_exposes_cluster_feature_reduction_when_cluster_on_pca_enabled( @@ -339,9 +338,9 @@ def test_pipeline_exposes_cluster_feature_reduction_when_cluster_on_pca_enabled( "plugins": {"enabled": [], "params": {}}, } bundle = run_interpretability_pipeline(dataset, cfg, artifacts_dir=tmp_path / "artifacts") - assert bundle.cluster_feature_reduction is not None - assert bundle.cluster_feature_reduction.method == "pca" - assert bundle.cluster_feature_reduction.embedding.shape[1] == 1 + 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( @@ -351,11 +350,11 @@ def test_report_cli_smoke( ) -> None: instance_path, spatial_path = _write_minimal_tables(tmp_path) monkeypatch.setattr( - "src.interpretability.reporting.render.render_pdf_via_playwright", + "grass_mil.interpretability.reporting.render.render_pdf_via_playwright", lambda html_path, pdf_path: pdf_path, # type: ignore[lambda-assign] ) monkeypatch.setattr( - "src.interpretability.reporting.render.export_plotly_snapshots", + "grass_mil.interpretability.reporting.render.export_plotly_snapshots", lambda figures, output_dir, scale=2.0: {}, # type: ignore[lambda-assign] ) with open_dict(cfg_interpret): @@ -365,8 +364,7 @@ def test_report_cli_smoke( cfg_interpret.reduction.method = "pca" cfg_interpret.reduction.params = {"n_components": 2, "random_state": 42} cfg_interpret.plugins.enabled = [ - "cluster_profiles", - "attention_attribution", + "niche_profiles", "neighborhood_enrichment", "diff_neighborhood_enrichment", "filtration_curves", @@ -383,7 +381,7 @@ def test_report_cli_smoke( def test_pdf_renderer_raises_when_unavailable(monkeypatch, tmp_path: Path) -> None: - from src.interpretability.reporting.pdf import render_pdf_via_playwright + from grass_mil.interpretability.reporting.pdf import render_pdf_via_playwright real_import = __import__ @@ -398,3 +396,46 @@ def _fake_import(name, *args, **kwargs): 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 index 09a3728..846e7b5 100644 --- a/tests/interpretability/test_plugins_and_tier2.py +++ b/tests/interpretability/test_plugins_and_tier2.py @@ -6,12 +6,12 @@ import pandas as pd import pytest -from src.interpretability.contracts import ClusterSummary, InterpretabilityDataset -from src.interpretability.plugins.base import PluginContext -from src.interpretability.plugins.builtin import create_builtin_registry -from src.interpretability.plugins.registry import create_plugin_registry -from src.interpretability.tier2.filtration import compute_filtration_curves -from src.interpretability.tier2.neighborhood import ( +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, @@ -55,13 +55,13 @@ def _sample_dataset() -> InterpretabilityDataset: def test_plugin_registry_and_builtin_execution() -> None: registry = create_builtin_registry() - assert "cluster_profiles" in registry.list() + 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("cluster_profiles") - out = plugin.run(ds, PluginContext(state={"cluster_labels": labels})) - assert out.name == "cluster_profiles" + plugin = registry.get("niche_profiles") + out = plugin.run(ds, PluginContext(state={"niche_labels": labels})) + assert out.name == "niche_profiles" assert "composition" in out.payload @@ -72,7 +72,7 @@ def test_builtin_tissue_graph_plugin_execution() -> None: plugin = registry.get("tissue_graph") out = plugin.run( ds, - PluginContext(state={"cluster_labels": labels}), + PluginContext(state={"niche_labels": labels}), sample_column="sample_id", sample_value="sx", id_column="instance_id", @@ -93,7 +93,7 @@ def test_builtin_tissue_graph_plugin_requires_sample_value() -> None: with pytest.raises(ValueError, match="requires a non-empty sample_value"): plugin.run( ds, - PluginContext(state={"cluster_labels": labels}), + PluginContext(state={"niche_labels": labels}), sample_column="sample_id", id_column="instance_id", ) @@ -114,7 +114,7 @@ def test_builtin_tissue_graph_plugin_fails_when_coords_missing() -> None: with pytest.raises(ValueError, match="missing required columns"): plugin.run( missing_coords_ds, - PluginContext(state={"cluster_labels": labels}), + PluginContext(state={"niche_labels": labels}), sample_column="sample_id", sample_value="sx", id_column="instance_id", @@ -128,7 +128,7 @@ def test_builtin_diff_neighborhood_plugin_execution() -> None: plugin = registry.get("diff_neighborhood_enrichment") out = plugin.run( ds, - PluginContext(state={"cluster_labels": labels}), + PluginContext(state={"niche_labels": labels}), condition_column="condition", permutation_group_column="sample_id", n_perms=0, @@ -147,62 +147,28 @@ 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 = ClusterSummary( - cluster_labels=labels.copy(), + 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]), - cluster_counts=pd.Series({0: 2, 1: 4}), + niche_counts=pd.Series({0: 2, 1: 4}), ) def _should_not_run(*args, **kwargs): # type: ignore[no-untyped-def] - raise AssertionError("cluster_biomarker_summary should not be recomputed") + raise AssertionError("niche_composition_summary should not be recomputed") monkeypatch.setattr( - "src.interpretability.plugins.builtin.cluster_biomarker_summary", _should_not_run + "grass_mil.interpretability.plugins.builtin.niche_composition_summary", _should_not_run ) - plugin = registry.get("cluster_profiles") + plugin = registry.get("niche_profiles") out = plugin.run( ds, - PluginContext(state={"cluster_labels": labels, "cluster_summary": summary}), + PluginContext(state={"niche_labels": labels, "niche_summary": summary}), ) assert out.payload["composition"].equals(summary.composition) assert out.payload["enrichment"].equals(summary.enrichment) -def test_attention_attribution_reuses_context_base_summary(monkeypatch) -> None: - registry = create_builtin_registry() - ds = _sample_dataset() - labels = np.array([0, 0, 1, 1, 1, 0]) - base_summary = ClusterSummary( - cluster_labels=labels.copy(), - composition=pd.DataFrame({"comp_A": [0.6, 0.4]}, index=[0, 1]), - enrichment=pd.DataFrame({"comp_A": [0.2, -0.2]}, index=[0, 1]), - cluster_counts=pd.Series({0: 3, 1: 3}), - ) - attn_summary = ClusterSummary( - cluster_labels=labels.copy(), - composition=pd.DataFrame({"comp_A": [0.9, 0.1]}, index=[0, 1]), - enrichment=pd.DataFrame({"comp_A": [1.2, -1.2]}, index=[0, 1]), - cluster_counts=pd.Series({0: 2, 1: 4}), - weighted_scores=pd.Series({0: 0.3, 1: 0.7}), - mean_scores=pd.Series({0: 0.2, 1: 0.8}), - attention_present=pd.Series({0: 0.4, 1: 0.6}), - attention_lift_present=pd.Series({0: 1.1, 1: 0.9}), - ) - - monkeypatch.setattr( - "src.interpretability.plugins.builtin.cluster_attention_summary", - lambda *args, **kwargs: attn_summary, # type: ignore[no-untyped-def] - ) - plugin = registry.get("attention_attribution") - out = plugin.run( - ds, - PluginContext(state={"cluster_labels": labels, "cluster_summary": base_summary}), - ) - assert out.payload["weighted_scores"].equals(attn_summary.weighted_scores) - assert out.payload["mean_scores"].equals(attn_summary.mean_scores) - - def test_plugin_registry_instances_are_isolated() -> None: class _DummyPlugin: name = "dummy" @@ -277,11 +243,11 @@ 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["cluster_label"] = labels + frame["niche_label"] = labels enr = run_neighborhood_enrichment( frame, ds.spatial_table, # type: ignore[arg-type] - label_column="cluster_label", + label_column="niche_label", id_column="instance_id", n_perms=4, random_state=2, @@ -291,7 +257,7 @@ def test_tier2_neighborhood_and_filtration() -> None: frame, ds.spatial_table, # type: ignore[arg-type] thresholds=np.array([0.2, 0.5, 1.0]), - cluster_column="cluster_label", + niche_column="niche_label", cell_type_column="cell_type", id_column="instance_id", distance_column="distance", @@ -303,7 +269,7 @@ def test_filtration_subgraph_centric_counts_both_edge_endpoints() -> None: node_table = pd.DataFrame( { "instance_id": ["n0", "n1", "n2"], - "cluster_label": [0, 0, 1], + "niche_label": [0, 0, 1], "cell_type": ["A", "B", "A"], } ) @@ -318,17 +284,17 @@ def test_filtration_subgraph_centric_counts_both_edge_endpoints() -> None: node_table, spatial_table, thresholds=np.array([0.15, 0.25]), - cluster_column="cluster_label", + 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 cluster 0: A=1 (n0), B=1 (n1). + # 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 cluster 0: A=3, B=2. + # 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])) @@ -337,7 +303,7 @@ def test_neighborhood_undirected_option_adds_reverse_direction() -> None: node_table = pd.DataFrame( { "instance_id": ["n0", "n1"], - "cluster_label": ["A", "B"], + "niche_label": ["A", "B"], } ) spatial_table = pd.DataFrame( @@ -350,7 +316,7 @@ def test_neighborhood_undirected_option_adds_reverse_direction() -> None: directed = run_neighborhood_enrichment( node_table, spatial_table, - label_column="cluster_label", + label_column="niche_label", id_column="instance_id", n_perms=0, undirected=False, @@ -359,7 +325,7 @@ def test_neighborhood_undirected_option_adds_reverse_direction() -> None: undirected = run_neighborhood_enrichment( node_table, spatial_table, - label_column="cluster_label", + label_column="niche_label", id_column="instance_id", n_perms=0, undirected=True, @@ -375,7 +341,7 @@ def test_neighborhood_undirected_deduplicates_already_bidirectional_edges() -> N node_table = pd.DataFrame( { "instance_id": ["n0", "n1"], - "cluster_label": ["A", "B"], + "niche_label": ["A", "B"], } ) spatial_table = pd.DataFrame( @@ -387,7 +353,7 @@ def test_neighborhood_undirected_deduplicates_already_bidirectional_edges() -> N out = run_neighborhood_enrichment( node_table, spatial_table, - label_column="cluster_label", + label_column="niche_label", id_column="instance_id", n_perms=0, undirected=True, @@ -401,7 +367,7 @@ def test_neighborhood_ignores_weight_semantics() -> None: node_table = pd.DataFrame( { "instance_id": ["n0", "n1", "n2"], - "cluster_label": ["A", "B", "B"], + "niche_label": ["A", "B", "B"], } ) spatial_table = pd.DataFrame( @@ -414,7 +380,7 @@ def test_neighborhood_ignores_weight_semantics() -> None: out = run_neighborhood_enrichment( node_table, spatial_table, - label_column="cluster_label", + label_column="niche_label", id_column="instance_id", n_perms=0, undirected=False, @@ -428,7 +394,7 @@ def test_neighborhood_enrichment_mode_is_tunable() -> None: node_table = pd.DataFrame( { "instance_id": ["n0", "n1"], - "cluster_label": ["A", "B"], + "niche_label": ["A", "B"], } ) spatial_table = pd.DataFrame( @@ -440,7 +406,7 @@ def test_neighborhood_enrichment_mode_is_tunable() -> None: zscore = run_neighborhood_enrichment( node_table, spatial_table, - label_column="cluster_label", + label_column="niche_label", id_column="instance_id", n_perms=0, undirected=False, @@ -450,7 +416,7 @@ def test_neighborhood_enrichment_mode_is_tunable() -> None: difference = run_neighborhood_enrichment( node_table, spatial_table, - label_column="cluster_label", + label_column="niche_label", id_column="instance_id", n_perms=0, undirected=False, @@ -460,7 +426,7 @@ def test_neighborhood_enrichment_mode_is_tunable() -> None: log2fc = run_neighborhood_enrichment( node_table, spatial_table, - label_column="cluster_label", + label_column="niche_label", id_column="instance_id", n_perms=0, undirected=False, @@ -476,7 +442,7 @@ def test_neighborhood_enrichment_mode_rejects_unknown_values() -> None: node_table = pd.DataFrame( { "instance_id": ["n0", "n1"], - "cluster_label": ["A", "B"], + "niche_label": ["A", "B"], } ) spatial_table = pd.DataFrame( @@ -489,7 +455,7 @@ def test_neighborhood_enrichment_mode_rejects_unknown_values() -> None: run_neighborhood_enrichment( node_table, spatial_table, - label_column="cluster_label", + label_column="niche_label", id_column="instance_id", n_perms=0, warn_analytical=False, @@ -501,7 +467,7 @@ def test_diff_neighborhood_returns_pairwise_differences_with_pvalues() -> None: node_table = pd.DataFrame( { "instance_id": ["x0", "x1", "y0", "y1"], - "cluster_label": ["A", "B", "A", "B"], + "niche_label": ["A", "B", "A", "B"], "sample_id": ["sx", "sx", "sy", "sy"], "condition": ["X", "X", "Y", "Y"], } @@ -516,7 +482,7 @@ def test_diff_neighborhood_returns_pairwise_differences_with_pvalues() -> None: out = run_diff_neighborhood_enrichment( node_table, spatial_table, - label_column="cluster_label", + label_column="niche_label", condition_column="condition", id_column="instance_id", n_perms=8, @@ -550,13 +516,13 @@ def _fake_run_neighborhood(*args, **kwargs): # type: ignore[no-untyped-def] ) monkeypatch.setattr( - "src.interpretability.tier2.neighborhood.run_neighborhood_enrichment", + "grass_mil.interpretability.tier2.neighborhood.run_neighborhood_enrichment", _fake_run_neighborhood, ) node_table = pd.DataFrame( { "instance_id": ["x0", "x1", "y0", "y1"], - "cluster_label": ["A", "A", "A", "A"], + "niche_label": ["A", "A", "A", "A"], "sample_id": ["sx0", "sx1", "sy0", "sy1"], "condition": ["X", "X", "Y", "Y"], } @@ -565,7 +531,7 @@ def _fake_run_neighborhood(*args, **kwargs): # type: ignore[no-untyped-def] out = run_diff_neighborhood_enrichment( node_table, spatial_table, - label_column="cluster_label", + label_column="niche_label", condition_column="condition", permutation_group_column="sample_id", n_perms=4, @@ -580,7 +546,7 @@ def test_diff_neighborhood_respects_enrichment_mode() -> None: node_table = pd.DataFrame( { "instance_id": ["x0", "x1", "y0", "y1"], - "cluster_label": ["A", "B", "A", "B"], + "niche_label": ["A", "B", "A", "B"], "condition": ["X", "X", "Y", "Y"], } ) @@ -593,7 +559,7 @@ def test_diff_neighborhood_respects_enrichment_mode() -> None: zscore_out = run_diff_neighborhood_enrichment( node_table, spatial_table, - label_column="cluster_label", + label_column="niche_label", condition_column="condition", id_column="instance_id", n_perms=0, @@ -604,7 +570,7 @@ def test_diff_neighborhood_respects_enrichment_mode() -> None: diff_out = run_diff_neighborhood_enrichment( node_table, spatial_table, - label_column="cluster_label", + label_column="niche_label", condition_column="condition", id_column="instance_id", n_perms=0, @@ -620,7 +586,7 @@ def test_diff_neighborhood_requires_permutation_group_column_for_permutation_mod node_table = pd.DataFrame( { "instance_id": ["x0", "x1", "y0", "y1"], - "cluster_label": ["A", "B", "A", "B"], + "niche_label": ["A", "B", "A", "B"], "condition": ["X", "X", "Y", "Y"], } ) @@ -634,7 +600,7 @@ def test_diff_neighborhood_requires_permutation_group_column_for_permutation_mod run_diff_neighborhood_enrichment( node_table, spatial_table, - label_column="cluster_label", + label_column="niche_label", condition_column="condition", n_perms=4, warn_analytical=False, @@ -645,7 +611,7 @@ def test_diff_neighborhood_rejects_mixed_condition_groups() -> None: node_table = pd.DataFrame( { "instance_id": ["x0", "x1", "y0", "y1"], - "cluster_label": ["A", "B", "A", "B"], + "niche_label": ["A", "B", "A", "B"], "sample_id": ["s0", "s0", "s1", "s1"], "condition": ["X", "Y", "X", "Y"], } @@ -660,7 +626,7 @@ def test_diff_neighborhood_rejects_mixed_condition_groups() -> None: run_diff_neighborhood_enrichment( node_table, spatial_table, - label_column="cluster_label", + label_column="niche_label", condition_column="condition", permutation_group_column="sample_id", n_perms=4, @@ -672,7 +638,7 @@ def test_neighborhood_warns_in_analytical_mode() -> None: node_table = pd.DataFrame( { "instance_id": ["n0", "n1"], - "cluster_label": ["A", "B"], + "niche_label": ["A", "B"], } ) spatial_table = pd.DataFrame( @@ -685,7 +651,7 @@ def test_neighborhood_warns_in_analytical_mode() -> None: run_neighborhood_enrichment( node_table, spatial_table, - label_column="cluster_label", + label_column="niche_label", id_column="instance_id", n_perms=0, ) @@ -695,7 +661,7 @@ def test_diff_neighborhood_warns_about_analytical_baseline() -> None: node_table = pd.DataFrame( { "instance_id": ["x0", "x1", "y0", "y1"], - "cluster_label": ["A", "B", "A", "B"], + "niche_label": ["A", "B", "A", "B"], "condition": ["X", "X", "Y", "Y"], } ) @@ -709,7 +675,7 @@ def test_diff_neighborhood_warns_about_analytical_baseline() -> None: run_diff_neighborhood_enrichment( node_table, spatial_table, - label_column="cluster_label", + label_column="niche_label", condition_column="condition", id_column="instance_id", n_perms=0, @@ -721,26 +687,26 @@ def _legacy_filtration_curves( spatial_table: pd.DataFrame, *, thresholds: np.ndarray, - cluster_column: str, + 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)[cluster_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) - clusters = sorted(node_table[cluster_column].astype(str).unique()) + 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 clusters + c: {ct: np.zeros_like(thresholds, dtype=float) for ct in cell_types} for c in niches } - for cluster in clusters: - c_edges = edges[edges["source_cluster"] == str(cluster)] + 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) } @@ -758,7 +724,7 @@ def _legacy_filtration_curves( if ct is not None: threshold_counts[ct] += 1.0 for ct in cell_types: - curves[cluster][ct][i] = float(threshold_counts.get(ct, 0.0)) + curves[niche][ct][i] = float(threshold_counts.get(ct, 0.0)) return curves @@ -769,7 +735,7 @@ def test_filtration_vectorized_matches_legacy_counts() -> None: node_table = pd.DataFrame( { "instance_id": node_ids, - "cluster_label": rng.choice(["0", "1", "2"], size=n_nodes), + "niche_label": rng.choice(["0", "1", "2"], size=n_nodes), "cell_type": rng.choice(["A", "B", "C"], size=n_nodes), } ) @@ -787,7 +753,7 @@ def test_filtration_vectorized_matches_legacy_counts() -> None: node_table, spatial_table, thresholds=thresholds, - cluster_column="cluster_label", + niche_column="niche_label", cell_type_column="cell_type", id_column="instance_id", distance_column="distance", @@ -796,15 +762,15 @@ def test_filtration_vectorized_matches_legacy_counts() -> None: node_table, spatial_table, thresholds=thresholds, - cluster_column="cluster_label", + niche_column="niche_label", cell_type_column="cell_type", id_column="instance_id", distance_column="distance", scale_within_cluster=False, ) - for cluster, ct_map in legacy.items(): + for niche, ct_map in legacy.items(): for cell_type, values in ct_map.items(): - assert np.allclose(out.curves[cluster][cell_type], values) + assert np.allclose(out.curves[niche][cell_type], values) def test_filtration_runtime_sanity() -> None: @@ -814,7 +780,7 @@ def test_filtration_runtime_sanity() -> None: node_table = pd.DataFrame( { "instance_id": node_ids, - "cluster_label": rng.choice(["0", "1", "2", "3"], size=n_nodes), + "niche_label": rng.choice(["0", "1", "2", "3"], size=n_nodes), "cell_type": rng.choice(["A", "B", "C", "D"], size=n_nodes), } ) @@ -832,7 +798,7 @@ def test_filtration_runtime_sanity() -> None: node_table, spatial_table, thresholds=thresholds, - cluster_column="cluster_label", + niche_column="niche_label", cell_type_column="cell_type", id_column="instance_id", distance_column="distance", @@ -840,3 +806,150 @@ def test_filtration_runtime_sanity() -> None: elapsed = time.perf_counter() - t0 assert len(out.curves) > 0 assert elapsed < 8.0 + + +def test_builtin_per_niche_cell_type_enrichment_plugin_execution() -> None: + registry = create_builtin_registry() + ds = _sample_dataset() + labels = np.array([0, 0, 0, 1, 1, 1]) + plugin = registry.get("per_niche_cell_type_enrichment") + out = plugin.run( + ds, + PluginContext(state={"niche_labels": labels}), + n_perms=0, + enrichment_mode="obs-exp", + ) + assert out.name == "per_niche_cell_type_enrichment" + enrichment = out.payload["enrichment_by_niche"] + # Both niches are internally connected in the sample spatial table. + assert set(enrichment) == {"0", "1"} + for frame in enrichment.values(): + # Enrichment is a square cell-type x cell-type matrix. + assert list(frame.index) == list(frame.columns) + assert all(section.title.startswith("Cell-Type Enrichment") for section in out.sections) + + +def test_per_niche_cell_type_enrichment_skips_noise_cluster() -> None: + registry = create_builtin_registry() + ds = _sample_dataset() + labels = np.array([-1, -1, -1, 1, 1, 1]) + plugin = registry.get("per_niche_cell_type_enrichment") + + skipped = plugin.run( + ds, + PluginContext(state={"niche_labels": labels}), + skip_background=True, + n_perms=0, + ) + assert set(skipped.payload["enrichment_by_niche"]) == {"1"} + + retained = plugin.run( + ds, + PluginContext(state={"niche_labels": labels}), + skip_background=False, + n_perms=0, + ) + assert set(retained.payload["enrichment_by_niche"]) == {"-1", "1"} + + +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_morans_i_plugin_runs_over_the_instance_graph() -> None: + registry = create_builtin_registry() + ds = _sample_dataset() + labels = np.array([0, 0, 0, 1, 1, 1]) + + out = registry.get("morans_i").run( + ds, PluginContext(state={"niche_labels": labels}), n_perms=10, min_nodes=2 + ) + assert out.name == "morans_i" + assert set(out.payload["statistic"].index) == {0, 1} + assert "comp_A" in out.payload["statistic"].columns + assert (out.payload["qvalue"].to_numpy(dtype=float) >= 0).any() + + +def test_ripley_plugin_produces_centred_curves() -> None: + registry = create_builtin_registry() + ds = _sample_dataset() + labels = np.array([0, 0, 1, 1, 0, 1]) + + out = registry.get("ripley").run( + ds, PluginContext(state={"niche_labels": labels}), n_radii=6, min_count=2 + ) + assert out.name == "ripley" + curves = out.payload["curves"] + assert len(out.payload["radii"]) == 6 + assert curves.shape[0] == 6 + + +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 new file mode 100644 index 0000000..dadf4c3 --- /dev/null +++ b/tests/interpretability/test_spatial_statistics.py @@ -0,0 +1,262 @@ +"""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 ( + diff_morans_i_vs_reference, + 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 test_diff_morans_i_reports_a_z_statistic_against_the_reference() -> None: + nodes, edges = _autocorrelation_tables() + nodes["niche_label"] = nodes["niche_label"].map({"a": 0, "b": -1}) + result = run_morans_i( + nodes, + edges, + feature_columns=["feature"], + group_column="niche_label", + n_perms=50, + random_state=0, + ) + diff = diff_morans_i_vs_reference(result, reference_group=-1) + assert list(diff.index) == [0] + assert diff.loc[0, "feature"] > 0 + + with pytest.raises(ValueError, match="not present"): + diff_morans_i_vs_reference(result, reference_group=99) + + +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 index 135b074..11150f7 100644 --- a/tests/interpretability/test_tissue_graph.py +++ b/tests/interpretability/test_tissue_graph.py @@ -3,7 +3,7 @@ import pandas as pd import pytest -from src.interpretability.tier2.tissue_graph import ( +from grass_mil.interpretability.tier2.tissue_graph import ( build_tissue_graph_figure, prepare_tissue_graph_view, ) @@ -18,7 +18,7 @@ def _sample_tables() -> tuple[pd.DataFrame, pd.DataFrame]: "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], - "cluster_label": [0, 1, 1, 0, 2], + "niche_label": [0, 1, 1, 0, 2], } ) spatial_table = pd.DataFrame( @@ -41,7 +41,7 @@ def test_prepare_tissue_graph_view_happy_path_filters_to_selected_sample() -> No id_column="instance_id", x_column="center_x", y_column="center_y", - label_column="cluster_label", + label_column="niche_label", ) assert len(view.nodes) == 3 @@ -52,7 +52,7 @@ def test_prepare_tissue_graph_view_happy_path_filters_to_selected_sample() -> No assert view.metadata["sample_value"] == "sx" assert view.metadata["node_count"] == 3 assert view.metadata["edge_count"] == 2 - assert view.metadata["cluster_counts"] == {"0": 1, "1": 2} + assert view.metadata["niche_counts"] == {"0": 1, "1": 2} def test_prepare_tissue_graph_view_fails_for_missing_coords() -> None: @@ -66,7 +66,7 @@ def test_prepare_tissue_graph_view_fails_for_missing_coords() -> None: id_column="instance_id", x_column="center_x", y_column="center_y", - label_column="cluster_label", + label_column="niche_label", ) diff --git a/tests/interpretability/test_transfer.py b/tests/interpretability/test_transfer.py index ad9927d..7466bff 100644 --- a/tests/interpretability/test_transfer.py +++ b/tests/interpretability/test_transfer.py @@ -7,13 +7,17 @@ import pytest from sklearn.neighbors import KNeighborsClassifier -from src.interpretability.contracts import ClusteringResult, InterpretabilityDataset, ReportBundle -from src.interpretability.core.reduction import run_reduction -from src.interpretability.core.transfer import ( - apply_cluster_transfer, - fit_cluster_transfer_from_report_bundle, - load_cluster_transfer_bundle, - save_cluster_transfer_bundle, +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, ) @@ -47,16 +51,16 @@ def _build_report_bundle( return ReportBundle( dataset=dataset, reduction=pca_result, - cluster_feature_reduction=pca_result if cluster_on_pca else None, + niche_feature_reduction=pca_result if cluster_on_pca else None, clustering=clustering, - cluster_summary=None, + niche_summary=None, plugin_results={}, artifacts_dir=artifacts_dir, metadata={"cluster_on_pca": cluster_on_pca}, ) -def test_cluster_transfer_fit_and_apply_matches_direct_pca_knn_reference(tmp_path: Path) -> None: +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) @@ -68,14 +72,14 @@ def test_cluster_transfer_fit_and_apply_matches_direct_pca_knn_reference(tmp_pat pca_components=2, ) - bundle = fit_cluster_transfer_from_report_bundle(dataset, report_bundle, n_neighbors=3) + 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_cluster_transfer(bundle, query, id_column="instance_id") + 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.cluster_feature_reduction.embedding, labels) # type: ignore[union-attr] - ref_features = report_bundle.cluster_feature_reduction.fitted_object.transform( # type: ignore[union-attr] + 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) @@ -85,7 +89,7 @@ def test_cluster_transfer_fit_and_apply_matches_direct_pca_knn_reference(tmp_pat assert np.allclose(out["transfer_confidence"].to_numpy(dtype=float), ref_conf) -def test_cluster_transfer_preserves_noise_label_minus_one(tmp_path: Path) -> None: +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) @@ -95,15 +99,15 @@ def test_cluster_transfer_preserves_noise_label_minus_one(tmp_path: Path) -> Non artifacts_dir=tmp_path, cluster_on_pca=False, ) - bundle = fit_cluster_transfer_from_report_bundle(dataset, report_bundle, n_neighbors=1) - out = apply_cluster_transfer(bundle, table[["instance_id", "inst_emb_0", "inst_emb_1"]]) + 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_cluster_transfer_bundle_serialization_roundtrip(tmp_path: Path) -> None: +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) @@ -114,19 +118,19 @@ def test_cluster_transfer_bundle_serialization_roundtrip(tmp_path: Path) -> None cluster_on_pca=True, pca_components=2, ) - bundle = fit_cluster_transfer_from_report_bundle(dataset, report_bundle, n_neighbors=3) + bundle = fit_niche_transfer_from_report_bundle(dataset, report_bundle, n_neighbors=3) bundle_path = tmp_path / "cluster_transfer_bundle.joblib" - save_cluster_transfer_bundle(bundle, bundle_path) - restored = load_cluster_transfer_bundle(bundle_path) + 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_cluster_transfer(bundle, query) - right = apply_cluster_transfer(restored, query) + left = apply_niche_transfer(bundle, query) + right = apply_niche_transfer(restored, query) assert left.equals(right) -def test_cluster_transfer_query_column_validation_and_reordered_columns(tmp_path: Path) -> None: +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) @@ -136,19 +140,19 @@ def test_cluster_transfer_query_column_validation_and_reordered_columns(tmp_path artifacts_dir=tmp_path, cluster_on_pca=False, ) - bundle = fit_cluster_transfer_from_report_bundle(dataset, report_bundle, n_neighbors=3) + 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_cluster_transfer(bundle, ordered) - out_reordered = apply_cluster_transfer(bundle, reordered) + 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_cluster_transfer(bundle, table[["instance_id", "inst_emb_0"]]) + apply_niche_transfer(bundle, table[["instance_id", "inst_emb_0"]]) -def test_cluster_transfer_fit_fails_fast_for_missing_inputs(tmp_path: Path) -> None: +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) @@ -163,34 +167,34 @@ def test_cluster_transfer_fit_fails_fast_for_missing_inputs(tmp_path: Path) -> N missing_clustering = ReportBundle( dataset=dataset, reduction=valid.reduction, - cluster_feature_reduction=valid.cluster_feature_reduction, + niche_feature_reduction=valid.niche_feature_reduction, clustering=None, - cluster_summary=None, + niche_summary=None, plugin_results={}, artifacts_dir=tmp_path, metadata={"cluster_on_pca": True}, ) with pytest.raises(ValueError, match="missing clustering result"): - fit_cluster_transfer_from_report_bundle(dataset, missing_clustering, n_neighbors=3) + fit_niche_transfer_from_report_bundle(dataset, missing_clustering, n_neighbors=3) missing_cluster_pca = ReportBundle( dataset=dataset, reduction=valid.reduction, - cluster_feature_reduction=None, + niche_feature_reduction=None, clustering=valid.clustering, - cluster_summary=None, + niche_summary=None, plugin_results={}, artifacts_dir=tmp_path, metadata={"cluster_on_pca": True}, ) - with pytest.raises(ValueError, match="missing cluster_feature_reduction"): - fit_cluster_transfer_from_report_bundle(dataset, missing_cluster_pca, n_neighbors=3) + 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_cluster_transfer_from_report_bundle(dataset, valid, n_neighbors=len(table) + 1) + fit_niche_transfer_from_report_bundle(dataset, valid, n_neighbors=len(table) + 1) -def test_cluster_transfer_apply_fails_for_empty_input(tmp_path: Path) -> None: +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) @@ -200,7 +204,7 @@ def test_cluster_transfer_apply_fails_for_empty_input(tmp_path: Path) -> None: artifacts_dir=tmp_path, cluster_on_pca=False, ) - bundle = fit_cluster_transfer_from_report_bundle(dataset, report_bundle, n_neighbors=3) + 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_cluster_transfer(bundle, empty) + apply_niche_transfer(bundle, empty) diff --git a/tests/models/test_attention.py b/tests/models/test_attention.py index a8f6265..e513d29 100644 --- a/tests/models/test_attention.py +++ b/tests/models/test_attention.py @@ -3,7 +3,7 @@ def test_gated_attention_shape(): - from src.models.components.attention import AttnNetGated + 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) @@ -13,7 +13,7 @@ def test_gated_attention_shape(): def test_gated_projected_attention_shape(): - from src.models.components.attention import AttnNetGatedProjected + from grass_mil.models.components.attention import AttnNetGatedProjected module = AttnNetGatedProjected( input_dim=32, projection_dim=16, hidden_dim=8, dropout=False, n_classes=1 @@ -26,10 +26,43 @@ def test_gated_projected_attention_shape(): @pytest.mark.parametrize("enabled", [True, False]) def test_attention_optional_factory(enabled): - from src.models.components.factory import build_attention + 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 new file mode 100644 index 0000000..3ef132c --- /dev/null +++ b/tests/models/test_augmentations.py @@ -0,0 +1,243 @@ +"""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 index 9c000e1..886d446 100644 --- a/tests/models/test_backbones.py +++ b/tests/models/test_backbones.py @@ -17,7 +17,7 @@ def _toy_graph(num_nodes: int = 12): @pytest.mark.parametrize("conv_type", ["gin", "gcn", "gat", "graphsage", "gine"]) def test_encoder_forward_node_and_graph(conv_type): pytest.importorskip("torch_geometric") - from src.models.components.backbones import EncoderConfig, GNNEncoder + from grass_mil.models.components.backbones import EncoderConfig, GNNEncoder data = _toy_graph() cfg = EncoderConfig( @@ -44,7 +44,7 @@ def test_encoder_forward_node_and_graph(conv_type): def test_encoder_warns_when_edge_attr_enabled_on_unsupported_backbone(): pytest.importorskip("torch_geometric") - from src.models.components.backbones import EncoderConfig, GNNEncoder + from grass_mil.models.components.backbones import EncoderConfig, GNNEncoder with pytest.warns(UserWarning, match="use_edge_attr=True is ignored"): GNNEncoder( @@ -62,7 +62,7 @@ def test_encoder_warns_when_edge_attr_enabled_on_unsupported_backbone(): @pytest.mark.parametrize("jk", ["last", "concat", "max", "sum"]) def test_encoder_jk_modes(jk): pytest.importorskip("torch_geometric") - from src.models.components.backbones import EncoderConfig, GNNEncoder + from grass_mil.models.components.backbones import EncoderConfig, GNNEncoder data = _toy_graph() cfg = EncoderConfig( diff --git a/tests/models/test_bagging.py b/tests/models/test_bagging.py new file mode 100644 index 0000000..0f7934d --- /dev/null +++ b/tests/models/test_bagging.py @@ -0,0 +1,106 @@ +"""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_losses.py b/tests/models/test_losses.py index 512199e..4551a3f 100644 --- a/tests/models/test_losses.py +++ b/tests/models/test_losses.py @@ -3,7 +3,7 @@ def test_weighted_cross_entropy(): - from src.models.components.losses import WeightedCrossEntropyLoss + from grass_mil.models.components.losses import WeightedCrossEntropyLoss loss_fn = WeightedCrossEntropyLoss() logits = torch.randn(12, 3) @@ -15,7 +15,7 @@ def test_weighted_cross_entropy(): def test_weighted_bce(): - from src.models.components.losses import WeightedBCEWithLogitsLoss + from grass_mil.models.components.losses import WeightedBCEWithLogitsLoss loss_fn = WeightedBCEWithLogitsLoss() logits = torch.randn(12, 1) @@ -28,7 +28,7 @@ def test_weighted_bce(): @pytest.mark.parametrize("loss_name", ["mse", "huber"]) def test_regression_losses(loss_name): - from src.models.components.losses import WeightedHuberLoss, WeightedMSELoss + from grass_mil.models.components.losses import WeightedHuberLoss, WeightedMSELoss pred = torch.randn(12, 2) target = torch.randn(12, 2) @@ -40,7 +40,7 @@ def test_regression_losses(loss_name): def test_coxsgd_loss(): - from src.models.components.losses import CoxSGDLoss + from grass_mil.models.components.losses import CoxSGDLoss loss_fn = CoxSGDLoss(top_n=2, regularizer_weight=0.01) y_pred = torch.randn(16, 1) @@ -52,7 +52,7 @@ def test_coxsgd_loss(): def test_coxsgd_top_n_handles_tiny_batches(): - from src.models.components.losses import CoxSGDLoss + from grass_mil.models.components.losses import CoxSGDLoss loss_fn = CoxSGDLoss(top_n=10, regularizer_weight=0.01) y_pred = torch.randn(1, 1) diff --git a/tests/models/test_node_embeddings.py b/tests/models/test_node_embeddings.py new file mode 100644 index 0000000..d6a0663 --- /dev/null +++ b/tests/models/test_node_embeddings.py @@ -0,0 +1,181 @@ +"""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 index e4e3b86..aa90e3d 100644 --- a/tests/models/test_optimization.py +++ b/tests/models/test_optimization.py @@ -14,7 +14,6 @@ def _encoder_cfg(): "jk": "last", "act": "relu", "pooling": "mean", - "set2set_steps": 3, "gat_heads": 2, "use_edge_attr": False, "edge_weight_index": 0, @@ -44,7 +43,7 @@ def _cosine_step_cfg(): def test_supervised_mil_uses_explicit_backbone_and_attention_lrs(): pytest.importorskip("torch_geometric") - from src.models.supervised_module import SupervisedModule + from grass_mil.models.supervised_module import SupervisedModule module = SupervisedModule( encoder=_encoder_cfg(), @@ -86,7 +85,7 @@ def test_supervised_mil_uses_explicit_backbone_and_attention_lrs(): def test_supervised_mean_warmup_is_opt_in(): pytest.importorskip("torch_geometric") - from src.models.supervised_module import SupervisedModule + from grass_mil.models.supervised_module import SupervisedModule base_kwargs = dict( encoder=_encoder_cfg(), @@ -130,7 +129,7 @@ def test_supervised_mean_warmup_is_opt_in(): def test_bgrl_uses_warmup_then_base_scheduler(): pytest.importorskip("torch_geometric") - from src.models.bgrl_module import BGRLModule + from grass_mil.models.bgrl_module import BGRLModule module = BGRLModule( encoder=_encoder_cfg(), diff --git a/tests/models/test_runtime.py b/tests/models/test_runtime.py index 481ed27..570a2f0 100644 --- a/tests/models/test_runtime.py +++ b/tests/models/test_runtime.py @@ -3,7 +3,7 @@ def test_remap_encoder_keys_auto_bgrl_or_identity(): - from src.models.training.checkpoint_init import remap_encoder_keys + from grass_mil.models.training.checkpoint_init import remap_encoder_keys sd = { "online_encoder.layers.0.weight": torch.randn(4, 4), @@ -17,13 +17,13 @@ def test_remap_encoder_keys_auto_bgrl_or_identity(): def test_compute_supervised_loss_routes(): - from src.models.components.losses import ( + from grass_mil.models.components.losses import ( CoxSGDLoss, WeightedBCEWithLogitsLoss, WeightedCrossEntropyLoss, WeightedMSELoss, ) - from src.models.training.loss_utils import compute_supervised_loss + from grass_mil.models.training.loss_utils import compute_supervised_loss bce = compute_supervised_loss( loss_fn=WeightedBCEWithLogitsLoss(), @@ -69,14 +69,14 @@ def test_compute_supervised_loss_routes(): def test_validate_task_config_accepts_categorical(): - from src.models.training.builders import validate_task_config + 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 src.models.supervised_module import SupervisedModule + from grass_mil.models.supervised_module import SupervisedModule module = SupervisedModule( encoder={}, @@ -107,7 +107,7 @@ def _capture_log(name, value, **kwargs): def test_supervised_module_logs_epoch_validation_accuracy_from_aggregated_regions(): pytest.importorskip("lightning") - from src.models.supervised_module import SupervisedModule + from grass_mil.models.supervised_module import SupervisedModule module = SupervisedModule( encoder={}, @@ -143,7 +143,7 @@ def _capture_log(name, value, **kwargs): def test_supervised_module_logs_epoch_validation_regression_metrics(): pytest.importorskip("lightning") - from src.models.supervised_module import SupervisedModule + from grass_mil.models.supervised_module import SupervisedModule module = SupervisedModule( encoder={}, @@ -180,7 +180,7 @@ def _capture_log(name, value, **kwargs): def test_supervised_module_logs_epoch_validation_survival_c_index(): pytest.importorskip("lightning") - from src.models.supervised_module import SupervisedModule + from grass_mil.models.supervised_module import SupervisedModule module = SupervisedModule( encoder={}, @@ -216,8 +216,8 @@ def _capture_log(name, value, **kwargs): def test_attention_bag_aggregation_shapes(): import torch - from src.models.components.attention import AttnNetGatedProjected - from src.models.training.bagging import aggregate_bag_logits_attention + 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) @@ -237,57 +237,10 @@ def test_attention_bag_aggregation_shapes(): assert len(bag_indices) == 2 -def test_build_mil_aux_targets_attention_shaped(): - from src.models.training.loss_utils import build_mil_aux_targets - - bag_targets = torch.tensor([[1.0], [0.0]]) - bag_indices = [[0, 1], [2, 3, 4]] - bag_ids = ["r0", "r1"] - bag_attention = { - "r0": torch.tensor([0.25, 0.75]), - "r1": torch.tensor([0.2, 0.3, 0.5]), - } - targets, weights = build_mil_aux_targets( - bag_targets=bag_targets, - bag_indices=bag_indices, - bag_ids=bag_ids, - bag_attention=bag_attention, - target_mode="attention_shaped_ti", - ) - assert targets.shape == (5, 1) - assert weights.shape == (5, 1) - assert torch.all(targets[:2] >= 0.5) - assert torch.all(targets[2:] <= 0.5) - - -def test_compute_aux_and_entropy_terms(): - from src.models.training.loss_utils import ( - compute_aux_node_loss, - compute_entropy_regularization, - ) - - logits = torch.tensor([[0.2], [-0.3], [0.9]]) - targets = torch.tensor([[1.0], [0.0], [1.0]]) - weights = torch.tensor([[0.5], [0.25], [0.25]]) - bce = compute_aux_node_loss( - aux_logits=logits, aux_targets=targets, aux_weights=None, loss_mode="bce" - ) - weighted = compute_aux_node_loss( - aux_logits=logits, - aux_targets=targets, - aux_weights=weights, - loss_mode="weighted_bce", - ) - ent = compute_entropy_regularization(targets) - assert bce.ndim == 0 - assert weighted.ndim == 0 - assert ent.ndim == 0 - - def test_select_target_columns_handles_single_label_list(): from types import SimpleNamespace - from src.models.training.bagging import select_target_columns + from grass_mil.models.training.bagging import select_target_columns graph_y = torch.tensor([[1.0]]) batch = SimpleNamespace(graph_label_names=["label"]) @@ -299,7 +252,7 @@ def test_select_target_columns_handles_single_label_list(): def test_supervised_predict_step_allows_missing_targets(): pytest.importorskip("lightning") pytest.importorskip("torch_geometric") - from src.models.supervised_module import SupervisedModule + from grass_mil.models.supervised_module import SupervisedModule module = SupervisedModule( encoder={}, @@ -331,7 +284,7 @@ def _forward_bags(_batch, *, allow_missing_targets=False): def test_supervised_predict_step_can_disable_instance_payload(): pytest.importorskip("lightning") pytest.importorskip("torch_geometric") - from src.models.supervised_module import SupervisedModule + from grass_mil.models.supervised_module import SupervisedModule module = SupervisedModule( encoder={}, @@ -394,7 +347,7 @@ def _build_predict_group_metadata( def test_supervised_predict_step_can_emit_embedding_payload(): pytest.importorskip("lightning") pytest.importorskip("torch_geometric") - from src.models.supervised_module import SupervisedModule + from grass_mil.models.supervised_module import SupervisedModule module = SupervisedModule( encoder={}, @@ -454,7 +407,7 @@ class _Batch: def test_supervised_predict_step_omits_embedding_payload_by_default(): pytest.importorskip("lightning") pytest.importorskip("torch_geometric") - from src.models.supervised_module import SupervisedModule + from grass_mil.models.supervised_module import SupervisedModule module = SupervisedModule( encoder={}, @@ -508,7 +461,7 @@ class _Batch: def test_extract_bag_ids_falls_back_per_item_when_primary_missing(): from types import SimpleNamespace - from src.models.training.bagging import extract_bag_ids + 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") @@ -518,7 +471,7 @@ def test_extract_bag_ids_falls_back_per_item_when_primary_missing(): def test_extract_bag_ids_namespaces_primary_by_fallback_for_uniqueness(): from types import SimpleNamespace - from src.models.training.bagging import extract_bag_ids + 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") @@ -528,7 +481,7 @@ def test_extract_bag_ids_namespaces_primary_by_fallback_for_uniqueness(): def test_extract_bag_ids_patch_key_uses_sample_region_patch_hierarchy(): from types import SimpleNamespace - from src.models.training.bagging import extract_bag_ids + from grass_mil.models.training.bagging import extract_bag_ids batch = SimpleNamespace( patch_id=["p0", "p1"], @@ -542,7 +495,7 @@ def test_extract_bag_ids_patch_key_uses_sample_region_patch_hierarchy(): def test_extract_bag_ids_patch_key_fallback_uses_sample_region_when_patch_missing(): from types import SimpleNamespace - from src.models.training.bagging import extract_bag_ids + from grass_mil.models.training.bagging import extract_bag_ids batch = SimpleNamespace( patch_id=[None, ""], @@ -558,7 +511,7 @@ def test_predict_group_metadata_uses_batch_patch_ids_for_instance_patch_ids(): pytest.importorskip("torch_geometric") from types import SimpleNamespace - from src.models.supervised_module import SupervisedModule + from grass_mil.models.supervised_module import SupervisedModule module = SupervisedModule( encoder={}, @@ -585,7 +538,7 @@ def test_predict_group_metadata_uses_batch_patch_ids_for_instance_patch_ids(): def test_gather_bag_targets_requires_consistent_labels_within_bag(): from types import SimpleNamespace - from src.models.training.bagging import gather_bag_targets + 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), @@ -599,26 +552,10 @@ def test_gather_bag_targets_requires_consistent_labels_within_bag(): ) -def test_attention_shaped_aux_targets_are_detached(): - from src.models.training.loss_utils import build_mil_aux_targets - - bag_targets = torch.tensor([[1.0]], dtype=torch.float32) - bag_attention = {"b0": torch.tensor([0.2, 0.8], requires_grad=True)} - targets, weights = build_mil_aux_targets( - bag_targets=bag_targets, - bag_indices=[[0, 1]], - bag_ids=["b0"], - bag_attention=bag_attention, - target_mode="attention_shaped_ti", - ) - assert targets.requires_grad is False - assert weights.requires_grad is False - - def test_region_buffer_flushes_in_hyperbatch_chunks(): pytest.importorskip("lightning") pytest.importorskip("torch_geometric") - from src.models.supervised_module import SupervisedModule + from grass_mil.models.supervised_module import SupervisedModule module = SupervisedModule( encoder={}, @@ -658,7 +595,7 @@ def step(self): module.manual_backward = lambda loss: None # type: ignore[method-assign] module.log = lambda *args, **kwargs: None # type: ignore[method-assign] - module._region_total_loss_buffer = [ + module._region_loss_buffer = [ torch.tensor(1.0), torch.tensor(2.0), torch.tensor(3.0), @@ -667,12 +604,12 @@ def step(self): ] flushes = module._flush_region_buffer_if_needed(force=False) assert len(flushes) == 2 - assert len(module._region_total_loss_buffer) == 1 + 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_total_loss_buffer) == 0 + 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 index 9f90a20..25c851a 100644 --- a/tests/models/test_ssl.py +++ b/tests/models/test_ssl.py @@ -13,7 +13,7 @@ def forward(self, x): def test_bgrl_forward_and_update(): - from src.models.components.ssl import BGRL, MLPPredictor + from grass_mil.models.components.ssl import BGRL, MLPPredictor encoder = _TinyEncoder() predictor = MLPPredictor(input_size=8, output_size=8, hidden_size=16) @@ -33,8 +33,8 @@ def test_bgrl_forward_and_update(): @pytest.mark.parametrize("enabled", [True, False]) def test_ssl_optional_factory(enabled): - from src.models.components.backbones import EncoderConfig, GNNEncoder - from src.models.components.factory import build_ssl + 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( @@ -52,8 +52,8 @@ def test_ssl_optional_factory(enabled): def test_ssl_factory_requires_hidden_size(): - from src.models.components.backbones import EncoderConfig, GNNEncoder - from src.models.components.factory import build_ssl + 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"): @@ -61,8 +61,8 @@ def test_ssl_factory_requires_hidden_size(): def test_ssl_factory_accepts_hydra_target_in_predictor_cfg(): - from src.models.components.backbones import EncoderConfig, GNNEncoder - from src.models.components.factory import build_ssl + 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( @@ -71,7 +71,7 @@ def test_ssl_factory_accepts_hydra_target_in_predictor_cfg(): { "method": "bgrl", "predictor": { - "_target_": "src.models.components.ssl.MLPPredictor", + "_target_": "grass_mil.models.components.ssl.MLPPredictor", "hidden_size": 16, }, }, diff --git a/tests/test_eval.py b/tests/test_eval.py index 1d76cf2..f1c2d9d 100644 --- a/tests/test_eval.py +++ b/tests/test_eval.py @@ -5,11 +5,10 @@ import torch from hydra.core.hydra_config import HydraConfig from omegaconf import DictConfig, open_dict -from src.eval import evaluate -from src.train import train +from grass_mil.eval import evaluate +from grass_mil.train import train -@pytest.mark.slow 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`. diff --git a/tests/test_loocv.py b/tests/test_loocv.py index f099d3e..86cf868 100644 --- a/tests/test_loocv.py +++ b/tests/test_loocv.py @@ -1,15 +1,15 @@ from pathlib import Path import pytest -from hydra import compose, initialize +from hydra import compose, initialize_config_module from hydra.core.global_hydra import GlobalHydra from omegaconf import DictConfig, open_dict -from src.loocv import _fold_slug, _resolve_selected_folds, _to_float_metrics, run_loocv +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(version_base="1.3", config_path="../configs"): + 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) @@ -96,8 +96,8 @@ def test_loocv_runs_with_safe_isolation_flags(tmp_path: Path, monkeypatch) -> No def _fake_train(_cfg): return {"val/loss": 1.23}, {} - monkeypatch.setattr("src.loocv._discover_viable_fold_ids", lambda _cfg: ["s0", "s1"]) - monkeypatch.setattr("src.loocv.train", _fake_train) + 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" @@ -128,10 +128,10 @@ def _fake_train(_cfg): return {"val/loss": 0.5}, {} monkeypatch.setattr( - "src.loocv._discover_viable_fold_ids", + "grass_mil.loocv._discover_viable_fold_ids", lambda _cfg: ["sampleA::region1", "sampleA//region1"], ) - monkeypatch.setattr("src.loocv.train", _fake_train) + monkeypatch.setattr("grass_mil.loocv.train", _fake_train) summary = run_loocv(cfg) assert len(summary["folds"]) == 2 @@ -148,9 +148,9 @@ def test_loocv_fails_fast_when_selected_fold_not_viable(tmp_path: Path, monkeypa cfg.loocv.per_fold_processed_dir = True cfg.loocv.force_precompute_per_fold = False - monkeypatch.setattr("src.loocv._discover_viable_fold_ids", lambda _cfg: ["s0"]) + monkeypatch.setattr("grass_mil.loocv._discover_viable_fold_ids", lambda _cfg: ["s0"]) monkeypatch.setattr( - "src.loocv.train", + "grass_mil.loocv.train", lambda _cfg: pytest.fail("train() should not be called when preflight fails"), ) @@ -176,8 +176,10 @@ def _fake_train(_cfg): call_count["n"] += 1 return {"val/loss": 0.75}, {} - monkeypatch.setattr("src.loocv._discover_viable_fold_ids", lambda _cfg: ["s0", "s1", "s2"]) - monkeypatch.setattr("src.loocv.train", _fake_train) + 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) diff --git a/tests/test_packaging.py b/tests/test_packaging.py new file mode 100644 index 0000000..5a0de9f --- /dev/null +++ b/tests/test_packaging.py @@ -0,0 +1,68 @@ +"""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_spatial_omics_datamodule.py b/tests/test_spatial_omics_datamodule.py index c21759b..5e2ce54 100644 --- a/tests/test_spatial_omics_datamodule.py +++ b/tests/test_spatial_omics_datamodule.py @@ -63,7 +63,7 @@ def test_spatial_omics_datamodule_precompute(tmp_path: Path) -> None: graph_labels_path = raw_dir / "graph_labels.csv" graph_labels_path.write_text("region_id,outcome\nregion_1,1\n") - from src.data.spatial_omics_datamodule import SpatialOmicsDataModule + from grass_mil.data.spatial_omics_datamodule import SpatialOmicsDataModule dm = SpatialOmicsDataModule( data_dir=str(data_dir), @@ -149,8 +149,8 @@ def test_spatial_omics_datamodule_precompute(tmp_path: Path) -> None: assert hasattr(data, "edge_attr") assert hasattr(data, "graph_y") assert data.x.shape[1] == 2 # gene1, gene2 - assert hasattr(data, "categorical_index") - assert data.categorical_index.shape[1] == 1 + 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()) @@ -197,7 +197,7 @@ def test_spatial_omics_datamodule_indices_no_molecular(tmp_path: Path) -> None: f"sample_1,{csv_path},csv,region_1,{polygon_path}\n" ) - from src.data.spatial_omics_datamodule import SpatialOmicsDataModule + from grass_mil.data.spatial_omics_datamodule import SpatialOmicsDataModule dm = SpatialOmicsDataModule( data_dir=str(data_dir), @@ -265,8 +265,8 @@ def test_spatial_omics_datamodule_indices_no_molecular(tmp_path: Path) -> None: data = dm.train_dataloader().dataset[0] assert hasattr(data, "x") assert data.x.shape[1] == 0 - assert hasattr(data, "categorical_index") - assert data.categorical_index.shape == (2, 1) + assert hasattr(data, "categorical_codes") + assert data.categorical_codes.shape == (2, 1) def test_h5ad_obsm_spatial_coords(tmp_path: Path) -> None: @@ -284,7 +284,7 @@ def test_h5ad_obsm_spatial_coords(tmp_path: Path) -> None: p = tmp_path / "test.h5ad" adata.write_h5ad(p) - from src.data.components.loaders import H5adConfig, load_h5ad_table + from grass_mil.data.components.loaders import H5adConfig, load_h5ad_table table = load_h5ad_table( p, @@ -332,7 +332,7 @@ def test_datamodule_requires_persisted_splits(tmp_path: Path) -> None: ) ) - from src.data.spatial_omics_datamodule import SpatialOmicsDataModule + from grass_mil.data.spatial_omics_datamodule import SpatialOmicsDataModule dm = SpatialOmicsDataModule( data_dir=str(tmp_path), @@ -417,7 +417,7 @@ def test_datamodule_uses_persisted_splits(tmp_path: Path) -> None: ) ) - from src.data.spatial_omics_datamodule import SpatialOmicsDataModule + from grass_mil.data.spatial_omics_datamodule import SpatialOmicsDataModule dm = SpatialOmicsDataModule( data_dir=str(tmp_path), @@ -474,7 +474,7 @@ def test_datamodule_uses_persisted_splits(tmp_path: Path) -> None: def _make_unit( table_idx: int, sample_id: str, region_id: str, patch_idx: int, split: str = "train" ): - from src.data.components.precompute import GraphUnitSpec + from grass_mil.data.components.precompute import GraphUnitSpec return GraphUnitSpec( table_idx=table_idx, @@ -488,7 +488,7 @@ def _make_unit( def test_region_split_groups_all_patches_per_region() -> None: - from src.data.components.precompute import _assign_split_labels + from grass_mil.data.components.precompute import _assign_split_labels units = [ _make_unit(0, "s0", "r0", 0), @@ -508,7 +508,7 @@ def test_region_split_groups_all_patches_per_region() -> None: def test_loocv_holdout_region_is_test_only_with_region_validation() -> None: - from src.data.components.precompute import LoocvConfig, _assign_split_labels + from grass_mil.data.components.precompute import LoocvConfig, _assign_split_labels units = [ _make_unit(0, "s0", "r0", 0), @@ -540,7 +540,7 @@ def test_loocv_holdout_region_is_test_only_with_region_validation() -> None: def test_loocv_patch_validation_uses_only_train_fold_items() -> None: - from src.data.components.precompute import LoocvConfig, _assign_split_labels + from grass_mil.data.components.precompute import LoocvConfig, _assign_split_labels units = [ _make_unit(0, "s0", "r0", 0), @@ -574,7 +574,7 @@ def test_loocv_patch_validation_uses_only_train_fold_items() -> None: def test_loocv_heldout_validation_allows_zero_val_ratio() -> None: pytest.importorskip("lightning") pytest.importorskip("torch_geometric") - from src.data.components.precompute import LoocvConfig, _assign_split_labels + from grass_mil.data.components.precompute import LoocvConfig, _assign_split_labels units = [ _make_unit(0, "s0", "r0", 0), @@ -604,7 +604,7 @@ def test_loocv_heldout_validation_allows_zero_val_ratio() -> None: def test_loocv_heldout_validation_forces_one_for_nonzero_ratio() -> None: pytest.importorskip("lightning") pytest.importorskip("torch_geometric") - from src.data.components.precompute import LoocvConfig, _assign_split_labels + from grass_mil.data.components.precompute import LoocvConfig, _assign_split_labels units = [ _make_unit(0, "s0", "r0", 0), @@ -631,7 +631,7 @@ def test_loocv_heldout_validation_forces_one_for_nonzero_ratio() -> None: def test_loocv_patch_validation_allows_zero_val_ratio() -> None: pytest.importorskip("lightning") pytest.importorskip("torch_geometric") - from src.data.components.precompute import LoocvConfig, _assign_split_labels + from grass_mil.data.components.precompute import LoocvConfig, _assign_split_labels units = [ _make_unit(0, "s0", "r0", 0), @@ -661,7 +661,7 @@ def test_loocv_patch_validation_allows_zero_val_ratio() -> None: def test_loocv_patch_validation_forces_one_for_nonzero_ratio() -> None: pytest.importorskip("lightning") pytest.importorskip("torch_geometric") - from src.data.components.precompute import LoocvConfig, _assign_split_labels + from grass_mil.data.components.precompute import LoocvConfig, _assign_split_labels units = [ _make_unit(0, "s0", "r0", 0), @@ -686,7 +686,7 @@ def test_loocv_patch_validation_forces_one_for_nonzero_ratio() -> None: def test_loocv_supports_sample_fold_unit() -> None: - from src.data.components.precompute import LoocvConfig, _assign_split_labels + from grass_mil.data.components.precompute import LoocvConfig, _assign_split_labels units = [ _make_unit(0, "s0", "r0", 0), @@ -718,7 +718,7 @@ 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 src.data.spatial_omics_datamodule import SpatialOmicsDataModule + from grass_mil.data.spatial_omics_datamodule import SpatialOmicsDataModule return SpatialOmicsDataModule( data_dir=str(tmp_path), @@ -844,7 +844,7 @@ def _build_minimal_datamodule( force_precompute: bool, edge_features: list[str], ): - from src.data.spatial_omics_datamodule import SpatialOmicsDataModule + from grass_mil.data.spatial_omics_datamodule import SpatialOmicsDataModule return SpatialOmicsDataModule( data_dir=str(data_dir), @@ -1008,3 +1008,163 @@ def test_precompute_reuse_rejects_coord_scale_um_mismatch(tmp_path: Path) -> Non 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 index c8ef99f..dfad5b1 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -7,7 +7,7 @@ from hydra.core.hydra_config import HydraConfig from omegaconf import DictConfig, OmegaConf, open_dict -from src.train import ( +from grass_mil.train import ( _filter_val_monitor_callbacks, _requires_zero_validation, train, @@ -124,7 +124,6 @@ def test_train_fast_dev_run_gpu(cfg_train: DictConfig) -> None: @RunIf(min_gpus=1) -@pytest.mark.slow def test_train_epoch_gpu_amp(cfg_train: DictConfig) -> None: """Train 1 epoch on GPU with mixed-precision. @@ -138,7 +137,6 @@ def test_train_epoch_gpu_amp(cfg_train: DictConfig) -> None: train(cfg_train) -@pytest.mark.slow def test_train_epoch_double_val_loop(cfg_train: DictConfig) -> None: """Train 1 epoch with validation loop twice per epoch. @@ -151,7 +149,6 @@ def test_train_epoch_double_val_loop(cfg_train: DictConfig) -> None: train(cfg_train) -@pytest.mark.slow def test_train_ddp_sim(cfg_train: DictConfig) -> None: """Simulate DDP (Distributed Data Parallel) on 2 CPU processes. @@ -168,7 +165,6 @@ def test_train_ddp_sim(cfg_train: DictConfig) -> None: train(cfg_train) -@pytest.mark.slow def test_train_resume(tmp_path: Path, cfg_train: DictConfig) -> None: """Run 1 epoch, finish, and resume for another epoch. diff --git a/tests/test_train_regimes.py b/tests/test_train_regimes.py index 62a4abe..ae54855 100644 --- a/tests/test_train_regimes.py +++ b/tests/test_train_regimes.py @@ -4,19 +4,30 @@ import rootutils from hydra.core.hydra_config import HydraConfig from omegaconf import OmegaConf, open_dict -from src.train import train +from grass_mil.train import train +from tests.helpers.config_paths import CONFIGS_DIR def _load_model_cfg(name: str): - return OmegaConf.load(Path(f"configs/model/{name}.yaml")) + return OmegaConf.load(CONFIGS_DIR / "model" / f"{name}.yaml") def _load_task_cfg(name: str): - return OmegaConf.load(Path(f"configs/task/{name}.yaml")).task + 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 src.models.bgrl_module import BGRLModule + from grass_mil.models.bgrl_module import BGRLModule module = BGRLModule( encoder={ @@ -73,26 +84,24 @@ def test_training_regimes_fast_dev_run( ) cfg_train.task = _load_task_cfg(task_name) cfg_train.model = _load_model_cfg(model_name) - cfg_train.optim = OmegaConf.load(Path("configs/optim/adamw.yaml")) + 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(Path("configs/scheduler") / scheduler_name) + 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 - cfg_train.model.task.node_aux.enabled = True - cfg_train.model.task.node_aux.target_mode = "attention_shaped_ti" - cfg_train.model.task.node_aux.loss_mode = "bce" - cfg_train.model.task.node_aux.weight = 0.2 - cfg_train.model.task.entropy_reg.enabled = True - cfg_train.model.task.entropy_reg.mode = "attention_shaped_target" - cfg_train.model.task.entropy_reg.weight = 0.01 if use_pretrained: cfg_train.model.init_from_ckpt = _make_pretrain_ckpt(tmp_path) cfg_train.model.init_strict = False @@ -123,7 +132,7 @@ def test_training_regimes_fast_dev_run( def test_runtime_shadow_path_with_real_datamodule(cfg_train): root = Path(rootutils.find_root(indicator=".project-root")) with open_dict(cfg_train): - cfg_train.data = OmegaConf.load(Path("configs/data/spatial_omics.yaml")) + cfg_train.data = OmegaConf.load(CONFIGS_DIR / "data" / "spatial_omics.yaml") cfg_train.data.raw_manifest_path = str(root / "data" / "dummy" / "manifest.csv") cfg_train.data.processed_dir = str(root / "data" / "dummy" / "processed") cfg_train.data.num_workers = 0 @@ -135,8 +144,8 @@ def test_runtime_shadow_path_with_real_datamodule(cfg_train): cfg_train.task = _load_task_cfg("pretrain_bgrl") cfg_train.model = _load_model_cfg("bgrl_module") - cfg_train.optim = OmegaConf.load(Path("configs/optim/adamw.yaml")) - cfg_train.scheduler = OmegaConf.load(Path("configs/scheduler/cosine_step.yaml")) + cfg_train.optim = OmegaConf.load(CONFIGS_DIR / "optim" / "adamw.yaml") + cfg_train.scheduler = OmegaConf.load(CONFIGS_DIR / "scheduler" / "cosine_step.yaml") cfg_train.model.task = cfg_train.task cfg_train.model.optim = cfg_train.optim cfg_train.model.scheduler = cfg_train.scheduler @@ -154,3 +163,56 @@ def test_runtime_shadow_path_with_real_datamodule(cfg_train): HydraConfig().set_config(cfg_train) metric_dict, _ = train(cfg_train) assert "train/loss" in metric_dict + + +def test_survival_regime_fast_dev_run(cfg_train, tmp_path): + """The survival preset must train end-to-end and log a c-index.""" + 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, + "survival": True, + } + ) + cfg_train.task = _load_task_cfg("finetune_survival") + cfg_train.model = _load_model_cfg("supervised_module") + overrides = _load_task_model_overrides("finetune_survival") + if overrides is not None: + cfg_train.model = OmegaConf.merge(cfg_train.model, overrides) + cfg_train.optim = OmegaConf.load(CONFIGS_DIR / "optim" / "adamw.yaml") + cfg_train.scheduler = OmegaConf.load(CONFIGS_DIR / "scheduler" / "cosine_epoch.yaml") + 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 + cfg_train.trainer.fast_dev_run = True + cfg_train.paths.output_dir = str(tmp_path) + cfg_train.paths.log_dir = str(tmp_path) + + HydraConfig().set_config(cfg_train) + metric_dict, _ = train(cfg_train) + assert metric_dict is not None + + +def test_survival_task_config_validation() -> None: + from grass_mil.models.training.builders import validate_task_config + + validate_task_config( + {"target_type": "survival", "target_columns": ["time", "event"], "loss": "survival_coxsgd"} + ) + + with pytest.raises(ValueError, match="exactly two target columns"): + validate_task_config({"target_type": "survival", "target_columns": ["time"]}) + + with pytest.raises(ValueError, match="survival_coxsgd"): + validate_task_config( + { + "target_type": "survival", + "target_columns": ["time", "event"], + "loss": "categorical_bce", + } + )