From 0418ace2f1d39aee25d44f5363e52e474326aacd Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 5 Mar 2026 16:25:43 +0100 Subject: [PATCH 01/25] Dataset prep documentation --- README.md | 4 ++ docs/dataset_preparation.md | 138 ++++++++++++++++++++++++++++++++++++ 2 files changed, 142 insertions(+) create mode 100644 docs/dataset_preparation.md diff --git a/README.md b/README.md index a732ab8..91e7e96 100644 --- a/README.md +++ b/README.md @@ -67,6 +67,9 @@ 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 @@ -86,6 +89,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) diff --git a/docs/dataset_preparation.md b/docs/dataset_preparation.md new file mode 100644 index 0000000..f83611d --- /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 +python src/train.py \ + 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` From 2df5753ca6f4d7e38ba4d94a86ce59fd3d6c8925 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Wed, 12 Aug 2026 19:16:28 +0200 Subject: [PATCH 02/25] Wire up per-cluster enrichment plugin and multi-attribute scatters These were an uncommitted work-in-progress that broke 5 interpretability tests: build_multi_attribute_scatters coerced every extra column with to_numpy(dtype=float), which raises on `condition` (routinely a string label). Unlike the composite-heatmap call site it was not wrapped, so the whole report build failed. - Add _resolve_color_values: numeric columns pass through, non-numeric ones are factorized with their labels kept as colorbar ticks, so categorical attributes stay plottable instead of raising. - Expose per_cluster_cell_type_enrichment params in the plugins config; the plugin was registered but unreachable by configuration. It stays out of `enabled` to preserve the conservative default. - Cover both with tests, including a string-valued condition column and the noise-cluster skip behaviour. Co-Authored-By: Claude Opus 5 (1M context) --- configs/interpretability/plugins/default.yaml | 12 + src/interpretability/plugins/builtin.py | 78 ++++ .../reporting/plotly_builders.py | 333 ++++++++++++++++++ .../test_pipeline_and_report.py | 43 +++ .../test_plugins_and_tier2.py | 44 +++ 5 files changed, 510 insertions(+) diff --git a/configs/interpretability/plugins/default.yaml b/configs/interpretability/plugins/default.yaml index 3681346..3e2bec3 100644 --- a/configs/interpretability/plugins/default.yaml +++ b/configs/interpretability/plugins/default.yaml @@ -28,6 +28,18 @@ params: threshold_start: 0.0 threshold_stop: 55.0 threshold_count: 500 + per_cluster_cell_type_enrichment: + cell_type_column: cell_type + cluster_column: cluster_label + # Clusters with fewer than 2 instances, and HDBSCAN noise (-1) when + # skip_noise is true, are omitted. + skip_noise: 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 tissue_graph: sample_column: sample_id sample_value: null diff --git a/src/interpretability/plugins/builtin.py b/src/interpretability/plugins/builtin.py index 2bff268..7534685 100644 --- a/src/interpretability/plugins/builtin.py +++ b/src/interpretability/plugins/builtin.py @@ -353,6 +353,83 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: return PluginResult(name=self.name, payload=payload, sections=sections) +@dataclass +class PerClusterCellTypeEnrichmentPlugin(InterpretabilityPlugin): + """Compute cell-type neighborhood enrichment within each cluster separately. + + For each cluster, subsets the node table and spatial table to nodes belonging + to that cluster, then runs neighborhood enrichment with label_column=cell_type. + This answers: "within cluster C, which cell types are spatially co-located?" + """ + + name: str = "per_cluster_cell_type_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 per-cluster cell-type enrichment.") + cell_type_column = str( + params.get("cell_type_column", dataset.cell_type_column or "cell_type") + ) + cluster_column = str(params.get("cluster_column", "cluster_label")) + id_column = str(params.get("id_column", dataset.id_column)) + + table = dataset.instance_table.copy() + table[cluster_column] = np.asarray(context.state["cluster_labels"]) + + unique_clusters = sorted(set(table[cluster_column])) + skip_noise = bool(params.get("skip_noise", True)) + + enrichment_by_cluster: Dict[str, pd.DataFrame] = {} + sections: List[ReportSection] = [] + + for cluster_id in unique_clusters: + if skip_noise and cluster_id == -1: + continue + cluster_mask = table[cluster_column] == cluster_id + cluster_table = table[cluster_mask].copy() + if len(cluster_table) < 2: + continue + + cluster_node_ids = set(cluster_table[id_column].astype(str)) + spatial = dataset.spatial_table.copy() + cluster_spatial = spatial[ + spatial["source_id"].astype(str).isin(cluster_node_ids) + & spatial["target_id"].astype(str).isin(cluster_node_ids) + ] + if cluster_spatial.empty: + continue + + try: + result = run_neighborhood_enrichment( + cluster_table, + cluster_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_cluster[str(cluster_id)] = result.enrichment + sections.append( + ReportSection( + title=f"Cell-Type Enrichment — Cluster {cluster_id}", + description=( + f"Cell-type neighborhood enrichment within cluster {cluster_id}." + ), + tables={"enrichment": result.enrichment}, + ) + ) + except (ValueError, KeyError): + continue + + payload: Dict[str, Any] = {"enrichment_by_cluster": enrichment_by_cluster} + return PluginResult(name=self.name, payload=payload, sections=sections) + + def register_builtin_plugins(registry: PluginRegistry) -> None: registry.register(ClusterProfilesPlugin()) registry.register(AttentionAttributionPlugin()) @@ -360,6 +437,7 @@ def register_builtin_plugins(registry: PluginRegistry) -> None: registry.register(DiffNeighborhoodEnrichmentPlugin()) registry.register(FiltrationCurvesPlugin()) registry.register(TissueGraphPlugin()) + registry.register(PerClusterCellTypeEnrichmentPlugin()) def create_builtin_registry() -> PluginRegistry: diff --git a/src/interpretability/reporting/plotly_builders.py b/src/interpretability/reporting/plotly_builders.py index fc82448..1c4f453 100644 --- a/src/interpretability/reporting/plotly_builders.py +++ b/src/interpretability/reporting/plotly_builders.py @@ -155,6 +155,28 @@ def build_plugin_figures(bundle: ReportBundle) -> Dict[str, "go.Figure"]: template="plotly_white", ) figures[f"{name}_{pair}"] = fig + elif "enrichment_by_cluster" in result.payload and isinstance( + result.payload["enrichment_by_cluster"], dict + ): + for cluster_id, enr in result.payload["enrichment_by_cluster"].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} — Cluster {cluster_id}", + template="plotly_white", + ) + figures[f"{name}_cluster_{_safe_slug(cluster_id)}"] = fig elif "tissue_graph_view" in result.payload and isinstance( result.payload["tissue_graph_view"], TissueGraphView ): @@ -177,12 +199,323 @@ def build_plugin_figures(bundle: ReportBundle) -> Dict[str, "go.Figure"]: 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 cluster_feature_reduction) and the main + reduction (e.g. UMAP), each colored by cluster labels and any extra columns + found in the instance_table. + + Args: + bundle: ReportBundle with reduction and optionally cluster_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) + ) + + noise_mask = labels != -1 if -1 in labels else None + + reductions: List[tuple[str, np.ndarray, str, str]] = [] + if bundle.cluster_feature_reduction is not None: + emb = np.asarray(bundle.cluster_feature_reduction.embedding) + if emb.shape[1] >= 2: + fitted = bundle.cluster_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: + # Cluster 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()} — Cluster Labels", + xaxis_title=x_label, + yaxis_title=y_label, + template="plotly_white", + ) + figures.append((f"{red_name}_cluster_all", fig)) + + # Cluster scatter (no noise) + if noise_mask is not None: + fig = go.Figure( + data=[ + go.Scattergl( + x=emb[noise_mask, 0], + y=emb[noise_mask, 1], + mode="markers", + marker={ + "size": 3, + "color": labels[noise_mask], + "colorscale": "Viridis", + "showscale": True, + }, + ) + ] + ) + fig.update_layout( + title=f"{red_name.upper()} — Cluster Labels (no noise)", + xaxis_title=x_label, + yaxis_title=y_label, + template="plotly_white", + ) + figures.append((f"{red_name}_cluster_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 noise_mask is not None: + fig = go.Figure( + data=[ + go.Scattergl( + x=emb[noise_mask, 0], + y=emb[noise_mask, 1], + mode="markers", + marker={ + "size": 3, + "color": values[noise_mask], + "colorscale": cscale, + "showscale": True, + "colorbar": colorbar, + }, + ) + ] + ) + fig.update_layout( + title=f"{red_name.upper()} — {col} (no noise)", + 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_cluster_heatmap(bundle: ReportBundle) -> "go.Figure": + """Build a composite heatmap with enrichment + predictions + attention lift + abundance. + + Replicates the notebook's 4-panel biomarker summary figure using Plotly subplots. + """ + go = _require_plotly() + from plotly.subplots import make_subplots + + summary = bundle.cluster_summary + if summary is None: + raise ValueError("Cluster summary is missing.") + + enr = summary.enrichment + cluster_ids = [str(i) for i in enr.index] + feature_names = [str(c) for c in enr.columns] + + has_attention = summary.mean_scores is not None and summary.attention_lift_present is not None + + n_cols = 4 if has_attention else 2 + widths = [13, 1, 1, 1] if has_attention else [13, 1] + + fig = make_subplots( + rows=1, + cols=n_cols, + shared_yaxes=True, + column_widths=widths, + horizontal_spacing=0.02, + ) + + # Panel 1: Enrichment heatmap + fig.add_trace( + go.Heatmap( + z=enr.to_numpy(), + x=feature_names, + y=cluster_ids, + colorscale="RdBu", + zmid=0.0, + colorbar={"title": "Enrichment", "x": 0.65, "len": 0.9}, + ), + row=1, + col=1, + ) + + # Panel 2: Mean predictions + if summary.mean_scores is not None: + preds = summary.mean_scores.reindex(enr.index).to_numpy().reshape(-1, 1) + else: + preds = summary.cluster_counts.reindex(enr.index).to_numpy().reshape(-1, 1).astype(float) + fig.add_trace( + go.Heatmap( + z=preds, + x=["Predictions"], + y=cluster_ids, + colorscale="RdBu_r", + colorbar={"title": "Pred", "x": 0.78, "len": 0.9}, + ), + row=1, + col=2, + ) + + if has_attention: + # Panel 3: Attention lift (presence conditioned) + lift = summary.attention_lift_present.reindex(enr.index).to_numpy().reshape(-1, 1) + fig.add_trace( + go.Heatmap( + z=lift, + x=["Attn Lift"], + y=cluster_ids, + colorscale="RdGy", + zmid=1.0, + zmin=0.7, + zmax=1.3, + colorbar={"title": "Attn Lift", "x": 0.90, "len": 0.9}, + ), + row=1, + col=3, + ) + + # Panel 4: Log abundance + counts = summary.cluster_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=cluster_ids, + colorscale="Viridis", + colorbar={"title": "Log N", "x": 1.02, "len": 0.9}, + ), + row=1, + col=4, + ) + + fig.update_layout( + title="Cluster Summary: Enrichment + Predictions + Attention + Abundance", + template="plotly_white", + height=max(400, 50 * len(cluster_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.cluster_summary is not None: figures.append(("cluster_enrichment", build_cluster_enrichment_heatmap(bundle))) + + # Multi-attribute scatters (PCA + UMAP colored by various attributes) + if bundle.reduction is not None or bundle.cluster_feature_reduction is not None: + figures.extend(build_multi_attribute_scatters(bundle)) + + # Composite cluster heatmap (enrichment + predictions + attention + abundance) + if bundle.cluster_summary is not None: + try: + figures.append(("composite_cluster_summary", build_composite_cluster_heatmap(bundle))) + except (ValueError, KeyError): + pass + for name, fig in build_plugin_figures(bundle).items(): figures.append((name, fig)) return figures diff --git a/tests/interpretability/test_pipeline_and_report.py b/tests/interpretability/test_pipeline_and_report.py index 08f3eae..3a33e92 100644 --- a/tests/interpretability/test_pipeline_and_report.py +++ b/tests/interpretability/test_pipeline_and_report.py @@ -398,3 +398,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 src.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..c40e6b6 100644 --- a/tests/interpretability/test_plugins_and_tier2.py +++ b/tests/interpretability/test_plugins_and_tier2.py @@ -840,3 +840,47 @@ def test_filtration_runtime_sanity() -> None: elapsed = time.perf_counter() - t0 assert len(out.curves) > 0 assert elapsed < 8.0 + + +def test_builtin_per_cluster_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_cluster_cell_type_enrichment") + out = plugin.run( + ds, + PluginContext(state={"cluster_labels": labels}), + n_perms=0, + enrichment_mode="obs-exp", + ) + assert out.name == "per_cluster_cell_type_enrichment" + enrichment = out.payload["enrichment_by_cluster"] + # Both clusters 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_cluster_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_cluster_cell_type_enrichment") + + skipped = plugin.run( + ds, + PluginContext(state={"cluster_labels": labels}), + skip_noise=True, + n_perms=0, + ) + assert set(skipped.payload["enrichment_by_cluster"]) == {"1"} + + retained = plugin.run( + ds, + PluginContext(state={"cluster_labels": labels}), + skip_noise=False, + n_perms=0, + ) + assert set(retained.payload["enrichment_by_cluster"]) == {"-1", "1"} From bc58f8c8e1f93bcb373be617f04cd8f557bb777b Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Wed, 12 Aug 2026 19:22:55 +0200 Subject: [PATCH 03/25] Repackage src/ as an installable grass_mil package The tree was importable only as `src.*` from the repository root, and pyproject.toml had no [project] table or build backend, so `grass-mil` could not be installed or depended on. - Move all modules to src/grass_mil/ and rewrite `from src.*` imports plus the `_target_:` strings in the config tree. - Move configs/ into the package. Hydra resolves config_path relative to the decorated function's file, so the tree now resolves identically from a checkout and from an installed wheel, with no importlib.resources shim. - Add [build-system] (hatchling) and [project], with squidpy/scanpy, plotly and playwright demoted from hard requirements to extras. - Namespace the console scripts (grass-mil-train, ...); the previous `train_command`/`eval_command` names would collide on any user's PATH. - Replace the bare rootutils.setup_root calls with _paths.setup_project_root, which tolerates a missing .project-root marker, and let paths.root_dir fall back to the Hydra runtime cwd. An installed wheel has no marker above it. - Point tests at the packaged config tree instead of CWD-relative paths. - Add tests/test_packaging.py and a CI job that pip-installs outside conda and runs the CLI from an unrelated directory, so packaging cannot rot silently. No behaviour change: the same 243 tests pass before and after, plus 9 new packaging tests. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/test.yml | 35 +- Makefile | 2 +- README.md | 56 +- configs/model/attention/gated.yaml | 5 - configs/model/loss/categorical_bce.yaml | 2 - configs/model/loss/categorical_ce.yaml | 2 - configs/model/loss/regression_huber.yaml | 2 - configs/model/loss/regression_mse.yaml | 1 - configs/model/loss/survival_coxsgd.yaml | 3 - docs/capabilities_and_modes.md | 10 +- docs/dataset_preparation.md | 2 +- docs/hyperparameter_reference.md | 734 +++++++++--------- docs/inference_utility_suite.md | 16 +- docs/interpretability_extension_guide.md | 8 +- docs/interpretability_suite.md | 60 +- docs/repository_map.md | 94 +-- docs/workflows_train_infer.md | 72 +- pyproject.toml | 44 ++ setup.py | 19 - src/{ => grass_mil}/__init__.py | 0 src/grass_mil/_paths.py | 29 + .../grass_mil/configs}/__init__.py | 0 .../grass_mil/configs}/callbacks/default.yaml | 0 .../configs}/callbacks/early_stopping.yaml | 0 .../configs}/callbacks/model_checkpoint.yaml | 0 .../configs}/callbacks/model_summary.yaml | 0 .../grass_mil/configs}/callbacks/none.yaml | 0 .../configs}/callbacks/rich_progress_bar.yaml | 0 .../configs}/data/sampler/identity.yaml | 0 .../configs}/data/sampler/shadow_custom.yaml | 0 .../configs}/data/sampler/shadow_native.yaml | 0 .../configs}/data/spatial_omics.yaml | 2 +- .../grass_mil/configs}/debug/default.yaml | 0 .../grass_mil/configs}/debug/fdr.yaml | 0 .../grass_mil/configs}/debug/limit.yaml | 0 .../grass_mil/configs}/debug/overfit.yaml | 0 .../grass_mil/configs}/debug/profiler.yaml | 0 {configs => src/grass_mil/configs}/eval.yaml | 0 .../grass_mil/configs}/experiment/debug.yaml | 0 .../notebook_data_shadow_custom.yaml | 0 .../experiment/notebook_model_components.yaml | 0 .../experiment/notebook_training_logic.yaml | 0 .../grass_mil/configs}/extras/default.yaml | 0 .../grass_mil/configs}/hydra/default.yaml | 0 .../inference/aggregation/default.yaml | 0 .../inference/embeddings/default.yaml | 0 .../inference/interpretability/default.yaml | 0 .../configs}/inference/metrics/default.yaml | 0 .../grass_mil/configs}/inference/predict.yaml | 0 .../clustering/agglomerative.yaml | 0 .../interpretability/clustering/dbscan.yaml | 0 .../interpretability/clustering/hdbscan.yaml | 0 .../interpretability/clustering/kmeans.yaml | 0 .../interpretability/clustering/optics.yaml | 0 .../interpretability/clustering/spectral.yaml | 0 .../interpretability/plugins/default.yaml | 0 .../interpretability/reduction/pca.yaml | 0 .../interpretability/reduction/tsne.yaml | 0 .../interpretability/reduction/umap.yaml | 0 .../configs}/interpretability/report.yaml | 0 .../grass_mil/configs}/local/.gitkeep | 0 .../grass_mil/configs}/logger/aim.yaml | 0 .../grass_mil/configs}/logger/comet.yaml | 0 .../grass_mil/configs}/logger/csv.yaml | 0 .../configs}/logger/many_loggers.yaml | 0 .../grass_mil/configs}/logger/mlflow.yaml | 0 .../grass_mil/configs}/logger/neptune.yaml | 0 .../configs}/logger/tensorboard.yaml | 0 .../grass_mil/configs}/logger/wandb.yaml | 0 {configs => src/grass_mil/configs}/loocv.yaml | 0 .../configs/model/attention/gated.yaml | 5 + .../model/attention/gated_projected.yaml | 2 +- .../grass_mil/configs}/model/bgrl_module.yaml | 2 +- .../grass_mil/configs}/model/components.yaml | 10 +- .../grass_mil/configs}/model/encoder/gat.yaml | 2 +- .../grass_mil/configs}/model/encoder/gcn.yaml | 2 +- .../grass_mil/configs}/model/encoder/gin.yaml | 2 +- .../configs}/model/encoder/gine.yaml | 2 +- .../configs}/model/encoder/graphsage.yaml | 2 +- .../grass_mil/configs}/model/heads/graph.yaml | 2 +- .../grass_mil/configs}/model/heads/node.yaml | 2 +- .../configs/model/loss/categorical_bce.yaml | 2 + .../configs/model/loss/categorical_ce.yaml | 2 + .../configs/model/loss/regression_huber.yaml | 2 + .../configs/model/loss/regression_mse.yaml | 1 + .../configs/model/loss/survival_coxsgd.yaml | 3 + .../grass_mil/configs}/model/ssl/bgrl.yaml | 2 +- .../configs}/model/supervised_module.yaml | 2 +- .../grass_mil/configs}/optim/adamw.yaml | 0 .../grass_mil/configs}/paths/default.yaml | 6 +- .../configs}/scheduler/cosine_epoch.yaml | 0 .../configs}/scheduler/cosine_step.yaml | 0 .../configs}/task/finetune_mean.yaml | 0 .../grass_mil/configs}/task/finetune_mil.yaml | 0 .../configs}/task/pretrain_bgrl.yaml | 0 {configs => src/grass_mil/configs}/train.yaml | 0 .../grass_mil/configs}/trainer/cpu.yaml | 0 .../grass_mil/configs}/trainer/ddp.yaml | 0 .../grass_mil/configs}/trainer/ddp_sim.yaml | 0 .../grass_mil/configs}/trainer/default.yaml | 0 .../grass_mil/configs}/trainer/gpu.yaml | 0 .../grass_mil/configs}/trainer/mps.yaml | 0 src/{ => grass_mil}/data/__init__.py | 0 .../data/components/__init__.py | 0 .../data/components/datasets.py | 0 .../data/components/feature_reducers.py | 0 .../data/components/graph_builders.py | 0 .../data/components/loaders.py | 0 .../data/components/patching.py | 0 .../data/components/precompute.py | 0 .../data/components/samplers.py | 0 .../data/components/spatial_types.py | 0 .../data/components/transforms.py | 0 .../data/spatial_omics_datamodule.py | 0 src/{ => grass_mil}/eval.py | 9 +- src/grass_mil/inference/__init__.py | 11 + src/{ => grass_mil}/inference/aggregation.py | 2 +- src/{ => grass_mil}/inference/collectors.py | 2 +- src/{ => grass_mil}/inference/embeddings.py | 4 +- .../inference/interpretability_export.py | 2 +- src/{ => grass_mil}/inference/io.py | 2 +- src/{ => grass_mil}/inference/metrics.py | 2 +- src/{ => grass_mil}/inference/predict.py | 21 +- src/{ => grass_mil}/inference/schemas.py | 0 .../interpretability/__init__.py | 18 +- .../interpretability/contracts.py | 0 .../interpretability/core/__init__.py | 13 +- .../interpretability/core/biomarkers.py | 2 +- .../interpretability/core/clustering.py | 2 +- .../interpretability/core/data.py | 2 +- .../interpretability/core/reduction.py | 2 +- .../interpretability/core/transfer.py | 4 +- .../interpretability/pipeline.py | 16 +- .../interpretability/plugins/__init__.py | 15 + .../interpretability/plugins/base.py | 2 +- .../interpretability/plugins/builtin.py | 14 +- .../interpretability/plugins/registry.py | 2 +- .../interpretability/report_cli.py | 13 +- .../interpretability/reporting/__init__.py | 3 + .../interpretability/reporting/html.py | 2 +- .../interpretability/reporting/pdf.py | 0 .../reporting/plotly_builders.py | 7 +- .../interpretability/reporting/render.py | 12 +- .../interpretability/reporting/sections.py | 2 +- .../interpretability/reporting/snapshot.py | 0 .../interpretability/tier2/__init__.py | 6 +- .../interpretability/tier2/filtration.py | 2 +- .../interpretability/tier2/neighborhood.py | 2 +- .../interpretability/tier2/tissue_graph.py | 0 src/{ => grass_mil}/loocv.py | 13 +- src/{ => grass_mil}/models/__init__.py | 0 src/{ => grass_mil}/models/bgrl_module.py | 0 .../models/components/__init__.py | 0 .../models/components/attention.py | 0 .../models/components/backbones.py | 4 +- .../models/components/factory.py | 0 .../models/components/heads.py | 0 .../models/components/losses.py | 0 .../models/components/pooling.py | 4 +- src/{ => grass_mil}/models/components/ssl.py | 0 .../models/supervised_module.py | 6 +- .../models/training/__init__.py | 0 .../models/training/bagging.py | 0 .../models/training/builders.py | 2 +- .../models/training/checkpoint_init.py | 0 .../models/training/loss_utils.py | 0 .../models/training/optimization.py | 0 .../models/training/ssl_runtime.py | 0 src/{ => grass_mil}/train.py | 9 +- src/grass_mil/utils/__init__.py | 5 + src/{ => grass_mil}/utils/instantiators.py | 2 +- src/{ => grass_mil}/utils/logging_utils.py | 2 +- src/{ => grass_mil}/utils/pylogger.py | 0 src/{ => grass_mil}/utils/rich_utils.py | 2 +- src/{ => grass_mil}/utils/utils.py | 2 +- src/inference/__init__.py | 11 - src/interpretability/plugins/__init__.py | 12 - src/interpretability/reporting/__init__.py | 3 - src/utils/__init__.py | 5 - tests/configs/test_component_instantiation.py | 22 +- tests/conftest.py | 10 +- .../data/test_datamodule_sampler_exposure.py | 4 +- tests/data/test_loader_flags.py | 10 +- .../data/test_precompute_manifest_regions.py | 2 +- tests/data/test_precompute_reducer_modes.py | 12 +- tests/data/test_samplers.py | 40 +- tests/eval/test_aggregation.py | 4 +- tests/eval/test_collectors.py | 2 +- tests/eval/test_embeddings.py | 4 +- tests/eval/test_interpretability_export.py | 4 +- tests/eval/test_io.py | 2 +- tests/eval/test_metrics.py | 4 +- tests/eval/test_predict_cli.py | 16 +- tests/helpers/config_paths.py | 12 + tests/interpretability/test_core.py | 6 +- .../test_pipeline_and_report.py | 34 +- .../test_plugins_and_tier2.py | 18 +- tests/interpretability/test_tissue_graph.py | 2 +- tests/interpretability/test_transfer.py | 10 +- tests/models/test_attention.py | 6 +- tests/models/test_backbones.py | 6 +- tests/models/test_losses.py | 10 +- tests/models/test_optimization.py | 6 +- tests/models/test_runtime.py | 50 +- tests/models/test_ssl.py | 16 +- tests/test_eval.py | 4 +- tests/test_loocv.py | 24 +- tests/test_packaging.py | 68 ++ tests/test_spatial_omics_datamodule.py | 32 +- tests/test_train.py | 2 +- tests/test_train_regimes.py | 19 +- 211 files changed, 1102 insertions(+), 888 deletions(-) delete mode 100644 configs/model/attention/gated.yaml delete mode 100644 configs/model/loss/categorical_bce.yaml delete mode 100644 configs/model/loss/categorical_ce.yaml delete mode 100644 configs/model/loss/regression_huber.yaml delete mode 100644 configs/model/loss/regression_mse.yaml delete mode 100644 configs/model/loss/survival_coxsgd.yaml delete mode 100644 setup.py rename src/{ => grass_mil}/__init__.py (100%) create mode 100644 src/grass_mil/_paths.py rename {configs => src/grass_mil/configs}/__init__.py (100%) rename {configs => src/grass_mil/configs}/callbacks/default.yaml (100%) rename {configs => src/grass_mil/configs}/callbacks/early_stopping.yaml (100%) rename {configs => src/grass_mil/configs}/callbacks/model_checkpoint.yaml (100%) rename {configs => src/grass_mil/configs}/callbacks/model_summary.yaml (100%) rename {configs => src/grass_mil/configs}/callbacks/none.yaml (100%) rename {configs => src/grass_mil/configs}/callbacks/rich_progress_bar.yaml (100%) rename {configs => src/grass_mil/configs}/data/sampler/identity.yaml (100%) rename {configs => src/grass_mil/configs}/data/sampler/shadow_custom.yaml (100%) rename {configs => src/grass_mil/configs}/data/sampler/shadow_native.yaml (100%) rename {configs => src/grass_mil/configs}/data/spatial_omics.yaml (96%) rename {configs => src/grass_mil/configs}/debug/default.yaml (100%) rename {configs => src/grass_mil/configs}/debug/fdr.yaml (100%) rename {configs => src/grass_mil/configs}/debug/limit.yaml (100%) rename {configs => src/grass_mil/configs}/debug/overfit.yaml (100%) rename {configs => src/grass_mil/configs}/debug/profiler.yaml (100%) rename {configs => src/grass_mil/configs}/eval.yaml (100%) rename {configs => src/grass_mil/configs}/experiment/debug.yaml (100%) rename {configs => src/grass_mil/configs}/experiment/notebook_data_shadow_custom.yaml (100%) rename {configs => src/grass_mil/configs}/experiment/notebook_model_components.yaml (100%) rename {configs => src/grass_mil/configs}/experiment/notebook_training_logic.yaml (100%) rename {configs => src/grass_mil/configs}/extras/default.yaml (100%) rename {configs => src/grass_mil/configs}/hydra/default.yaml (100%) rename {configs => src/grass_mil/configs}/inference/aggregation/default.yaml (100%) rename {configs => src/grass_mil/configs}/inference/embeddings/default.yaml (100%) rename {configs => src/grass_mil/configs}/inference/interpretability/default.yaml (100%) rename {configs => src/grass_mil/configs}/inference/metrics/default.yaml (100%) rename {configs => src/grass_mil/configs}/inference/predict.yaml (100%) rename {configs => src/grass_mil/configs}/interpretability/clustering/agglomerative.yaml (100%) rename {configs => src/grass_mil/configs}/interpretability/clustering/dbscan.yaml (100%) rename {configs => src/grass_mil/configs}/interpretability/clustering/hdbscan.yaml (100%) rename {configs => src/grass_mil/configs}/interpretability/clustering/kmeans.yaml (100%) rename {configs => src/grass_mil/configs}/interpretability/clustering/optics.yaml (100%) rename {configs => src/grass_mil/configs}/interpretability/clustering/spectral.yaml (100%) rename {configs => src/grass_mil/configs}/interpretability/plugins/default.yaml (100%) rename {configs => src/grass_mil/configs}/interpretability/reduction/pca.yaml (100%) rename {configs => src/grass_mil/configs}/interpretability/reduction/tsne.yaml (100%) rename {configs => src/grass_mil/configs}/interpretability/reduction/umap.yaml (100%) rename {configs => src/grass_mil/configs}/interpretability/report.yaml (100%) rename {configs => src/grass_mil/configs}/local/.gitkeep (100%) rename {configs => src/grass_mil/configs}/logger/aim.yaml (100%) rename {configs => src/grass_mil/configs}/logger/comet.yaml (100%) rename {configs => src/grass_mil/configs}/logger/csv.yaml (100%) rename {configs => src/grass_mil/configs}/logger/many_loggers.yaml (100%) rename {configs => src/grass_mil/configs}/logger/mlflow.yaml (100%) rename {configs => src/grass_mil/configs}/logger/neptune.yaml (100%) rename {configs => src/grass_mil/configs}/logger/tensorboard.yaml (100%) rename {configs => src/grass_mil/configs}/logger/wandb.yaml (100%) rename {configs => src/grass_mil/configs}/loocv.yaml (100%) create mode 100644 src/grass_mil/configs/model/attention/gated.yaml rename {configs => src/grass_mil/configs}/model/attention/gated_projected.yaml (52%) rename {configs => src/grass_mil/configs}/model/bgrl_module.yaml (90%) rename {configs => src/grass_mil/configs}/model/components.yaml (70%) rename {configs => src/grass_mil/configs}/model/encoder/gat.yaml (78%) rename {configs => src/grass_mil/configs}/model/encoder/gcn.yaml (78%) rename {configs => src/grass_mil/configs}/model/encoder/gin.yaml (78%) rename {configs => src/grass_mil/configs}/model/encoder/gine.yaml (78%) rename {configs => src/grass_mil/configs}/model/encoder/graphsage.yaml (78%) rename {configs => src/grass_mil/configs}/model/heads/graph.yaml (52%) rename {configs => src/grass_mil/configs}/model/heads/node.yaml (53%) create mode 100644 src/grass_mil/configs/model/loss/categorical_bce.yaml create mode 100644 src/grass_mil/configs/model/loss/categorical_ce.yaml create mode 100644 src/grass_mil/configs/model/loss/regression_huber.yaml create mode 100644 src/grass_mil/configs/model/loss/regression_mse.yaml create mode 100644 src/grass_mil/configs/model/loss/survival_coxsgd.yaml rename {configs => src/grass_mil/configs}/model/ssl/bgrl.yaml (58%) rename {configs => src/grass_mil/configs}/model/supervised_module.yaml (93%) rename {configs => src/grass_mil/configs}/optim/adamw.yaml (100%) rename {configs => src/grass_mil/configs}/paths/default.yaml (65%) rename {configs => src/grass_mil/configs}/scheduler/cosine_epoch.yaml (100%) rename {configs => src/grass_mil/configs}/scheduler/cosine_step.yaml (100%) rename {configs => src/grass_mil/configs}/task/finetune_mean.yaml (100%) rename {configs => src/grass_mil/configs}/task/finetune_mil.yaml (100%) rename {configs => src/grass_mil/configs}/task/pretrain_bgrl.yaml (100%) rename {configs => src/grass_mil/configs}/train.yaml (100%) rename {configs => src/grass_mil/configs}/trainer/cpu.yaml (100%) rename {configs => src/grass_mil/configs}/trainer/ddp.yaml (100%) rename {configs => src/grass_mil/configs}/trainer/ddp_sim.yaml (100%) rename {configs => src/grass_mil/configs}/trainer/default.yaml (100%) rename {configs => src/grass_mil/configs}/trainer/gpu.yaml (100%) rename {configs => src/grass_mil/configs}/trainer/mps.yaml (100%) rename src/{ => grass_mil}/data/__init__.py (100%) rename src/{ => grass_mil}/data/components/__init__.py (100%) rename src/{ => grass_mil}/data/components/datasets.py (100%) rename src/{ => grass_mil}/data/components/feature_reducers.py (100%) rename src/{ => grass_mil}/data/components/graph_builders.py (100%) rename src/{ => grass_mil}/data/components/loaders.py (100%) rename src/{ => grass_mil}/data/components/patching.py (100%) rename src/{ => grass_mil}/data/components/precompute.py (100%) rename src/{ => grass_mil}/data/components/samplers.py (100%) rename src/{ => grass_mil}/data/components/spatial_types.py (100%) rename src/{ => grass_mil}/data/components/transforms.py (100%) rename src/{ => grass_mil}/data/spatial_omics_datamodule.py (100%) rename src/{ => grass_mil}/eval.py (90%) create mode 100644 src/grass_mil/inference/__init__.py rename src/{ => grass_mil}/inference/aggregation.py (99%) rename src/{ => grass_mil}/inference/collectors.py (99%) rename src/{ => grass_mil}/inference/embeddings.py (87%) rename src/{ => grass_mil}/inference/interpretability_export.py (99%) rename src/{ => grass_mil}/inference/io.py (97%) rename src/{ => grass_mil}/inference/metrics.py (99%) rename src/{ => grass_mil}/inference/predict.py (97%) rename src/{ => grass_mil}/inference/schemas.py (100%) rename src/{ => grass_mil}/interpretability/__init__.py (60%) rename src/{ => grass_mil}/interpretability/contracts.py (100%) rename src/{ => grass_mil}/interpretability/core/__init__.py (63%) rename src/{ => grass_mil}/interpretability/core/biomarkers.py (98%) rename src/{ => grass_mil}/interpretability/core/clustering.py (97%) rename src/{ => grass_mil}/interpretability/core/data.py (96%) rename src/{ => grass_mil}/interpretability/core/reduction.py (95%) rename src/{ => grass_mil}/interpretability/core/transfer.py (98%) rename src/{ => grass_mil}/interpretability/pipeline.py (93%) create mode 100644 src/grass_mil/interpretability/plugins/__init__.py rename src/{ => grass_mil}/interpretability/plugins/base.py (83%) rename src/{ => grass_mil}/interpretability/plugins/builtin.py (97%) rename src/{ => grass_mil}/interpretability/plugins/registry.py (95%) rename src/{ => grass_mil}/interpretability/report_cli.py (85%) create mode 100644 src/grass_mil/interpretability/reporting/__init__.py rename src/{ => grass_mil}/interpretability/reporting/html.py (97%) rename src/{ => grass_mil}/interpretability/reporting/pdf.py (100%) rename src/{ => grass_mil}/interpretability/reporting/plotly_builders.py (99%) rename src/{ => grass_mil}/interpretability/reporting/render.py (79%) rename src/{ => grass_mil}/interpretability/reporting/sections.py (92%) rename src/{ => grass_mil}/interpretability/reporting/snapshot.py (100%) rename src/{ => grass_mil}/interpretability/tier2/__init__.py (61%) rename src/{ => grass_mil}/interpretability/tier2/filtration.py (98%) rename src/{ => grass_mil}/interpretability/tier2/neighborhood.py (99%) rename src/{ => grass_mil}/interpretability/tier2/tissue_graph.py (100%) rename src/{ => grass_mil}/loocv.py (96%) rename src/{ => grass_mil}/models/__init__.py (100%) rename src/{ => grass_mil}/models/bgrl_module.py (100%) rename src/{ => grass_mil}/models/components/__init__.py (100%) rename src/{ => grass_mil}/models/components/attention.py (100%) rename src/{ => grass_mil}/models/components/backbones.py (98%) rename src/{ => grass_mil}/models/components/factory.py (100%) rename src/{ => grass_mil}/models/components/heads.py (100%) rename src/{ => grass_mil}/models/components/losses.py (100%) rename src/{ => grass_mil}/models/components/pooling.py (92%) rename src/{ => grass_mil}/models/components/ssl.py (100%) rename src/{ => grass_mil}/models/supervised_module.py (99%) rename src/{ => grass_mil}/models/training/__init__.py (100%) rename src/{ => grass_mil}/models/training/bagging.py (100%) rename src/{ => grass_mil}/models/training/builders.py (98%) rename src/{ => grass_mil}/models/training/checkpoint_init.py (100%) rename src/{ => grass_mil}/models/training/loss_utils.py (100%) rename src/{ => grass_mil}/models/training/optimization.py (100%) rename src/{ => grass_mil}/models/training/ssl_runtime.py (100%) rename src/{ => grass_mil}/train.py (95%) create mode 100644 src/grass_mil/utils/__init__.py rename src/{ => grass_mil}/utils/instantiators.py (97%) rename src/{ => grass_mil}/utils/logging_utils.py (97%) rename src/{ => grass_mil}/utils/pylogger.py (100%) rename src/{ => grass_mil}/utils/rich_utils.py (98%) rename src/{ => grass_mil}/utils/utils.py (98%) delete mode 100644 src/inference/__init__.py delete mode 100644 src/interpretability/plugins/__init__.py delete mode 100644 src/interpretability/reporting/__init__.py delete mode 100644 src/utils/__init__.py create mode 100644 tests/helpers/config_paths.py create mode 100644 tests/test_packaging.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index c130688..22caa9b 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -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 + + packaging: + name: Install from source and run the CLI + runs-on: ubuntu-latest + timeout-minutes: 30 + + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + 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..0ce3e7b 100644 --- a/Makefile +++ b/Makefile @@ -27,4 +27,4 @@ test-full: ## Run all tests pytest train: ## Train the model - python src/train.py + grass-mil-train diff --git a/README.md b/README.md index 91e7e96..9341e95 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 @@ -73,13 +87,13 @@ For the full onboarding checklist and dataset schema details, see ### 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 \ @@ -103,10 +117,10 @@ python src/eval.py \ | 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` | Attention-weighted bag logits (+ optional MIL auxiliary terms) | +| 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) @@ -133,20 +147,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 @@ -155,13 +169,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 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/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..9edb981 100644 --- a/docs/capabilities_and_modes.md +++ b/docs/capabilities_and_modes.md @@ -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 index f83611d..9f042c1 100644 --- a/docs/dataset_preparation.md +++ b/docs/dataset_preparation.md @@ -108,7 +108,7 @@ Use a CSV file to attach supervision targets at `sample`, `region`, or `patch` s Use a fresh processed directory when onboarding new data: ```bash -python src/train.py \ +grass-mil-train \ task=finetune_mean \ model=supervised_module \ data=spatial_omics \ diff --git a/docs/hyperparameter_reference.md b/docs/hyperparameter_reference.md index 1abbc74..a6897b7 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_index` | `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,50 @@ 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` | +| `task.node_aux.enabled` | `false` | bool | `true/false` | MIL + binary training stage | `grass-mil-train task.node_aux.enabled=true` | +| `task.node_aux.target_mode` | `attention_shaped_ti` | str | `graph_label`, `attention_shaped_ti` | node auxiliary targets | `grass-mil-train task.node_aux.target_mode=graph_label` | +| `task.node_aux.loss_mode` | `bce` | str | `bce`, `weighted_bce` | node auxiliary loss | `grass-mil-train task.node_aux.loss_mode=weighted_bce` | +| `task.node_aux.weight` | `1.0` | float | any float | node auxiliary weighting | `grass-mil-train task.node_aux.weight=0.2` | +| `task.entropy_reg.enabled` | `false` | bool | `true/false` | MIL + binary training stage | `grass-mil-train task.entropy_reg.enabled=true` | +| `task.entropy_reg.mode` | `attention_shaped_target` | str | `attention`, `attention_shaped_target` | entropy regularization | `grass-mil-train task.entropy_reg.mode=attention` | +| `task.entropy_reg.weight` | `0.01` | float | any float | entropy weighting | `grass-mil-train task.entropy_reg.weight=0.005` | +| `task.loss_weights.region` | `1.0` | float | any float | MIL total loss composition | `grass-mil-train task.loss_weights.region=1.0` | +| `task.loss_weights.node_aux` | `1.0` | float | any float | MIL total loss composition | `grass-mil-train task.loss_weights.node_aux=0.5` | +| `task.loss_weights.entropy` | `0.0` | float | any float | MIL total loss composition | `grass-mil-train task.loss_weights.entropy=0.1` | ## 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 +285,45 @@ 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`,`attention`,`set2set` | graph embedding | `grass-mil-train model.encoder.pooling=set2set` | +| `model.encoder.set2set_steps` | `3` | int | `>=1` | `pooling=set2set` | `grass-mil-train model.encoder.set2set_steps=4` | +| `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 +331,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 +349,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`) @@ -448,20 +448,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 +469,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 +494,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 +551,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 +562,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 +653,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 +677,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 +726,19 @@ 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` | +| `task.node_aux.target_mode` must be `graph_label` or `attention_shaped_ti` | `src/grass_mil/models/training/loss_utils.py` | raises `ValueError` | +| `task.node_aux.loss_mode` must be `bce` or `weighted_bce` | `src/grass_mil/models/training/loss_utils.py` | raises `ValueError` | +| `task.entropy_reg.mode` must be `attention` or `attention_shaped_target` | `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 +750,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..0090e2d 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: @@ -64,7 +64,7 @@ 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: @@ -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..d4a6616 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 @@ -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 ( +from grass_mil.interpretability.core.data import load_interpretability_dataset +from grass_mil.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 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( @@ -374,7 +374,7 @@ transferred = apply_cluster_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, ) @@ -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,14 +438,14 @@ 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]" @@ -454,7 +454,7 @@ python src/interpretability/report_cli.py \ Render a single selected tissue graph colored by cluster 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 diff --git a/docs/repository_map.md b/docs/repository_map.md index a8b19d5..f2309a5 100644 --- a/docs/repository_map.md +++ b/docs/repository_map.md @@ -160,71 +160,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..63e9c66 100644 --- a/docs/workflows_train_infer.md +++ b/docs/workflows_train_infer.md @@ -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,7 +364,7 @@ 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 \ @@ -378,10 +378,10 @@ python src/train.py \ ### 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 ``` @@ -391,7 +391,7 @@ python src/train.py task=finetune_mil \ Region accumulation (manual optimizer stepping over region hyperbatches): ```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 ``` @@ -399,7 +399,7 @@ python src/train.py task=finetune_mil \ Node auxiliary loss: ```bash -python src/train.py task=finetune_mil \ +grass-mil-train task=finetune_mil \ task.node_aux.enabled=true \ task.node_aux.target_mode=attention_shaped_ti \ task.node_aux.loss_mode=weighted_bce \ @@ -409,7 +409,7 @@ python src/train.py task=finetune_mil \ Entropy regularization: ```bash -python src/train.py task=finetune_mil \ +grass-mil-train task=finetune_mil \ task.entropy_reg.enabled=true \ task.entropy_reg.mode=attention \ task.entropy_reg.weight=0.01 @@ -418,7 +418,7 @@ python src/train.py task=finetune_mil \ 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 +428,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 +459,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 +473,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 +578,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 +627,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/pyproject.toml b/pyproject.toml index 36a180f..3b70fcf 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", 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/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/configs/data/sampler/shadow_custom.yaml b/src/grass_mil/configs/data/sampler/shadow_custom.yaml similarity index 100% rename from configs/data/sampler/shadow_custom.yaml rename to src/grass_mil/configs/data/sampler/shadow_custom.yaml diff --git a/configs/data/sampler/shadow_native.yaml b/src/grass_mil/configs/data/sampler/shadow_native.yaml similarity index 100% rename from configs/data/sampler/shadow_native.yaml rename to src/grass_mil/configs/data/sampler/shadow_native.yaml diff --git a/configs/data/spatial_omics.yaml b/src/grass_mil/configs/data/spatial_omics.yaml similarity index 96% rename from configs/data/spatial_omics.yaml rename to src/grass_mil/configs/data/spatial_omics.yaml index a7ae84c..d9fc3f5 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 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 100% rename from configs/inference/interpretability/default.yaml rename to src/grass_mil/configs/inference/interpretability/default.yaml 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/configs/interpretability/plugins/default.yaml b/src/grass_mil/configs/interpretability/plugins/default.yaml similarity index 100% rename from configs/interpretability/plugins/default.yaml rename to src/grass_mil/configs/interpretability/plugins/default.yaml 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..922a481 --- /dev/null +++ b/src/grass_mil/configs/model/attention/gated.yaml @@ -0,0 +1,5 @@ +_target_: grass_mil.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/src/grass_mil/configs/model/attention/gated_projected.yaml similarity index 52% rename from configs/model/attention/gated_projected.yaml rename to src/grass_mil/configs/model/attention/gated_projected.yaml index 76f09b5..38a187a 100644 --- a/configs/model/attention/gated_projected.yaml +++ b/src/grass_mil/configs/model/attention/gated_projected.yaml @@ -1,4 +1,4 @@ -_target_: src.models.components.attention.AttnNetGatedProjected +_target_: grass_mil.models.components.attention.AttnNetGatedProjected input_dim: 128 projection_dim: 64 hidden_dim: 16 diff --git a/configs/model/bgrl_module.yaml b/src/grass_mil/configs/model/bgrl_module.yaml similarity index 90% rename from configs/model/bgrl_module.yaml rename to src/grass_mil/configs/model/bgrl_module.yaml index 4c60dbf..28d9487 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: diff --git a/configs/model/components.yaml b/src/grass_mil/configs/model/components.yaml similarity index 70% rename from configs/model/components.yaml rename to src/grass_mil/configs/model/components.yaml index e2d1729..bad9e15 100644 --- a/configs/model/components.yaml +++ b/src/grass_mil/configs/model/components.yaml @@ -6,7 +6,7 @@ model: use_ssl: false encoder: - _target_: src.models.components.backbones.EncoderConfig + _target_: grass_mil.models.components.backbones.EncoderConfig input_dim: 64 hidden_dim: 128 out_dim: 128 @@ -24,7 +24,7 @@ model: edge_attr_dim: null attention: - _target_: src.models.components.attention.AttnNetGatedProjected + _target_: grass_mil.models.components.attention.AttnNetGatedProjected input_dim: 128 projection_dim: 64 hidden_dim: 16 @@ -34,13 +34,13 @@ model: ssl: 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 graph_head: - _target_: src.models.components.heads.GraphPredictionHead + _target_: grass_mil.models.components.heads.GraphPredictionHead input_dim: 128 output_dim: 1 hidden_dim: 128 @@ -48,7 +48,7 @@ model: dropout: 0.0 node_head: - _target_: src.models.components.heads.NodePredictionHead + _target_: grass_mil.models.components.heads.NodePredictionHead input_dim: 128 output_dim: 1 hidden_dim: 128 diff --git a/configs/model/encoder/gat.yaml b/src/grass_mil/configs/model/encoder/gat.yaml similarity index 78% rename from configs/model/encoder/gat.yaml rename to src/grass_mil/configs/model/encoder/gat.yaml index 5c99954..a74ecd1 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 diff --git a/configs/model/encoder/gcn.yaml b/src/grass_mil/configs/model/encoder/gcn.yaml similarity index 78% rename from configs/model/encoder/gcn.yaml rename to src/grass_mil/configs/model/encoder/gcn.yaml index 9cf83e8..33046bd 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 diff --git a/configs/model/encoder/gin.yaml b/src/grass_mil/configs/model/encoder/gin.yaml similarity index 78% rename from configs/model/encoder/gin.yaml rename to src/grass_mil/configs/model/encoder/gin.yaml index 4b2dab7..9b5f85d 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 diff --git a/configs/model/encoder/gine.yaml b/src/grass_mil/configs/model/encoder/gine.yaml similarity index 78% rename from configs/model/encoder/gine.yaml rename to src/grass_mil/configs/model/encoder/gine.yaml index 932c5c3..3feca52 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 diff --git a/configs/model/encoder/graphsage.yaml b/src/grass_mil/configs/model/encoder/graphsage.yaml similarity index 78% rename from configs/model/encoder/graphsage.yaml rename to src/grass_mil/configs/model/encoder/graphsage.yaml index d5c3110..74d8884 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 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/configs/model/heads/node.yaml b/src/grass_mil/configs/model/heads/node.yaml similarity index 53% rename from configs/model/heads/node.yaml rename to src/grass_mil/configs/model/heads/node.yaml index 9b0be66..0ab830e 100644 --- a/configs/model/heads/node.yaml +++ b/src/grass_mil/configs/model/heads/node.yaml @@ -1,4 +1,4 @@ -_target_: src.models.components.heads.NodePredictionHead +_target_: grass_mil.models.components.heads.NodePredictionHead 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 93% rename from configs/model/supervised_module.yaml rename to src/grass_mil/configs/model/supervised_module.yaml index 10aa191..2e92750 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: 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 100% rename from configs/task/finetune_mil.yaml rename to src/grass_mil/configs/task/finetune_mil.yaml diff --git a/configs/task/pretrain_bgrl.yaml b/src/grass_mil/configs/task/pretrain_bgrl.yaml similarity index 100% rename from configs/task/pretrain_bgrl.yaml rename to src/grass_mil/configs/task/pretrain_bgrl.yaml 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/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 100% rename from src/data/components/datasets.py rename to src/grass_mil/data/components/datasets.py 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 100% rename from src/data/components/precompute.py rename to src/grass_mil/data/components/precompute.py diff --git a/src/data/components/samplers.py b/src/grass_mil/data/components/samplers.py similarity index 100% rename from src/data/components/samplers.py rename to src/grass_mil/data/components/samplers.py 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 100% rename from src/data/components/transforms.py rename to src/grass_mil/data/components/transforms.py 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 99% rename from src/inference/aggregation.py rename to src/grass_mil/inference/aggregation.py index 19e7f99..fffb8e6 100644 --- a/src/inference/aggregation.py +++ b/src/grass_mil/inference/aggregation.py @@ -5,7 +5,7 @@ import torch -from src.inference.schemas import AggregatedPredictionPayload, BatchPredictionPayload +from grass_mil.inference.schemas import AggregatedPredictionPayload, BatchPredictionPayload def _aggregate_rows(values: torch.Tensor, mode: str) -> torch.Tensor: 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 99% rename from src/inference/interpretability_export.py rename to src/grass_mil/inference/interpretability_export.py index 0cf6792..ba9f49e 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]: 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 99% rename from src/inference/metrics.py rename to src/grass_mil/inference/metrics.py index 95f6444..8bd6932 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: diff --git a/src/inference/predict.py b/src/grass_mil/inference/predict.py similarity index 97% rename from src/inference/predict.py rename to src/grass_mil/inference/predict.py index 3b20041..d18d499 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, @@ -465,7 +466,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/interpretability/__init__.py b/src/grass_mil/interpretability/__init__.py similarity index 60% rename from src/interpretability/__init__.py rename to src/grass_mil/interpretability/__init__.py index a736047..1decef7 100644 --- a/src/interpretability/__init__.py +++ b/src/grass_mil/interpretability/__init__.py @@ -1,24 +1,24 @@ -from src.interpretability.core.biomarkers import ( +from grass_mil.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 ( +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_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 ( +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 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/interpretability/contracts.py b/src/grass_mil/interpretability/contracts.py similarity index 100% rename from src/interpretability/contracts.py rename to src/grass_mil/interpretability/contracts.py diff --git a/src/interpretability/core/__init__.py b/src/grass_mil/interpretability/core/__init__.py similarity index 63% rename from src/interpretability/core/__init__.py rename to src/grass_mil/interpretability/core/__init__.py index 8e92430..60c422a 100644 --- a/src/interpretability/core/__init__.py +++ b/src/grass_mil/interpretability/core/__init__.py @@ -1,12 +1,15 @@ -from src.interpretability.core.biomarkers import ( +from grass_mil.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 ( +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_cluster_transfer, fit_cluster_transfer_from_report_bundle, load_cluster_transfer_bundle, diff --git a/src/interpretability/core/biomarkers.py b/src/grass_mil/interpretability/core/biomarkers.py similarity index 98% rename from src/interpretability/core/biomarkers.py rename to src/grass_mil/interpretability/core/biomarkers.py index daede4a..b5c9dcb 100644 --- a/src/interpretability/core/biomarkers.py +++ b/src/grass_mil/interpretability/core/biomarkers.py @@ -3,7 +3,7 @@ import numpy as np import pandas as pd -from src.interpretability.contracts import ClusterSummary +from grass_mil.interpretability.contracts import ClusterSummary def _resolve_variance_estimator(variance_estimator: str) -> int: 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..61ea0f6 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.interpretability.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..013e20b 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.interpretability.contracts import EmbeddingSet, InterpretabilityDataset def load_table(path: Path) -> pd.DataFrame: diff --git a/src/interpretability/core/reduction.py b/src/grass_mil/interpretability/core/reduction.py similarity index 95% rename from src/interpretability/core/reduction.py rename to src/grass_mil/interpretability/core/reduction.py index 88fca8a..6306ff5 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.interpretability.contracts import ReductionResult def run_reduction( diff --git a/src/interpretability/core/transfer.py b/src/grass_mil/interpretability/core/transfer.py similarity index 98% rename from src/interpretability/core/transfer.py rename to src/grass_mil/interpretability/core/transfer.py index 0edf4b3..9d93cde 100644 --- a/src/interpretability/core/transfer.py +++ b/src/grass_mil/interpretability/core/transfer.py @@ -8,12 +8,12 @@ import pandas as pd from sklearn.neighbors import KNeighborsClassifier -from src.interpretability.contracts import ( +from grass_mil.interpretability.contracts import ( ClusterTransferBundle, 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 diff --git a/src/interpretability/pipeline.py b/src/grass_mil/interpretability/pipeline.py similarity index 93% rename from src/interpretability/pipeline.py rename to src/grass_mil/interpretability/pipeline.py index 2e16b93..27af971 100644 --- a/src/interpretability/pipeline.py +++ b/src/grass_mil/interpretability/pipeline.py @@ -7,19 +7,19 @@ import numpy as np import pandas as pd -from src.interpretability.contracts import ( +from grass_mil.interpretability.contracts import ( ClusterSummary, 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 cluster_biomarker_summary +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( 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 83% rename from src/interpretability/plugins/base.py rename to src/grass_mil/interpretability/plugins/base.py index 525524f..04ea2ce 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.interpretability.contracts import InterpretabilityDataset, PluginResult @dataclass diff --git a/src/interpretability/plugins/builtin.py b/src/grass_mil/interpretability/plugins/builtin.py similarity index 97% rename from src/interpretability/plugins/builtin.py rename to src/grass_mil/interpretability/plugins/builtin.py index 7534685..e5e3ed9 100644 --- a/src/interpretability/plugins/builtin.py +++ b/src/grass_mil/interpretability/plugins/builtin.py @@ -6,23 +6,23 @@ import numpy as np import pandas as pd -from src.interpretability.contracts import ( +from grass_mil.interpretability.contracts import ( ClusterSummary, PluginResult, ReportSection, ) -from src.interpretability.core.biomarkers import ( +from grass_mil.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 ( +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.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 prepare_tissue_graph_view +from grass_mil.interpretability.tier2.tissue_graph import prepare_tissue_graph_view def _resolve_filtration_thresholds(params: Dict[str, Any]) -> np.ndarray: 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..cfd2e7c 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.interpretability.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/interpretability/reporting/plotly_builders.py b/src/grass_mil/interpretability/reporting/plotly_builders.py similarity index 99% rename from src/interpretability/reporting/plotly_builders.py rename to src/grass_mil/interpretability/reporting/plotly_builders.py index 1c4f453..7644c5d 100644 --- a/src/interpretability/reporting/plotly_builders.py +++ b/src/grass_mil/interpretability/reporting/plotly_builders.py @@ -5,8 +5,11 @@ 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 +from grass_mil.interpretability.contracts import ReportBundle +from grass_mil.interpretability.tier2.tissue_graph import ( + TissueGraphView, + build_tissue_graph_figure, +) if TYPE_CHECKING: import plotly.graph_objects as go 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..8cce6d2 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.interpretability.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 92% rename from src/interpretability/reporting/sections.py rename to src/grass_mil/interpretability/reporting/sections.py index f83eb93..2f8bab1 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.interpretability.contracts import ReportBundle, ReportSection def build_report_sections(bundle: ReportBundle) -> List[ReportSection]: 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/interpretability/tier2/filtration.py b/src/grass_mil/interpretability/tier2/filtration.py similarity index 98% rename from src/interpretability/tier2/filtration.py rename to src/grass_mil/interpretability/tier2/filtration.py index a2d2761..6130b9b 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.interpretability.contracts import FiltrationCurvesResult def compute_filtration_curves( diff --git a/src/interpretability/tier2/neighborhood.py b/src/grass_mil/interpretability/tier2/neighborhood.py similarity index 99% rename from src/interpretability/tier2/neighborhood.py rename to src/grass_mil/interpretability/tier2/neighborhood.py index c9fd9f5..1c0cdfe 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.interpretability.contracts import NeighborhoodEnrichmentResult def _coerce_edges(spatial_table: pd.DataFrame) -> pd.DataFrame: diff --git a/src/interpretability/tier2/tissue_graph.py b/src/grass_mil/interpretability/tier2/tissue_graph.py similarity index 100% rename from src/interpretability/tier2/tissue_graph.py rename to src/grass_mil/interpretability/tier2/tissue_graph.py 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 100% rename from src/models/__init__.py rename to src/grass_mil/models/__init__.py diff --git a/src/models/bgrl_module.py b/src/grass_mil/models/bgrl_module.py similarity index 100% rename from src/models/bgrl_module.py rename to src/grass_mil/models/bgrl_module.py diff --git a/src/models/components/__init__.py b/src/grass_mil/models/components/__init__.py similarity index 100% rename from src/models/components/__init__.py rename to src/grass_mil/models/components/__init__.py diff --git a/src/models/components/attention.py b/src/grass_mil/models/components/attention.py similarity index 100% rename from src/models/components/attention.py rename to src/grass_mil/models/components/attention.py diff --git a/src/models/components/backbones.py b/src/grass_mil/models/components/backbones.py similarity index 98% rename from src/models/components/backbones.py rename to src/grass_mil/models/components/backbones.py index d7ee95e..a57dbf9 100644 --- a/src/models/components/backbones.py +++ b/src/grass_mil/models/components/backbones.py @@ -12,7 +12,9 @@ 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"] diff --git a/src/models/components/factory.py b/src/grass_mil/models/components/factory.py similarity index 100% rename from src/models/components/factory.py rename to src/grass_mil/models/components/factory.py diff --git a/src/models/components/heads.py b/src/grass_mil/models/components/heads.py similarity index 100% rename from src/models/components/heads.py rename to src/grass_mil/models/components/heads.py 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/models/components/pooling.py b/src/grass_mil/models/components/pooling.py similarity index 92% rename from src/models/components/pooling.py rename to src/grass_mil/models/components/pooling.py index 1fb8769..221daf4 100644 --- a/src/models/components/pooling.py +++ b/src/grass_mil/models/components/pooling.py @@ -13,7 +13,9 @@ global_mean_pool, ) except Exception as exc: # pragma: no cover - raise ImportError("torch_geometric is required for src.models.components.pooling") from exc + raise ImportError( + "torch_geometric is required for grass_mil.models.components.pooling" + ) from exc class GraphPooling(nn.Module): 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 99% rename from src/models/supervised_module.py rename to src/grass_mil/models/supervised_module.py index 2b46678..d9b1b13 100644 --- a/src/models/supervised_module.py +++ b/src/grass_mil/models/supervised_module.py @@ -5,9 +5,9 @@ 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, diff --git a/src/models/training/__init__.py b/src/grass_mil/models/training/__init__.py similarity index 100% rename from src/models/training/__init__.py rename to src/grass_mil/models/training/__init__.py diff --git a/src/models/training/bagging.py b/src/grass_mil/models/training/bagging.py similarity index 100% rename from src/models/training/bagging.py rename to src/grass_mil/models/training/bagging.py diff --git a/src/models/training/builders.py b/src/grass_mil/models/training/builders.py similarity index 98% rename from src/models/training/builders.py rename to src/grass_mil/models/training/builders.py index 1ed51ba..643864b 100644 --- a/src/models/training/builders.py +++ b/src/grass_mil/models/training/builders.py @@ -5,7 +5,7 @@ import lightning as L import torch -from src.models.components import ( +from grass_mil.models.components import ( EncoderConfig, build_attention, build_encoder, diff --git a/src/models/training/checkpoint_init.py b/src/grass_mil/models/training/checkpoint_init.py similarity index 100% rename from src/models/training/checkpoint_init.py rename to src/grass_mil/models/training/checkpoint_init.py diff --git a/src/models/training/loss_utils.py b/src/grass_mil/models/training/loss_utils.py similarity index 100% rename from src/models/training/loss_utils.py rename to src/grass_mil/models/training/loss_utils.py 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/models/training/ssl_runtime.py b/src/grass_mil/models/training/ssl_runtime.py similarity index 100% rename from src/models/training/ssl_runtime.py rename to src/grass_mil/models/training/ssl_runtime.py 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/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/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/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..7a8167a 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", @@ -33,7 +35,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 +44,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 +62,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..9a9ed6f 100644 --- a/tests/data/test_samplers.py +++ b/tests/data/test_samplers.py @@ -48,7 +48,7 @@ def _attach_cell_type_metadata(data, labels): 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( diff --git a/tests/eval/test_aggregation.py b/tests/eval/test_aggregation.py index e0f46c4..8da1548 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: 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..082122c 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: 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..15ea7ea 100644 --- a/tests/eval/test_metrics.py +++ b/tests/eval/test_metrics.py @@ -4,8 +4,8 @@ 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: diff --git a/tests/eval/test_predict_cli.py b/tests/eval/test_predict_cli.py index 97c53ae..c775add 100644 --- a/tests/eval/test_predict_cli.py +++ b/tests/eval/test_predict_cli.py @@ -10,18 +10,18 @@ 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 @@ -186,7 +186,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 +243,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 +345,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 +465,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 +611,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/interpretability/test_core.py b/tests/interpretability/test_core.py index ba65521..dd5fdce 100644 --- a/tests/interpretability/test_core.py +++ b/tests/interpretability/test_core.py @@ -4,12 +4,12 @@ import pandas as pd import pytest -from src.interpretability.core.biomarkers import ( +from grass_mil.interpretability.core.biomarkers import ( cluster_attention_summary, cluster_biomarker_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: diff --git a/tests/interpretability/test_pipeline_and_report.py b/tests/interpretability/test_pipeline_and_report.py index 3a33e92..41157e0 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.interpretability.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]: @@ -80,11 +80,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( @@ -126,7 +126,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 +177,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 +222,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( @@ -299,7 +299,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": { @@ -351,11 +353,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): @@ -383,7 +385,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__ @@ -401,7 +403,7 @@ def _fake_import(name, *args, **kwargs): def test_resolve_color_values_handles_numeric_and_categorical_columns() -> None: - from src.interpretability.reporting.plotly_builders import _resolve_color_values + 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]) diff --git a/tests/interpretability/test_plugins_and_tier2.py b/tests/interpretability/test_plugins_and_tier2.py index c40e6b6..a734915 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.interpretability.contracts import ClusterSummary, 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, @@ -158,7 +158,7 @@ def _should_not_run(*args, **kwargs): # type: ignore[no-untyped-def] raise AssertionError("cluster_biomarker_summary should not be recomputed") monkeypatch.setattr( - "src.interpretability.plugins.builtin.cluster_biomarker_summary", _should_not_run + "grass_mil.interpretability.plugins.builtin.cluster_biomarker_summary", _should_not_run ) plugin = registry.get("cluster_profiles") out = plugin.run( @@ -191,7 +191,7 @@ def test_attention_attribution_reuses_context_base_summary(monkeypatch) -> None: ) monkeypatch.setattr( - "src.interpretability.plugins.builtin.cluster_attention_summary", + "grass_mil.interpretability.plugins.builtin.cluster_attention_summary", lambda *args, **kwargs: attn_summary, # type: ignore[no-untyped-def] ) plugin = registry.get("attention_attribution") @@ -550,7 +550,7 @@ 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( diff --git a/tests/interpretability/test_tissue_graph.py b/tests/interpretability/test_tissue_graph.py index 135b074..ed0ce5c 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, ) diff --git a/tests/interpretability/test_transfer.py b/tests/interpretability/test_transfer.py index ad9927d..cfa1397 100644 --- a/tests/interpretability/test_transfer.py +++ b/tests/interpretability/test_transfer.py @@ -7,9 +7,13 @@ 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 ( +from grass_mil.interpretability.contracts import ( + ClusteringResult, + InterpretabilityDataset, + ReportBundle, +) +from grass_mil.interpretability.core.reduction import run_reduction +from grass_mil.interpretability.core.transfer import ( apply_cluster_transfer, fit_cluster_transfer_from_report_bundle, load_cluster_transfer_bundle, diff --git a/tests/models/test_attention.py b/tests/models/test_attention.py index a8f6265..dfb62e4 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,7 +26,7 @@ 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: 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_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_optimization.py b/tests/models/test_optimization.py index e4e3b86..59156bf 100644 --- a/tests/models/test_optimization.py +++ b/tests/models/test_optimization.py @@ -44,7 +44,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 +86,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 +130,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..ef2f551 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) @@ -238,7 +238,7 @@ def test_attention_bag_aggregation_shapes(): def test_build_mil_aux_targets_attention_shaped(): - from src.models.training.loss_utils import build_mil_aux_targets + from grass_mil.models.training.loss_utils import build_mil_aux_targets bag_targets = torch.tensor([[1.0], [0.0]]) bag_indices = [[0, 1], [2, 3, 4]] @@ -261,7 +261,7 @@ def test_build_mil_aux_targets_attention_shaped(): def test_compute_aux_and_entropy_terms(): - from src.models.training.loss_utils import ( + from grass_mil.models.training.loss_utils import ( compute_aux_node_loss, compute_entropy_regularization, ) @@ -287,7 +287,7 @@ def test_compute_aux_and_entropy_terms(): 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 +299,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 +331,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 +394,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 +454,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 +508,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 +518,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 +528,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 +542,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 +558,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 +585,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), @@ -600,7 +600,7 @@ 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 + from grass_mil.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)} @@ -618,7 +618,7 @@ def test_attention_shaped_aux_targets_are_detached(): 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={}, 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..d66be8b 100644 --- a/tests/test_eval.py +++ b/tests/test_eval.py @@ -5,8 +5,8 @@ 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 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..bb98d73 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), @@ -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), @@ -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), diff --git a/tests/test_train.py b/tests/test_train.py index c8ef99f..595859b 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, diff --git a/tests/test_train_regimes.py b/tests/test_train_regimes.py index 62a4abe..6436988 100644 --- a/tests/test_train_regimes.py +++ b/tests/test_train_regimes.py @@ -4,19 +4,20 @@ 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 _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,11 +74,11 @@ 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) cfg_train.model.encoder.input_dim = 8 cfg_train.model.task = cfg_train.task cfg_train.model.optim = cfg_train.optim @@ -123,7 +124,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 +136,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 From 5f9cd672ce81e6762dd9cc78505c59e6ae9cd719 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Wed, 12 Aug 2026 19:25:59 +0200 Subject: [PATCH 04/25] Stop dropping node-level tensors when sampling, rename categorical codes Two bugs that together removed cell type from every sampled subgraph. Both had to be fixed before the encoder can consume cell-type codes. 1. _apply_transform_to_each_subgraph discarded every node-length tensor as "already represented by x". That is false for categorical_index (the cell-type codes) and pos (coordinates), neither of which is recoverable from x. They are now sliced to the subgraph like any other node attribute. 2. Fixing (1) exposed the reason it had gone unnoticed: PyTorch Geometric's Data.__cat_dim__ returns -1 for any key whose name contains "index", so `categorical_index` was concatenated along the feature dimension when graphs are batched. With equal-sized graphs that silently corrupts the codes; with unequal ones it raises. Renamed to `categorical_codes`, which PyG batches along the node dimension like `pos`. Graphs written by an older precompute are migrated on load, so existing processed caches keep working without a forced re-precompute. Regression tests cover subgraph survival, correct batching, and the migration. Co-Authored-By: Claude Opus 5 (1M context) --- docs/hyperparameter_reference.md | 2 +- docs/model_component_contracts.md | 2 +- src/grass_mil/data/components/datasets.py | 19 +++++- src/grass_mil/data/components/precompute.py | 4 +- src/grass_mil/data/components/samplers.py | 22 ++++--- src/grass_mil/data/components/transforms.py | 6 +- src/grass_mil/models/supervised_module.py | 12 ++-- tests/data/test_samplers.py | 71 ++++++++++++++++++++- tests/test_spatial_omics_datamodule.py | 45 +++++++++++-- 9 files changed, 155 insertions(+), 28 deletions(-) diff --git a/docs/hyperparameter_reference.md b/docs/hyperparameter_reference.md index a6897b7..232f91a 100644 --- a/docs/hyperparameter_reference.md +++ b/docs/hyperparameter_reference.md @@ -82,7 +82,7 @@ grass-mil-train key.path=value | 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` | `grass-mil-train 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 diff --git a/docs/model_component_contracts.md b/docs/model_component_contracts.md index 168821c..a2d5638 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: diff --git a/src/grass_mil/data/components/datasets.py b/src/grass_mil/data/components/datasets.py index b5827c6..ced8651 100644 --- a/src/grass_mil/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/precompute.py b/src/grass_mil/data/components/precompute.py index f6b7108..41285ac 100644 --- a/src/grass_mil/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/grass_mil/data/components/samplers.py b/src/grass_mil/data/components/samplers.py index b5ce44b..6311dd6 100644 --- a/src/grass_mil/data/components/samplers.py +++ b/src/grass_mil/data/components/samplers.py @@ -103,29 +103,29 @@ def from_dict(cls, cfg: Optional[Dict[str, Any]]) -> "RuntimeShadowConfig": def _build_weighted_node_idx( data: Data, *, property_name: str, weight_mode: str, min_weight: float ) -> 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 @@ -813,8 +813,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/transforms.py b/src/grass_mil/data/components/transforms.py index 558e262..b4faf55 100644 --- a/src/grass_mil/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/grass_mil/models/supervised_module.py b/src/grass_mil/models/supervised_module.py index d9b1b13..a65c78f 100644 --- a/src/grass_mil/models/supervised_module.py +++ b/src/grass_mil/models/supervised_module.py @@ -249,17 +249,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) diff --git a/tests/data/test_samplers.py b/tests/data/test_samplers.py index 9a9ed6f..d639c8c 100644 --- a/tests/data/test_samplers.py +++ b/tests/data/test_samplers.py @@ -42,7 +42,7 @@ 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 @@ -545,3 +545,72 @@ 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]), + ) diff --git a/tests/test_spatial_omics_datamodule.py b/tests/test_spatial_omics_datamodule.py index bb98d73..21dbd28 100644 --- a/tests/test_spatial_omics_datamodule.py +++ b/tests/test_spatial_omics_datamodule.py @@ -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()) @@ -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: @@ -1008,3 +1008,40 @@ 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 From 845226273cdbdf1b007e6b2cf9cb16c6d3d4a201 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Wed, 12 Aug 2026 19:34:02 +0200 Subject: [PATCH 05/25] Add cell-type embedding and GINE-on-edge-length to the encoder The encoder could not express the published input scheme. It had no categorical embedding, so cell identity -- the only node input for some cohorts -- had no path into the model at all, and `gine` fed the whole edge_attr vector into the convolution rather than the edge length alone. - NodeInputEmbedding implements h0 = E[t] + Wx, summing a learned cell-type embedding with a projection of the continuous features. It handles input_dim=0, where precompute legitimately emits a zero-column feature matrix and the embedding is the entire input representation. - EncoderConfig.edge_feature_index narrows edge_attr to one scalar column for GINE message passing, matching the current NSCLC/LUAD models. No edge-type embedding: methods.tex still describes the older edge-type convention and is stale on this point. - Out-of-range categorical codes raise instead of being clamped; clamping would silently train on the wrong cell type. - Vocabulary size and column index are inferred from the processed dataset's label map and categorical slices, with explicit config winning -- the same contract as encoder.input_dim. - New gine_length encoder preset (LayerNorm, max pooling) is the production arm; existing configs are unaffected because categorical_embedding defaults to null and reproduces the previous input projection exactly. Covered by unit tests (sum-not-concat, zero-width x, reserved unassigned row, out-of-range rejection, and a bit-identical assertion that non-selected edge columns cannot influence the output) and an end-to-end fast_dev_run on a cell-type-only cohort. Co-Authored-By: Claude Opus 5 (1M context) --- src/grass_mil/configs/model/bgrl_module.yaml | 6 + src/grass_mil/configs/model/components.yaml | 6 + .../configs/model/encoder/gine_length.yaml | 37 +++ .../configs/model/supervised_module.yaml | 6 + src/grass_mil/models/bgrl_module.py | 11 +- src/grass_mil/models/components/__init__.py | 3 + src/grass_mil/models/components/backbones.py | 55 ++++- src/grass_mil/models/components/embeddings.py | 213 ++++++++++++++++++ src/grass_mil/models/supervised_module.py | 2 + src/grass_mil/models/training/__init__.py | 4 + src/grass_mil/models/training/builders.py | 78 ++++++- tests/models/test_node_embeddings.py | 181 +++++++++++++++ tests/test_spatial_omics_datamodule.py | 123 ++++++++++ 13 files changed, 709 insertions(+), 16 deletions(-) create mode 100644 src/grass_mil/configs/model/encoder/gine_length.yaml create mode 100644 src/grass_mil/models/components/embeddings.py create mode 100644 tests/models/test_node_embeddings.py diff --git a/src/grass_mil/configs/model/bgrl_module.yaml b/src/grass_mil/configs/model/bgrl_module.yaml index 28d9487..9a50d40 100644 --- a/src/grass_mil/configs/model/bgrl_module.yaml +++ b/src/grass_mil/configs/model/bgrl_module.yaml @@ -17,6 +17,12 @@ encoder: 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 index bad9e15..1939422 100644 --- a/src/grass_mil/configs/model/components.yaml +++ b/src/grass_mil/configs/model/components.yaml @@ -22,6 +22,12 @@ model: 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 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..150103c --- /dev/null +++ b/src/grass_mil/configs/model/encoder/gine_length.yaml @@ -0,0 +1,37 @@ +# 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 +set2set_steps: 3 +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/src/grass_mil/configs/model/supervised_module.yaml b/src/grass_mil/configs/model/supervised_module.yaml index 2e92750..6df50df 100644 --- a/src/grass_mil/configs/model/supervised_module.yaml +++ b/src/grass_mil/configs/model/supervised_module.yaml @@ -20,6 +20,12 @@ encoder: 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: diff --git a/src/grass_mil/models/bgrl_module.py b/src/grass_mil/models/bgrl_module.py index 6c936e8..bbac5cb 100644 --- a/src/grass_mil/models/bgrl_module.py +++ b/src/grass_mil/models/bgrl_module.py @@ -11,6 +11,7 @@ CosineWarmup, augment_graph, infer_encoder_input_dim, + resolve_encoder_cfg, instantiate_optimizer, instantiate_scheduler_with_warmup, load_state_dict_with_optional_mapping, @@ -56,12 +57,9 @@ def __init__( 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)) self.ssl_model = build_ssl(True, self.encoder, dict(self.hparams.ssl)) if self.ssl_model is None: @@ -97,6 +95,7 @@ def _encode_tuple(self, batch): batch.edge_index, getattr(batch, "edge_attr", None), getattr(batch, "batch", None), + getattr(batch, "categorical_codes", None), ) def _step(self, batch) -> torch.Tensor: diff --git a/src/grass_mil/models/components/__init__.py b/src/grass_mil/models/components/__init__.py index 134b8cb..f06bbb5 100644 --- a/src/grass_mil/models/components/__init__.py +++ b/src/grass_mil/models/components/__init__.py @@ -1,5 +1,6 @@ from .attention import AttnNetGated, AttnNetGatedProjected from .backbones import EncoderConfig, GNNEncoder +from .embeddings import CategoricalEmbeddingConfig, NodeInputEmbedding from .factory import ( build_attention, build_encoder, @@ -25,7 +26,9 @@ "AttnNetGatedProjected", "BGRL", "CoxSGDLoss", + "CategoricalEmbeddingConfig", "EncoderConfig", + "NodeInputEmbedding", "GNNEncoder", "GraphPooling", "GraphPredictionHead", diff --git a/src/grass_mil/models/components/backbones.py b/src/grass_mil/models/components/backbones.py index a57dbf9..532d782 100644 --- a/src/grass_mil/models/components/backbones.py +++ b/src/grass_mil/models/components/backbones.py @@ -7,6 +7,7 @@ import torch from torch import nn +from .embeddings import CategoricalEmbeddingConfig, NodeInputEmbedding from .pooling import GraphPooling try: @@ -39,6 +40,11 @@ class EncoderConfig: 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] = None def _get_activation(name: str) -> nn.Module: @@ -73,10 +79,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)]) @@ -143,6 +149,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}'. " @@ -176,9 +198,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": @@ -197,9 +239,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/grass_mil/models/supervised_module.py b/src/grass_mil/models/supervised_module.py index a65c78f..2b822e3 100644 --- a/src/grass_mil/models/supervised_module.py +++ b/src/grass_mil/models/supervised_module.py @@ -183,6 +183,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) @@ -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( diff --git a/src/grass_mil/models/training/__init__.py b/src/grass_mil/models/training/__init__.py index 5343d9c..e84d850 100644 --- a/src/grass_mil/models/training/__init__.py +++ b/src/grass_mil/models/training/__init__.py @@ -9,7 +9,9 @@ ) from .builders import ( build_supervised_components, + infer_categorical_binding, infer_encoder_input_dim, + resolve_encoder_cfg, validate_task_config, ) from .checkpoint_init import ( @@ -41,7 +43,9 @@ "maybe_sample_indices", "select_target_columns", "build_supervised_components", + "infer_categorical_binding", "infer_encoder_input_dim", + "resolve_encoder_cfg", "validate_task_config", "load_state_dict_with_optional_mapping", "remap_encoder_keys", diff --git a/src/grass_mil/models/training/builders.py b/src/grass_mil/models/training/builders.py index 643864b..34cc45e 100644 --- a/src/grass_mil/models/training/builders.py +++ b/src/grass_mil/models/training/builders.py @@ -40,6 +40,49 @@ def infer_encoder_input_dim(module: L.LightningModule, fallback: int = 0) -> int 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"}: @@ -55,15 +98,42 @@ def validate_task_config(task_cfg: Dict[str, Any]) -> None: ) -def _resolve_encoder_cfg(encoder_cfg: Dict[str, Any], inferred_input_dim: int) -> Dict[str, Any]: +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: + 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." ) - cfg["input_dim"] = inferred_input_dim + # 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: + resolved = dict(categorical_cfg) + label = str(resolved.get("label", "cell_type")) + if module is not None and ( + resolved.get("num_embeddings") is None or resolved.get("column_index") is None + ): + num_embeddings, column_index = infer_categorical_binding(module, label) + if resolved.get("num_embeddings") is None: + resolved["num_embeddings"] = num_embeddings + if resolved.get("column_index") is None: + resolved["column_index"] = column_index + if resolved.get("num_embeddings") is None: + raise ValueError( + f"Could not infer the vocabulary size for categorical label '{label}' from the " + "datamodule. Set model.encoder.categorical_embedding.num_embeddings explicitly, " + "or include the label in data.categorical_features.include_labels." + ) + cfg["categorical_embedding"] = resolved return cfg @@ -79,7 +149,7 @@ def build_supervised_components( 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) + 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) 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/test_spatial_omics_datamodule.py b/tests/test_spatial_omics_datamodule.py index 21dbd28..5e2ce54 100644 --- a/tests/test_spatial_omics_datamodule.py +++ b/tests/test_spatial_omics_datamodule.py @@ -1045,3 +1045,126 @@ def test_graph_dataset_migrates_legacy_categorical_index(tmp_path): 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 From 95b50dc402c258052a0dac9be8c15b0184efeaee Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 09:49:26 +0200 Subject: [PATCH 06/25] Make MIL attention per-class and assert the attribution identity Bag pooling squeezed the attention head to a single channel and broadcast it across every class logit, so the published per-class form could not be expressed. The attribution suite rests on the identity L_c = sum_i A[i,c] * l[i,c] which needs one attention channel and one logit per class; a shared channel cannot express it. - Softmax down the instance axis per class column. Attention width must be 1 (shared, previous behaviour) or one per class; any other width now raises via a shared contract instead of silently broadcasting into the wrong shape. - Fix the shape ripples together, since they failed in opposite ways: inference/aggregation.py flattened attention silently, while interpretability_export.py hard-rejected anything wider than one column. - Export attention_c{j} and the raw attention_logit_c{j} alongside the canonical `attention` column. The raw logits are comparable across bags, which the within-bag normalised scores are not. - validate_attention_width runs at build time too, so a mismatched config fails before the first batch rather than part-way through it. - The width contract lives in a new leaf module (grass_mil.contracts) because models and inference already import each other through the supervised module. BREAKING: task=finetune_mil is now categorical/CE with graph_head.output_dim=2 and attention.n_classes=2. A single-logit BCE head cannot express the attribution identity. finetune_mean is unchanged. Also changed: AggregatedPredictionPayload.bag_attention rows are now [n_instances, n_channels] rather than [n_instances]. Co-Authored-By: Claude Opus 5 (1M context) --- .../inference/interpretability/default.yaml | 5 + src/grass_mil/configs/task/finetune_mil.yaml | 17 ++- src/grass_mil/contracts.py | 27 +++++ src/grass_mil/inference/aggregation.py | 13 ++- .../inference/interpretability_export.py | 42 +++++-- src/grass_mil/inference/predict.py | 1 + src/grass_mil/models/supervised_module.py | 14 ++- src/grass_mil/models/training/bagging.py | 30 ++++- src/grass_mil/models/training/builders.py | 6 + tests/eval/test_aggregation.py | 52 ++++++++- tests/eval/test_interpretability_export.py | 47 +++++++- tests/models/test_bagging.py | 107 ++++++++++++++++++ tests/test_train_regimes.py | 13 +++ 13 files changed, 352 insertions(+), 22 deletions(-) create mode 100644 src/grass_mil/contracts.py create mode 100644 tests/models/test_bagging.py diff --git a/src/grass_mil/configs/inference/interpretability/default.yaml b/src/grass_mil/configs/inference/interpretability/default.yaml index 61c55f8..f34e7a6 100644 --- a/src/grass_mil/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/src/grass_mil/configs/task/finetune_mil.yaml b/src/grass_mil/configs/task/finetune_mil.yaml index 63c8857..2452891 100644 --- a/src/grass_mil/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 @@ -40,3 +44,12 @@ task: 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/contracts.py b/src/grass_mil/contracts.py new file mode 100644 index 0000000..837de9c --- /dev/null +++ b/src/grass_mil/contracts.py @@ -0,0 +1,27 @@ +"""Invariants shared across the model and inference layers. + +Kept dependency-free and outside both subpackages so either side 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 + +__all__ = ["validate_attention_width"] + + +def validate_attention_width(*, attention_width: int, num_classes: int) -> None: + """Reject attention widths that cannot be broadcast against class logits. + + Bag pooling computes ``L_c = sum_i A[i,c] * l[i,c]``. Only two widths are + meaningful: one shared attention channel broadcast across all classes, or + one channel per class (the published form, and the one that makes the + attribution identity exact). Any other width would silently broadcast into + the wrong shape. + """ + if attention_width not in (1, num_classes): + raise ValueError( + f"Attention emits {attention_width} channel(s) but the head emits " + f"{num_classes} class logit(s). Set model.attention.n_classes to 1 " + f"(one shared channel) or to {num_classes} (one per class)." + ) diff --git a/src/grass_mil/inference/aggregation.py b/src/grass_mil/inference/aggregation.py index fffb8e6..7e74ec0 100644 --- a/src/grass_mil/inference/aggregation.py +++ b/src/grass_mil/inference/aggregation.py @@ -6,6 +6,7 @@ import torch from grass_mil.inference.schemas import AggregatedPredictionPayload, BatchPredictionPayload +from grass_mil.contracts import validate_attention_width def _aggregate_rows(values: torch.Tensor, mode: str) -> torch.Tensor: @@ -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/grass_mil/inference/interpretability_export.py b/src/grass_mil/inference/interpretability_export.py index ba9f49e..b2dfd14 100644 --- a/src/grass_mil/inference/interpretability_export.py +++ b/src/grass_mil/inference/interpretability_export.py @@ -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/grass_mil/inference/predict.py b/src/grass_mil/inference/predict.py index d18d499..88f51bb 100644 --- a/src/grass_mil/inference/predict.py +++ b/src/grass_mil/inference/predict.py @@ -422,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") diff --git a/src/grass_mil/models/supervised_module.py b/src/grass_mil/models/supervised_module.py index 2b822e3..2cb6fd7 100644 --- a/src/grass_mil/models/supervised_module.py +++ b/src/grass_mil/models/supervised_module.py @@ -323,6 +323,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, @@ -425,7 +433,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) @@ -460,7 +470,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 diff --git a/src/grass_mil/models/training/bagging.py b/src/grass_mil/models/training/bagging.py index 691077f..d4135d8 100644 --- a/src/grass_mil/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 index 34cc45e..07d0a50 100644 --- a/src/grass_mil/models/training/builders.py +++ b/src/grass_mil/models/training/builders.py @@ -5,6 +5,7 @@ import lightning as L import torch +from grass_mil.contracts import validate_attention_width from grass_mil.models.components import ( EncoderConfig, build_attention, @@ -159,6 +160,11 @@ def build_supervised_components( 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) diff --git a/tests/eval/test_aggregation.py b/tests/eval/test_aggregation.py index 8da1548..b24a2f7 100644 --- a/tests/eval/test_aggregation.py +++ b/tests/eval/test_aggregation.py @@ -71,7 +71,57 @@ 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_channels]; a shared channel keeps one column. + assert out.bag_attention[0].shape == (2, 1) + assert torch.allclose(out.bag_attention[0], torch.tensor([[0.000335], [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_interpretability_export.py b/tests/eval/test_interpretability_export.py index 082122c..46f6ef2 100644 --- a/tests/eval/test_interpretability_export.py +++ b/tests/eval/test_interpretability_export.py @@ -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/models/test_bagging.py b/tests/models/test_bagging.py new file mode 100644 index 0000000..b40f6a4 --- /dev/null +++ b/tests/models/test_bagging.py @@ -0,0 +1,107 @@ +"""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_shared_attention_channel_broadcasts_across_classes() -> None: + """n_classes=1 keeps the previous single-channel behaviour.""" + logits = torch.tensor([[1.0, 10.0], [3.0, 30.0]]) + attn_logits = torch.tensor([[1.0], [9.0]]) + + bag_logits, _, bag_attention, _ = _pool(logits=logits, attn_logits=attn_logits) + + attention = bag_attention["bag"] + assert attention.shape == (2, 1) + expected = (logits * attention).sum(dim=0, keepdim=True) + torch.testing.assert_close(bag_logits, expected, rtol=0, atol=1e-12) + + +def test_pooling_rejects_attention_width_that_is_neither_shared_nor_per_class() -> None: + with pytest.raises(ValueError, match="Attention emits 3 channel"): + _pool(logits=torch.randn(4, 2), attn_logits=torch.randn(4, 3)) + + +@pytest.mark.parametrize("width", [1, 4]) +def test_validate_attention_width_accepts_shared_and_per_class(width: int) -> None: + validate_attention_width(attention_width=width, num_classes=4) + + +@pytest.mark.parametrize("width", [0, 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/test_train_regimes.py b/tests/test_train_regimes.py index 6436988..58c7d8b 100644 --- a/tests/test_train_regimes.py +++ b/tests/test_train_regimes.py @@ -16,6 +16,16 @@ def _load_task_cfg(name: str): return OmegaConf.load(CONFIGS_DIR / "task" / f"{name}.yaml").task +def _load_task_model_overrides(name: str): + """Model overrides a task preset injects into the global package. + + Hydra applies these during composition; this test composes manually, so it + must apply them too or the head width and the task objective disagree. + """ + cfg = OmegaConf.load(CONFIGS_DIR / "task" / f"{name}.yaml") + return cfg.get("model", None) + + def _make_pretrain_ckpt(path: Path) -> str: from grass_mil.models.bgrl_module import BGRLModule @@ -79,6 +89,9 @@ def test_training_regimes_fast_dev_run( "cosine_step.yaml" if task_name == "pretrain_bgrl" else "cosine_epoch.yaml" ) cfg_train.scheduler = OmegaConf.load(CONFIGS_DIR / "scheduler" / scheduler_name) + task_model_overrides = _load_task_model_overrides(task_name) + if task_model_overrides is not None: + cfg_train.model = OmegaConf.merge(cfg_train.model, task_model_overrides) cfg_train.model.encoder.input_dim = 8 cfg_train.model.task = cfg_train.task cfg_train.model.optim = cfg_train.optim From a8d10ed1d4fd922c0c405933b8c9cc389559554f Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 09:58:19 +0200 Subject: [PATCH 07/25] Add hull-interior seed sampling and ego-graph radius cutoff Two sampling controls the clean repo lacked. Both default to off, so existing configs are unaffected. Interior seeds (ported from nsclc_tuning_study/boundary_sampling.py): ego-graphs rooted at the edge of a section are truncated by the section boundary rather than by biology, so they are not comparable to interior ones. The mask excludes only the n_hops graph-hop rim around the convex hull; internal holes and low-cellularity gaps are deliberately not boundary. The mask definition is kept identical to the legacy one so models stay comparable, but the traversal is expressed as sparse frontier expansion instead of a Python BFS. Interior restriction is applied to the candidate pool *before* class weighting, so class frequencies are computed over eligible nodes only. Not ported: the Voronoi-bbox boundary mode (it exists to bridge networkx graphs to relabelled tensors, which is unnecessary here since `pos` is available), the distant-edge fallback, merged_allowed_seed_indices (it serves a merge-all-regions pretraining path this repo does not have), and the NSCLC_BOUNDARY_MODE environment variable, which becomes plain config. Ego radius: r_k = radius_per_hop * depth + radius_offset, applied as a post-hoc subgraph transform through the existing hook. Distances are in datamodule coordinate units, not pixels -- the published 75/55 are pixel values at the NSCLC imaging scale, so the defaults are documented as needing conversion per cohort rather than being silently correct. Kept from the legacy implementation: the empty-interior fallback to all nodes, the largest-interior-component filter, and the atomic np.save/os.replace cache write (a half-written cache file must not be visible to another worker). Co-Authored-By: Claude Opus 5 (1M context) --- .../configs/data/sampler/shadow_custom.yaml | 20 ++ .../configs/data/sampler/shadow_native.yaml | 20 ++ src/grass_mil/configs/data/spatial_omics.yaml | 12 + src/grass_mil/data/components/ego_radius.py | 180 ++++++++++++ src/grass_mil/data/components/samplers.py | 112 ++++++- .../data/components/seed_sampling.py | 277 ++++++++++++++++++ tests/data/test_samplers.py | 106 +++++++ tests/data/test_seed_sampling.py | 223 ++++++++++++++ 8 files changed, 946 insertions(+), 4 deletions(-) create mode 100644 src/grass_mil/data/components/ego_radius.py create mode 100644 src/grass_mil/data/components/seed_sampling.py create mode 100644 tests/data/test_seed_sampling.py diff --git a/src/grass_mil/configs/data/sampler/shadow_custom.yaml b/src/grass_mil/configs/data/sampler/shadow_custom.yaml index 6206abf..8176031 100644 --- a/src/grass_mil/configs/data/sampler/shadow_custom.yaml +++ b/src/grass_mil/configs/data/sampler/shadow_custom.yaml @@ -11,3 +11,23 @@ runtime: 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 index b1df0f3..97971fb 100644 --- a/src/grass_mil/configs/data/sampler/shadow_native.yaml +++ b/src/grass_mil/configs/data/sampler/shadow_native.yaml @@ -11,3 +11,23 @@ runtime: property_name: cell_type weight_mode: inverse min_weight: 1.0e-6 + # Restrict ego-graph seeds to the tissue interior. Ego-graphs rooted at the + # section edge are truncated by the boundary rather than by biology, so they + # are not comparable to interior ones. Excludes only the n_hops graph-hop rim + # around the convex hull; internal holes are not treated as boundary. + interior_seeds: + enabled: false + # null uses `depth`, so the rim matches the ego-graph radius. + n_hops: null + seed_tolerance_multiplier: 1.0 + keep_largest_component: true + # Directory for the on-disk index cache; null disables caching. + cache_dir: null + # Cap the ego-graph's physical extent: r_k = radius_per_hop * depth + offset. + # Distances are in datamodule coordinate units (um after coord_scale_um), NOT + # pixels -- convert the cohort's pixel values before setting these. + ego_radius: + enabled: false + radius_per_hop: 75.0 + radius_offset: 55.0 + keep_root_component: true diff --git a/src/grass_mil/configs/data/spatial_omics.yaml b/src/grass_mil/configs/data/spatial_omics.yaml index d9fc3f5..1ad5f44 100644 --- a/src/grass_mil/configs/data/spatial_omics.yaml +++ b/src/grass_mil/configs/data/spatial_omics.yaml @@ -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/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/grass_mil/data/components/samplers.py b/src/grass_mil/data/components/samplers.py index 6311dd6..5780700 100644 --- a/src/grass_mil/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,11 +111,18 @@ 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_codes") or not hasattr( data, "categorical_slices" @@ -129,6 +150,15 @@ def _build_weighted_node_idx( 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: 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/tests/data/test_samplers.py b/tests/data/test_samplers.py index d639c8c..45f370a 100644 --- a/tests/data/test_samplers.py +++ b/tests/data/test_samplers.py @@ -614,3 +614,109 @@ def test_node_level_codes_batch_along_the_node_dimension(): 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}) From ad9e49bc1956c793d9fe22c2716c177dbcd6a641 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 10:01:53 +0200 Subject: [PATCH 08/25] Port degree-importance BGRL augmentations, fix asymmetric edge dropping The SSL views were uniform elementwise masking, not the production augmentation. - drop_importance: nodes and edges are dropped with probability min((1 - importance) * mu, p_lambda), where importance is the mean-max normalised log-degree, so hubs are preserved and sparse periphery is masked. Dropped nodes have features zeroed and their cell type set to the reserved "unassigned" code rather than being left as a spurious real type. - feature_noise resolves its columns BY NAME against the datamodule's feature names. The legacy default noised a hardcoded index, which pointed at cell size in one cohort and at an unrelated column in another, feeding the encoder pure noise. Two bug classes carried across deliberately: - Edge dropping is symmetric. The legacy finetune transform left force_undirected=False while pretraining used True, so the two regimes augmented differently and message passing became direction-dependent. Both the new augmentations and the existing augment_graph now force it, with a test asserting is_undirected() over several seeds. - Every write clones first. The legacy code wrote through a view, so its unassigned-token fixup silently reverted; a test now asserts the source graph is untouched. Naming corrected without changing behaviour: the legacy `keep_prob` was used as `rand > p`, so it is the drop probability. Deviations recorded in the docstrings: methods.tex describes a min-max degree normalisation, but the final run code uses mean-max and the code is authoritative. Not ported: RewireEdgesAll, AddEdgesByCellType, ShufflePositions, Apoptosis, Mitosis, PhenotypeShift, DropFeaturesGRACE, and the O(E) Python remove_directed_edges (replaced by a vectorised mutual-edge mask). Co-Authored-By: Claude Opus 5 (1M context) --- .../configs/augmentation/bgrl_paper.yaml | 17 + .../configs/augmentation/uniform.yaml | 6 + src/grass_mil/models/bgrl_module.py | 56 +++- src/grass_mil/models/training/__init__.py | 14 + .../models/training/augmentations.py | 293 ++++++++++++++++++ src/grass_mil/models/training/ssl_runtime.py | 25 +- tests/models/test_augmentations.py | 173 +++++++++++ 7 files changed, 571 insertions(+), 13 deletions(-) create mode 100644 src/grass_mil/configs/augmentation/bgrl_paper.yaml create mode 100644 src/grass_mil/configs/augmentation/uniform.yaml create mode 100644 src/grass_mil/models/training/augmentations.py create mode 100644 tests/models/test_augmentations.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/src/grass_mil/models/bgrl_module.py b/src/grass_mil/models/bgrl_module.py index bbac5cb..5aa8cef 100644 --- a/src/grass_mil/models/bgrl_module.py +++ b/src/grass_mil/models/bgrl_module.py @@ -10,6 +10,7 @@ from .training import ( CosineWarmup, augment_graph, + build_augmentation, infer_encoder_input_dim, resolve_encoder_cfg, instantiate_optimizer, @@ -18,6 +19,20 @@ ) +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, @@ -53,6 +68,9 @@ def __init__( 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 def setup(self, stage: Optional[str] = None) -> None: if self._built: @@ -61,6 +79,17 @@ def setup(self, stage: Optional[str] = None) -> None: dict(self.hparams.encoder), infer_encoder_input_dim(self), module=self ) self.encoder = build_encoder(EncoderConfig(**encoder_cfg)) + if self._augmentation_cfg: + input_proj = getattr(self.encoder, "input_proj", None) + categorical = getattr(input_proj, "categorical", None) + self._augmentation = build_augmentation( + self._augmentation_cfg, + 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 + ), + ) 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.") @@ -98,17 +127,24 @@ def _encode_tuple(self, batch): getattr(batch, "categorical_codes", None), ) - 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, + def _make_views(self, batch): + if self._augmentation is not None: + return self._augmentation(batch, None), self._augmentation(batch, None) + return ( + augment_graph( + batch, + drop_edge_p=self._drop_edge_p1, + drop_feat_p=self._drop_feat_p1, + ), + augment_graph( + batch, + drop_edge_p=self._drop_edge_p2, + drop_feat_p=self._drop_feat_p2, + ), ) + + def _step(self, batch) -> torch.Tensor: + view1, view2 = self._make_views(batch) q1, y2 = self.ssl_model(self._encode_tuple(view1), self._encode_tuple(view2)) q2, y1 = self.ssl_model(self._encode_tuple(view2), self._encode_tuple(view1)) loss = ( diff --git a/src/grass_mil/models/training/__init__.py b/src/grass_mil/models/training/__init__.py index e84d850..40b4330 100644 --- a/src/grass_mil/models/training/__init__.py +++ b/src/grass_mil/models/training/__init__.py @@ -7,6 +7,14 @@ 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, @@ -43,6 +51,12 @@ "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", diff --git a/src/grass_mil/models/training/augmentations.py b/src/grass_mil/models/training/augmentations.py new file mode 100644 index 0000000..30392cf --- /dev/null +++ b/src/grass_mil/models/training/augmentations.py @@ -0,0 +1,293 @@ +"""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`` defaults to True: the legacy finetune-time transform + defaulted it to False while BGRL pretraining used True, so the two regimes + silently augmented differently. + """ + from torch_geometric.utils import dropout_edge + + if p <= 0.0: + return data + out = _clone(data) + edge_index, edge_mask = dropout_edge( + out.edge_index, + p=float(p), + force_undirected=force_undirected, + training=True, + ) + out.edge_index = edge_index + edge_attr = getattr(out, "edge_attr", None) + if edge_attr is not None: + if force_undirected: + # dropout_edge returns the mask over the original edges and then + # re-symmetrises, so the attributes are duplicated to match. + kept = edge_attr[edge_mask] + out.edge_attr = torch.cat([kept, kept], dim=0)[: edge_index.size(1)] + else: + out.edge_attr = edge_attr[edge_mask] + 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/grass_mil/models/training/ssl_runtime.py b/src/grass_mil/models/training/ssl_runtime.py index 72ad460..a690e19 100644 --- a/src/grass_mil/models/training/ssl_runtime.py +++ b/src/grass_mil/models/training/ssl_runtime.py @@ -30,14 +30,33 @@ def value(self, step: int) -> float: 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: +def augment_graph( + batch: Data, + *, + drop_edge_p: float, + drop_feat_p: float, + force_undirected: bool = True, +) -> Data: + """Elementwise feature masking plus uniform edge dropping. + + ``force_undirected`` defaults to True. The legacy transform left it False + here while BGRL pretraining used True, so finetuning silently augmented an + asymmetric graph and message passing became direction-dependent. + """ 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) + edge_index, edge_mask = dropout_edge( + aug.edge_index, p=drop_edge_p, force_undirected=force_undirected + ) 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] + kept = aug.edge_attr[edge_mask] + if force_undirected: + # dropout_edge re-symmetrises after masking, so attributes are + # duplicated to match the restored reverse edges. + kept = torch.cat([kept, kept], dim=0)[: edge_index.size(1)] + aug.edge_attr = kept return aug diff --git a/tests/models/test_augmentations.py b/tests/models/test_augmentations.py new file mode 100644 index 0000000..dad3a0c --- /dev/null +++ b/tests/models/test_augmentations.py @@ -0,0 +1,173 @@ +"""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) From 4c6b22bb64e69dde784995995fd7810c77b7c1c5 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 10:04:48 +0200 Subject: [PATCH 09/25] Add a survival task preset and vectorise the concordance index Survival was fully supported by the module and losses but had no task preset, so it could only be reached by hand-overriding four keys, and a mis-shaped config failed deep inside the Cox loss with an opaque shape error. - New task=finetune_survival, documenting that target_columns order is load-bearing ([time, event]) and that attention stays single-channel because a scalar log-hazard has no per-class decomposition. - validate_task_config now rejects survival configs with the wrong number of target columns or a non-Cox loss. - _survival_concordance_index was an O(n^2) Python double loop with two .item() calls per pair, i.e. a device synchronisation per pair. Replaced with a chunked tensor implementation that bounds peak memory at chunk_size * n booleans. No new dependency: lifelines/sksurv would contradict the lean packaging. A 4000-subject cohort now scores in ~0.05s. - Tests compare against a reference double loop over several seeds, check that chunk size cannot change the result, and cover the no-events, all-ties, perfect and reversed orderings. The synthetic datamodule gained a survival mode. Its follow-up times are spread and its events mixed deliberately: with every event at the largest time the partial likelihood is constant and carries no gradient, which is exactly how the first version of this test failed. Co-Authored-By: Claude Opus 5 (1M context) --- .../configs/task/finetune_survival.yaml | 43 +++++++ src/grass_mil/inference/metrics.py | 46 +++++--- src/grass_mil/models/training/builders.py | 14 +++ tests/eval/test_metrics.py | 107 ++++++++++++++++++ tests/helpers/synthetic_datamodule.py | 36 +++++- tests/test_train_regimes.py | 53 +++++++++ 6 files changed, 281 insertions(+), 18 deletions(-) create mode 100644 src/grass_mil/configs/task/finetune_survival.yaml 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/src/grass_mil/inference/metrics.py b/src/grass_mil/inference/metrics.py index 8bd6932..eaee6c2 100644 --- a/src/grass_mil/inference/metrics.py +++ b/src/grass_mil/inference/metrics.py @@ -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/grass_mil/models/training/builders.py b/src/grass_mil/models/training/builders.py index 07d0a50..75bcbe5 100644 --- a/src/grass_mil/models/training/builders.py +++ b/src/grass_mil/models/training/builders.py @@ -97,6 +97,20 @@ def validate_task_config(task_cfg: Dict[str, Any]) -> None: "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( diff --git a/tests/eval/test_metrics.py b/tests/eval/test_metrics.py index 15ea7ea..f8981bd 100644 --- a/tests/eval/test_metrics.py +++ b/tests/eval/test_metrics.py @@ -1,5 +1,8 @@ from __future__ import annotations +import math + +import pytest import torch from omegaconf import OmegaConf from sklearn.metrics import roc_auc_score @@ -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/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/test_train_regimes.py b/tests/test_train_regimes.py index 58c7d8b..8769f36 100644 --- a/tests/test_train_regimes.py +++ b/tests/test_train_regimes.py @@ -168,3 +168,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", + } + ) From 201eb5166494bbc533a267bbe107178fbb6f1ffc Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 10:07:10 +0200 Subject: [PATCH 10/25] Add exact additive attribution with bootstrap intervals The interpretability suite computed an attention lift but no margins, pooled over all bags with no region-equal weighting and no uncertainty. The method's attribution claim rests on the identity L_c = sum_i A[i,c] * l[i,c], which was never actually evaluated. - instance_margin_contributions gives M[i,c] = A[i,c] * (l[i,c] - beta_c), with the head bias removed so the margin reflects the instance-driven part of the decision rather than the head's prior. - cluster_attribution_summary reports signed margin share, |M|-weighted margin, margin sign consistency, abundance-corrected attention lift and prevalence, each with a 200-resample percentile bootstrap interval over regions. Regions are resampled, not instances: region sizes vary by orders of magnitude and pooling over instances lets one large region dominate. - identity_residual = max_r |sum_i M[i,c] - L_c| is reported so the additive assumption is checked rather than trusted, and is covered by a test that deliberately breaks it. One correction to the manuscript's description: the signed margin share is NOT bounded to [-1, 1]. It is a fraction of the bag margin, so a cluster can exceed 1 when another cluster pushes the opposite way. The genuine invariant is that a full partition sums to 1, which is what the test asserts. Regions whose margin is within margin_eps of zero are excluded rather than clipped, since the ratio is meaningless there. Co-Authored-By: Claude Opus 5 (1M context) --- .../interpretability/core/attribution.py | 307 ++++++++++++++++++ tests/interpretability/test_attribution.py | 181 +++++++++++ 2 files changed, 488 insertions(+) create mode 100644 src/grass_mil/interpretability/core/attribution.py create mode 100644 tests/interpretability/test_attribution.py diff --git a/src/grass_mil/interpretability/core/attribution.py b/src/grass_mil/interpretability/core/attribution.py new file mode 100644 index 0000000..710c14e --- /dev/null +++ b/src/grass_mil/interpretability/core/attribution.py @@ -0,0 +1,307 @@ +"""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:`cluster_attribution_summary` reports ``identity_residual`` so the +assumption is checked rather than trusted. + +For binary tasks the reported margin removes the head's per-class bias, so it +reflects the instance-driven part of the decision: + +.. math:: m_i = A_{i,1}(\\ell_{i,1} - \\beta_1) - A_{i,0}(\\ell_{i,0} - \\beta_0) + +Cluster-level summaries aggregate within cluster 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, Optional, Sequence, Tuple + +import numpy as np +import pandas as pd + +__all__ = [ + "AttributionResult", + "cluster_attribution_summary", + "instance_margin_contributions", + "percentile_bootstrap_ci", +] + + +@dataclass(frozen=True) +class AttributionResult: + """Per-instance contributions and their cluster-level summaries. + + Attributes: + per_instance: One row per instance with its contribution columns. + per_cluster: One row per cluster; 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_cluster: 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 _binary_margin(contributions: np.ndarray) -> np.ndarray: + """Class-1-vs-0 margin per instance, or the single column when C == 1.""" + if contributions.shape[1] == 1: + return contributions[:, 0] + return contributions[:, 1] - contributions[:, 0] + + +def cluster_attribution_summary( + table: pd.DataFrame, + cluster_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, +) -> AttributionResult: + """Summarise each cluster's contribution to the bag decisions. + + Args: + table: Instance table; must carry the attention, logit and bag columns. + cluster_labels: One cluster 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. + """ + 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(cluster_labels) != len(table): + raise ValueError( + f"cluster_labels has {len(cluster_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(), + "cluster_label": np.asarray(cluster_labels), + } + ) + for index in range(contributions.shape[1]): + per_instance[f"contribution_c{index}"] = contributions[:, index] + per_instance["margin"] = _binary_margin(contributions) + + identity_residual = _identity_residual( + per_instance[bag_id_column].to_numpy(), raw_contributions, bag_logits + ) + + summary = _summarise_clusters( + per_instance, + attention=attention, + bag_id_column=bag_id_column, + n_bootstrap=n_bootstrap, + random_state=random_state, + margin_eps=margin_eps, + ) + return AttributionResult( + per_instance=per_instance, + per_cluster=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_clusters( + per_instance: pd.DataFrame, + *, + attention: np.ndarray, + bag_id_column: str, + n_bootstrap: int, + random_state: int, + margin_eps: float, +) -> pd.DataFrame: + frame = per_instance.copy() + # A shared attention channel is broadcast; per-class attention is summed to + # one selection weight per instance. + frame["__attention"] = attention.sum(axis=1) if attention.ndim == 2 else attention + + bag_margin = frame.groupby(bag_id_column)["margin"].sum() + bag_size = frame.groupby(bag_id_column).size() + + cell_margin = frame.groupby([bag_id_column, "cluster_label"])["margin"].sum() + cell_count = frame.groupby([bag_id_column, "cluster_label"]).size() + cell_attention = frame.groupby([bag_id_column, "cluster_label"])["__attention"].sum() + + rng = np.random.default_rng(random_state) + rows = [] + for cluster in sorted(frame["cluster_label"].unique()): + cluster_margin = cell_margin.xs(cluster, level="cluster_label") + cluster_count = cell_count.xs(cluster, level="cluster_label") + cluster_attn = cell_attention.xs(cluster, level="cluster_label") + bags = cluster_margin.index + + totals = bag_margin.loc[bags].to_numpy(dtype=float) + margins = cluster_margin.to_numpy(dtype=float) + counts = cluster_count.to_numpy(dtype=float) + sizes = bag_size.loc[bags].to_numpy(dtype=float) + attn = cluster_attn.to_numpy(dtype=float) + + prevalence = counts / np.maximum(sizes, 1.0) + # Abundance-corrected attention lift: 1 is neutral, >1 preferentially + # attended relative to how common the cluster is in that bag. + lift = np.divide(attn, prevalence, out=np.full_like(attn, np.nan), where=prevalence > 0) + + # Shares are fractions of the bag margin and sum to 1 across a full + # partition, but an individual cluster is NOT bounded to [-1, 1]: it can + # exceed 1 when another cluster 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] = {"cluster_label": cluster, "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), + ("attention_lift", lift, 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("cluster_label").sort_index() diff --git a/tests/interpretability/test_attribution.py b/tests/interpretability/test_attribution.py new file mode 100644 index 0000000..ce3be01 --- /dev/null +++ b/tests/interpretability/test_attribution.py @@ -0,0 +1,181 @@ +"""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 cluster-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 ( + cluster_attribution_summary, + instance_margin_contributions, + 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 = cluster_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 = cluster_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 = cluster_attribution_summary( + table, + labels, + attention_columns=["attention_c0", "attention_c1"], + logit_columns=["logit_0", "logit_1"], + n_bootstrap=50, + ) + summary = result.per_cluster + + assert set(summary.index) == {0, 1, 2} + # Note: margin_signed_share is NOT bounded to [-1, 1]. It is a fraction of + # the bag margin, so a cluster pushing hard in one direction can exceed 1 + # when another cluster 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 ("attention_lift", "prevalence"): + assert (summary[f"{column}_lo"] <= summary[column] + 1e-9).all() + assert (summary[column] <= summary[f"{column}_hi"] + 1e-9).all() + + +def test_attention_lift_is_neutral_when_attention_matches_abundance() -> None: + """Uniform attention over a bag gives every cluster a lift of 1.""" + per_bag = 6 + rows = [] + for bag in range(4): + for i in range(per_bag): + rows.append( + { + "bag_id": f"bag_{bag}", + "attention_c0": 1.0 / per_bag, + "logit_0": float(i), + } + ) + table = pd.DataFrame(rows) + labels = np.array([i % 2 for i in range(len(table))]) + + result = cluster_attribution_summary( + table, + labels, + attention_columns=["attention_c0"], + logit_columns=["logit_0"], + n_bootstrap=20, + ) + np.testing.assert_allclose(result.per_cluster["attention_lift"].to_numpy(), 1.0, atol=1e-9) + + +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 = cluster_attribution_summary( + table, + labels, + attention_columns=["attention_c0", "attention_c1"], + logit_columns=["logit_0", "logit_1"], + n_bootstrap=0, + ) + assert result.per_cluster["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) From 46c05870e1ae79f57912d2ba752224dc7ac12a88 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 10:10:19 +0200 Subject: [PATCH 11/25] Add Moran's I, Ripley cross-L, and cross-space cluster agreement The three spatial/geometry analyses the report generator relies on had no equivalent here. - Moran's I over a binary connectivity matrix with a permutation null, BH-FDR adjustment, and the differential-vs-reference z statistic (dI / sqrt(var_k + var_ref)) used to contrast clusters against the noise cluster. - Ripley's centred cross-L via cKDTree pair counts, keeping the self-pair correction (a point is not its own neighbour: the count loses n_a and the denominator is n_a(n_a - 1)) and the region-equal aggregation. - Chance-corrected agreement (ARI/AMI/NMI) plus cluster Jaccard overlap, so "structure the encoder resolves beyond the composition baseline" is a measurement rather than an assertion. Implemented natively over scipy.sparse rather than through squidpy: the package stays installable without scanpy/numba, and tier2.neighborhood already sets that precedent. Two things stated rather than papered over: - SCOPE: the legacy 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. Defensible and arguably cleaner, but not numerically the same statistic. - A latent legacy bug is fixed in the port: the Ripley radius grid was derived from whichever region happened to be iterated first, so results depended on ordering. The grid is now explicit config, defaulting to the median extent, with a test asserting order-independence. Co-Authored-By: Claude Opus 5 (1M context) --- .../interpretability/core/agreement.py | 103 +++++++ .../interpretability/tier2/autocorrelation.py | 245 ++++++++++++++++ .../interpretability/tier2/ripley.py | 227 +++++++++++++++ .../test_spatial_statistics.py | 262 ++++++++++++++++++ 4 files changed, 837 insertions(+) create mode 100644 src/grass_mil/interpretability/core/agreement.py create mode 100644 src/grass_mil/interpretability/tier2/autocorrelation.py create mode 100644 src/grass_mil/interpretability/tier2/ripley.py create mode 100644 tests/interpretability/test_spatial_statistics.py diff --git a/src/grass_mil/interpretability/core/agreement.py b/src/grass_mil/interpretability/core/agreement.py new file mode 100644 index 0000000..062f9d7 --- /dev/null +++ b/src/grass_mil/interpretability/core/agreement.py @@ -0,0 +1,103 @@ +"""Agreement between partitions of the same instances. + +The suite clusters 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__ = ["ClusterAgreementResult", "cluster_jaccard_matrix", "compute_cluster_agreement"] + +_METRICS = ("ari", "ami", "nmi") + + +@dataclass(frozen=True) +class ClusterAgreementResult: + """Pairwise agreement between labelings. + + Attributes: + pairwise_metrics: One row per space pair, columns ``ari``/``ami``/``nmi``. + jaccard: ``(space_a, space_b) -> cluster x cluster`` Jaccard overlap. + """ + + pairwise_metrics: pd.DataFrame + jaccard: Dict[Tuple[str, str], pd.DataFrame] + + +def cluster_jaccard_matrix(labels_a: Sequence, labels_b: Sequence) -> pd.DataFrame: + """Jaccard overlap between every pair of clusters 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_cluster_agreement( + labelings: Mapping[str, Sequence], + *, + metrics: Sequence[str] = _METRICS, + include_jaccard: bool = True, +) -> ClusterAgreementResult: + """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("Cluster 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)] = cluster_jaccard_matrix(a, b) + + frame = pd.DataFrame(rows).set_index(["space_a", "space_b"]) + return ClusterAgreementResult(pairwise_metrics=frame, jaccard=jaccard) diff --git a/src/grass_mil/interpretability/tier2/autocorrelation.py b/src/grass_mil/interpretability/tier2/autocorrelation.py new file mode 100644 index 0000000..d16c876 --- /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 HDBSCAN noise cluster. + """ + 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/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/tests/interpretability/test_spatial_statistics.py b/tests/interpretability/test_spatial_statistics.py new file mode 100644 index 0000000..a0fa32d --- /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 ( + cluster_jaccard_matrix, + compute_cluster_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, "cluster_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="cluster_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"], + "cluster_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="cluster_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["cluster_label"] = nodes["cluster_label"].map({"a": 0, "b": -1}) + result = run_morans_i( + nodes, + edges, + feature_columns=["feature"], + group_column="cluster_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_cluster_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 clusters are numbered. + relabelled = compute_cluster_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_cluster_agreement({"only": [0, 1, 2]}) + with pytest.raises(ValueError, match="same instances"): + compute_cluster_agreement({"a": [0, 1, 2], "b": [0, 1]}) + with pytest.raises(ValueError, match="Unsupported agreement metric"): + compute_cluster_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 = cluster_jaccard_matrix(a, b) + + # Cluster 0 of `a` = {0,1}; cluster 1 of `b` = {1,2,3}; overlap {1}, union 4. + assert matrix.loc[0, 1] == pytest.approx(0.25) + # Cluster 1 of `a` = {2,3} is fully inside cluster 1 of `b`. + assert matrix.loc[1, 1] == pytest.approx(2 / 3) + + with pytest.raises(ValueError, match="must align"): + cluster_jaccard_matrix(np.array([0, 1]), np.array([0, 1, 2])) From 2880f7151b688d382041be4c45d1eadb61f11b0b Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 10:12:51 +0200 Subject: [PATCH 12/25] Wire the new analyses into plugins and expose the attention head space - Attention heads expose head_spaces(), giving a_pre_tanh (the argument of the tanh gate) as a first-class output. The legacy report reached into attention_net.attention_a[0] directly. Pre-tanh matters because saturation collapses exactly the extremes that distinguish strongly-attended instances. With the space exported as table columns, multi-space analysis needs no new pipeline code: re-run the report with a different embedding_prefixes. - New plugins: margin_attribution (the exact additive identity with bootstrap intervals), morans_i (+ differential vs the noise cluster), ripley, and cluster_agreement. - All four are configurable and stay out of the default `enabled` list; a new plugins=full preset turns everything on. margin_attribution is added alongside the existing attention_attribution rather than replacing it. Consolidating them would change the payload keys that build_composite_cluster_heatmap consumes, so that cut is deferred rather than bundled into an already-large change. Co-Authored-By: Claude Opus 5 (1M context) --- .../interpretability/plugins/default.yaml | 33 +++ .../interpretability/plugins/full.yaml | 15 ++ .../interpretability/plugins/builtin.py | 241 ++++++++++++++++++ src/grass_mil/models/components/attention.py | 28 ++ .../test_plugins_and_tier2.py | 102 ++++++++ tests/models/test_attention.py | 33 +++ 6 files changed, 452 insertions(+) create mode 100644 src/grass_mil/configs/interpretability/plugins/full.yaml diff --git a/src/grass_mil/configs/interpretability/plugins/default.yaml b/src/grass_mil/configs/interpretability/plugins/default.yaml index 3e2bec3..7d7b955 100644 --- a/src/grass_mil/configs/interpretability/plugins/default.yaml +++ b/src/grass_mil/configs/interpretability/plugins/default.yaml @@ -40,6 +40,39 @@ params: 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, removed so the margin reflects the + # instance-driven part of the decision. Read it from model_summary.json. + logit_bias: null + n_bootstrap: 200 + random_state: 0 + morans_i: + feature_prefix: comp_ + n_perms: 100 + random_state: 0 + min_nodes: 3 + # Contrast every cluster against this one; -1 is the HDBSCAN noise label. + reference_group: -1 + ripley: + label_column: cluster_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 + cluster_agreement: + active_name: active + # Columns holding partitions from other representation spaces, e.g. joined + # from a previous run's cluster_labels.csv. + label_columns: [] + metrics: [ari, ami, nmi] tissue_graph: sample_column: sample_id sample_value: null 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..266b93b --- /dev/null +++ b/src/grass_mil/configs/interpretability/plugins/full.yaml @@ -0,0 +1,15 @@ +# Every builtin analysis enabled. Slower than the default preset: +# Moran's I and Ripley scale with cluster size and permutation count. +defaults: + - default + - _self_ + +enabled: + - cluster_profiles + - attention_attribution + - margin_attribution + - neighborhood_enrichment + - per_cluster_cell_type_enrichment + - filtration_curves + - morans_i + - ripley diff --git a/src/grass_mil/interpretability/plugins/builtin.py b/src/grass_mil/interpretability/plugins/builtin.py index e5e3ed9..4d073e2 100644 --- a/src/grass_mil/interpretability/plugins/builtin.py +++ b/src/grass_mil/interpretability/plugins/builtin.py @@ -11,17 +11,24 @@ PluginResult, ReportSection, ) +from grass_mil.interpretability.core.agreement import compute_cluster_agreement +from grass_mil.interpretability.core.attribution import cluster_attribution_summary from grass_mil.interpretability.core.biomarkers import ( cluster_attention_summary, cluster_biomarker_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 @@ -430,6 +437,236 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: 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)] + + +@dataclass +class MarginAttributionPlugin(InterpretabilityPlugin): + """Exact additive attribution of bag decisions to instance clusters. + + 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 ["cluster_labels"] + + def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: + table = dataset.instance_table + labels = np.asarray(context.state["cluster_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 = cluster_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=params.get("logit_bias"), + 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)), + ) + payload = { + "per_cluster": result.per_cluster, + "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 " + "cluster with percentile bootstrap intervals over regions." + ), + tables={"per_cluster": result.per_cluster}, + ) + ] + return PluginResult(name=self.name, payload=payload, sections=sections) + + +@dataclass +class MoransIPlugin(InterpretabilityPlugin): + """Spatial autocorrelation per cluster, with a permutation null.""" + + name: str = "morans_i" + + 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("Moran's I requires a spatial table.") + table = dataset.instance_table.copy() + cluster_column = str(params.get("cluster_column", "cluster_label")) + table[cluster_column] = np.asarray(context.state["cluster_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=cluster_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 cluster 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 ["cluster_labels"] + + def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: + table = dataset.instance_table.copy() + label_column = str(params.get("label_column", "cluster_label")) + if label_column == "cluster_label": + table[label_column] = np.asarray(context.state["cluster_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 ClusterAgreementPlugin(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 cluster_labels.csv rather than re-running clustering. + """ + + name: str = "cluster_agreement" + + def required_inputs(self) -> List[str]: + return ["cluster_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["cluster_labels"]) + } + for column in params.get("label_columns", []) or []: + if column not in table.columns: + raise ValueError( + f"cluster_agreement label column {column!r} is not in the instance table." + ) + labelings[str(column)] = table[column].to_numpy() + + if len(labelings) < 2: + raise ValueError( + "cluster_agreement needs at least one additional labeling; set " + "params.cluster_agreement.label_columns." + ) + + result = compute_cluster_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 Cluster 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(ClusterProfilesPlugin()) registry.register(AttentionAttributionPlugin()) @@ -438,6 +675,10 @@ def register_builtin_plugins(registry: PluginRegistry) -> None: registry.register(FiltrationCurvesPlugin()) registry.register(TissueGraphPlugin()) registry.register(PerClusterCellTypeEnrichmentPlugin()) + registry.register(MarginAttributionPlugin()) + registry.register(MoransIPlugin()) + registry.register(RipleyPlugin()) + registry.register(ClusterAgreementPlugin()) def create_builtin_registry() -> PluginRegistry: diff --git a/src/grass_mil/models/components/attention.py b/src/grass_mil/models/components/attention.py index f586965..8266146 100644 --- a/src/grass_mil/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/tests/interpretability/test_plugins_and_tier2.py b/tests/interpretability/test_plugins_and_tier2.py index a734915..340e6d7 100644 --- a/tests/interpretability/test_plugins_and_tier2.py +++ b/tests/interpretability/test_plugins_and_tier2.py @@ -884,3 +884,105 @@ def test_per_cluster_cell_type_enrichment_skips_noise_cluster() -> None: n_perms=0, ) assert set(retained.payload["enrichment_by_cluster"]) == {"-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_cluster_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={"cluster_labels": labels}), n_bootstrap=25 + ) + assert out.name == "margin_attribution" + summary = out.payload["per_cluster"] + assert set(summary.index) == {0, 1} + for column in ("margin_signed_share", "attention_lift", "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={"cluster_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={"cluster_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={"cluster_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("cluster_agreement").run( + ds, + PluginContext(state={"cluster_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("cluster_agreement").run( + _sample_dataset(), PluginContext(state={"cluster_labels": np.zeros(6, dtype=int)}) + ) diff --git a/tests/models/test_attention.py b/tests/models/test_attention.py index dfb62e4..e513d29 100644 --- a/tests/models/test_attention.py +++ b/tests/models/test_attention.py @@ -33,3 +33,36 @@ def test_attention_optional_factory(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 From 4492a3a4e1b411a51ffd6953842447bfbaede3cf Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 10:16:09 +0200 Subject: [PATCH 13/25] Remove the auxiliary objectives and learned graph pooling The goal is region cross-entropy as the sole objective, but the module still carried the machinery for two auxiliary terms that the reported models never used, plus config surface implying they were part of the method. Removed: - task.node_aux (instance-pull) and task.entropy_reg, with build_mil_aux_targets, compute_aux_node_loss, compute_entropy_regularization and gather_instance_logits. - task.loss_weights, which is vacuous once those are gone: _compute_losses now returns the region loss directly instead of composing a weighted sum. - NodePredictionHead, build_node_head and configs/model/heads/node.yaml. It was only ever consumed by node_aux, and was never wired into SupervisedModule at all -- the "node" auxiliary actually used instance-level subgraph logits from the graph head. - GraphPooling's `attention` and `set2set` modes. Instance selection is the MIL head's job, pooling across ego-graphs; duplicating it inside each ego-graph adds parameters without a matching claim. GlobalAttention is also deprecated upstream. supervised_module.py drops from 1318 to ~1100 lines. The advanced-MIL test arm now exercises region accumulation alone, which is retained: it is a throughput mechanism rather than an objective, and removing it would change optimizer-step semantics for streamed bags. Co-Authored-By: Claude Opus 5 (1M context) --- src/grass_mil/configs/model/bgrl_module.yaml | 1 - src/grass_mil/configs/model/components.yaml | 9 - src/grass_mil/configs/model/encoder/gat.yaml | 1 - src/grass_mil/configs/model/encoder/gcn.yaml | 1 - src/grass_mil/configs/model/encoder/gin.yaml | 1 - src/grass_mil/configs/model/encoder/gine.yaml | 1 - .../configs/model/encoder/gine_length.yaml | 1 - .../configs/model/encoder/graphsage.yaml | 1 - src/grass_mil/configs/model/heads/node.yaml | 6 - .../configs/model/supervised_module.yaml | 1 - src/grass_mil/configs/task/finetune_mil.yaml | 17 -- src/grass_mil/models/__init__.py | 4 - src/grass_mil/models/components/__init__.py | 5 +- src/grass_mil/models/components/backbones.py | 2 - src/grass_mil/models/components/factory.py | 6 +- src/grass_mil/models/components/heads.py | 4 - src/grass_mil/models/components/pooling.py | 35 ++-- src/grass_mil/models/supervised_module.py | 156 +----------------- src/grass_mil/models/training/__init__.py | 8 - src/grass_mil/models/training/loss_utils.py | 57 ------- tests/configs/test_component_instantiation.py | 1 - tests/models/test_optimization.py | 1 - tests/models/test_runtime.py | 63 ------- tests/test_train_regimes.py | 9 +- 24 files changed, 29 insertions(+), 362 deletions(-) delete mode 100644 src/grass_mil/configs/model/heads/node.yaml diff --git a/src/grass_mil/configs/model/bgrl_module.yaml b/src/grass_mil/configs/model/bgrl_module.yaml index 9a50d40..1f4c1a8 100644 --- a/src/grass_mil/configs/model/bgrl_module.yaml +++ b/src/grass_mil/configs/model/bgrl_module.yaml @@ -12,7 +12,6 @@ encoder: jk: last act: relu pooling: mean - set2set_steps: 3 gat_heads: 4 use_edge_attr: false edge_weight_index: 0 diff --git a/src/grass_mil/configs/model/components.yaml b/src/grass_mil/configs/model/components.yaml index 1939422..2db45cd 100644 --- a/src/grass_mil/configs/model/components.yaml +++ b/src/grass_mil/configs/model/components.yaml @@ -17,7 +17,6 @@ model: jk: last act: relu pooling: mean - set2set_steps: 3 gat_heads: 4 use_edge_attr: false edge_weight_index: 0 @@ -52,11 +51,3 @@ model: hidden_dim: 128 num_layers: 2 dropout: 0.0 - - node_head: - _target_: grass_mil.models.components.heads.NodePredictionHead - input_dim: 128 - output_dim: 1 - hidden_dim: 128 - num_layers: 2 - dropout: 0.0 diff --git a/src/grass_mil/configs/model/encoder/gat.yaml b/src/grass_mil/configs/model/encoder/gat.yaml index a74ecd1..dd080d4 100644 --- a/src/grass_mil/configs/model/encoder/gat.yaml +++ b/src/grass_mil/configs/model/encoder/gat.yaml @@ -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/src/grass_mil/configs/model/encoder/gcn.yaml b/src/grass_mil/configs/model/encoder/gcn.yaml index 33046bd..75d9b27 100644 --- a/src/grass_mil/configs/model/encoder/gcn.yaml +++ b/src/grass_mil/configs/model/encoder/gcn.yaml @@ -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/gin.yaml b/src/grass_mil/configs/model/encoder/gin.yaml index 9b5f85d..ec63d74 100644 --- a/src/grass_mil/configs/model/encoder/gin.yaml +++ b/src/grass_mil/configs/model/encoder/gin.yaml @@ -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/src/grass_mil/configs/model/encoder/gine.yaml b/src/grass_mil/configs/model/encoder/gine.yaml index 3feca52..6d5546e 100644 --- a/src/grass_mil/configs/model/encoder/gine.yaml +++ b/src/grass_mil/configs/model/encoder/gine.yaml @@ -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 index 150103c..02ec00b 100644 --- a/src/grass_mil/configs/model/encoder/gine_length.yaml +++ b/src/grass_mil/configs/model/encoder/gine_length.yaml @@ -17,7 +17,6 @@ norm: layernorm jk: last act: relu pooling: max -set2set_steps: 3 gat_heads: 4 # Message passing reads exactly one edge column: the edge length. Index 0 of diff --git a/src/grass_mil/configs/model/encoder/graphsage.yaml b/src/grass_mil/configs/model/encoder/graphsage.yaml index 74d8884..88df378 100644 --- a/src/grass_mil/configs/model/encoder/graphsage.yaml +++ b/src/grass_mil/configs/model/encoder/graphsage.yaml @@ -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/src/grass_mil/configs/model/heads/node.yaml b/src/grass_mil/configs/model/heads/node.yaml deleted file mode 100644 index 0ab830e..0000000 --- a/src/grass_mil/configs/model/heads/node.yaml +++ /dev/null @@ -1,6 +0,0 @@ -_target_: grass_mil.models.components.heads.NodePredictionHead -input_dim: 128 -output_dim: 1 -hidden_dim: 128 -num_layers: 2 -dropout: 0.0 diff --git a/src/grass_mil/configs/model/supervised_module.yaml b/src/grass_mil/configs/model/supervised_module.yaml index 6df50df..e973ac3 100644 --- a/src/grass_mil/configs/model/supervised_module.yaml +++ b/src/grass_mil/configs/model/supervised_module.yaml @@ -15,7 +15,6 @@ encoder: jk: last act: relu pooling: mean - set2set_steps: 3 gat_heads: 4 use_edge_attr: false edge_weight_index: 0 diff --git a/src/grass_mil/configs/task/finetune_mil.yaml b/src/grass_mil/configs/task/finetune_mil.yaml index 2452891..e4001c1 100644 --- a/src/grass_mil/configs/task/finetune_mil.yaml +++ b/src/grass_mil/configs/task/finetune_mil.yaml @@ -28,23 +28,6 @@ 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. diff --git a/src/grass_mil/models/__init__.py b/src/grass_mil/models/__init__.py index 697d481..f3fab0b 100644 --- a/src/grass_mil/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/grass_mil/models/components/__init__.py b/src/grass_mil/models/components/__init__.py index f06bbb5..2d89b69 100644 --- a/src/grass_mil/models/components/__init__.py +++ b/src/grass_mil/models/components/__init__.py @@ -6,11 +6,10 @@ 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, @@ -33,7 +32,6 @@ "GraphPooling", "GraphPredictionHead", "MLPPredictor", - "NodePredictionHead", "WeightedBCEWithLogitsLoss", "WeightedCrossEntropyLoss", "WeightedHuberLoss", @@ -42,7 +40,6 @@ "build_encoder", "build_graph_head", "build_loss", - "build_node_head", "build_pooling", "build_ssl", ] diff --git a/src/grass_mil/models/components/backbones.py b/src/grass_mil/models/components/backbones.py index 532d782..1667db8 100644 --- a/src/grass_mil/models/components/backbones.py +++ b/src/grass_mil/models/components/backbones.py @@ -35,7 +35,6 @@ 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 @@ -100,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 diff --git a/src/grass_mil/models/components/factory.py b/src/grass_mil/models/components/factory.py index 8de3e66..21085d7 100644 --- a/src/grass_mil/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/grass_mil/models/components/heads.py b/src/grass_mil/models/components/heads.py index aad991f..8c5aaec 100644 --- a/src/grass_mil/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/grass_mil/models/components/pooling.py b/src/grass_mil/models/components/pooling.py index 221daf4..46b7e9e 100644 --- a/src/grass_mil/models/components/pooling.py +++ b/src/grass_mil/models/components/pooling.py @@ -5,43 +5,40 @@ from torch import nn try: - from torch_geometric.nn import ( - GlobalAttention, - Set2Set, - global_add_pool, - global_max_pool, - global_mean_pool, - ) + 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): - """Thin pooling wrapper with explicit output_dim contract.""" + """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, set2set_steps: int = 3): + 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 - 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}") + raise ValueError( + f"Unsupported graph pooling '{name}'. Expected one of {list(_SUPPORTED)}." + ) def forward(self, x, batch): return self.pool(x, batch) diff --git a/src/grass_mil/models/supervised_module.py b/src/grass_mil/models/supervised_module.py index 2cb6fd7..a59dfcb 100644 --- a/src/grass_mil/models/supervised_module.py +++ b/src/grass_mil/models/supervised_module.py @@ -12,16 +12,12 @@ 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,17 +94,11 @@ 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._manual_optimizer_steps = 0 @@ -637,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, *, @@ -731,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, @@ -756,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 @@ -776,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( @@ -825,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, @@ -837,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, @@ -857,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 # ------------------------------------------------------------------ @@ -886,11 +758,9 @@ def training_step(self, batch, batch_idx): region_losses: list[torch.Tensor] = [] total_losses: list[torch.Tensor] = [] - node_losses: list[torch.Tensor] = [] - entropies: 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], @@ -901,11 +771,7 @@ 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()) + total_losses.append(region_loss) if not total_losses: return torch.zeros((), device=self.device) @@ -933,16 +799,12 @@ def training_step(self, batch, batch_idx): ) mean_total_loss = step_total_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, 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 diff --git a/src/grass_mil/models/training/__init__.py b/src/grass_mil/models/training/__init__.py index 40b4330..4e56a4c 100644 --- a/src/grass_mil/models/training/__init__.py +++ b/src/grass_mil/models/training/__init__.py @@ -27,13 +27,9 @@ 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, @@ -63,13 +59,9 @@ "validate_task_config", "load_state_dict_with_optional_mapping", "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", diff --git a/src/grass_mil/models/training/loss_utils.py b/src/grass_mil/models/training/loss_utils.py index 12e8965..b78882d 100644 --- a/src/grass_mil/models/training/loss_utils.py +++ b/src/grass_mil/models/training/loss_utils.py @@ -37,63 +37,6 @@ def compute_supervised_loss( ) -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) diff --git a/tests/configs/test_component_instantiation.py b/tests/configs/test_component_instantiation.py index 7a8167a..d8682b6 100644 --- a/tests/configs/test_component_instantiation.py +++ b/tests/configs/test_component_instantiation.py @@ -22,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", diff --git a/tests/models/test_optimization.py b/tests/models/test_optimization.py index 59156bf..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, diff --git a/tests/models/test_runtime.py b/tests/models/test_runtime.py index ef2f551..9e8b3bb 100644 --- a/tests/models/test_runtime.py +++ b/tests/models/test_runtime.py @@ -237,53 +237,6 @@ def test_attention_bag_aggregation_shapes(): assert len(bag_indices) == 2 -def test_build_mil_aux_targets_attention_shaped(): - from grass_mil.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 grass_mil.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 @@ -599,22 +552,6 @@ def test_gather_bag_targets_requires_consistent_labels_within_bag(): ) -def test_attention_shaped_aux_targets_are_detached(): - from grass_mil.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") diff --git a/tests/test_train_regimes.py b/tests/test_train_regimes.py index 8769f36..ae54855 100644 --- a/tests/test_train_regimes.py +++ b/tests/test_train_regimes.py @@ -97,16 +97,11 @@ def test_training_regimes_fast_dev_run( 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 From 9506acbd37f2dae386f9136d581f277615784045 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 10:18:00 +0200 Subject: [PATCH 14/25] Document method fidelity and correct the gap audit - New docs/method_fidelity.md maps each part of the method to its implementation and to the test that holds it in place, so a reader can check what is actually enforced before trusting a number. It also records the four places where methods.tex and the run code disagree (edge length vs edge type, mean-max vs min-max degree normalisation, the signed-share bound, and the Moran's I scope), and lists the analyses methods.tex claims that no code implements. - Rewrote docs/legacy_gap_audit.md. It previously asserted that a list of legacy risks had been resolved by rebuilding from concepts, which was misleading: it audited engineering hygiene while several load-bearing pieces of the published method were missing entirely. It is now a status record, including the silently-wrong-answer bugs found while closing those gaps. - Contracts doc gains the encoder, attention and attribution contracts, including the identity sum_i A[i,c]*l[i,c] == L_c as an invariant that bag pooling changes must preserve. - README states the scope decision (the method, not a reproduction harness) instead of listing benchmark packs as a "known gap", and documents the survival preset. Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 23 ++++-- docs/legacy_gap_audit.md | 132 +++++++++++++++++------------- docs/method_fidelity.md | 78 ++++++++++++++++++ docs/model_component_contracts.md | 32 ++++++++ 4 files changed, 202 insertions(+), 63 deletions(-) create mode 100644 docs/method_fidelity.md diff --git a/README.md b/README.md index 9341e95..20926fb 100644 --- a/README.md +++ b/README.md @@ -111,14 +111,16 @@ grass-mil-eval \ - 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/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` | Attention-weighted bag logits (+ optional MIL auxiliary terms) | +| 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_*` | @@ -194,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. Curated first-class hyperparameter sweep recipes in `configs/hparams_search/`. -2. Turnkey dataset-specific benchmark packs beyond the generic Hydra/config contracts. +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. + +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/docs/legacy_gap_audit.md b/docs/legacy_gap_audit.md index 760a0ba..2c4a343 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. + +## Still open + +- `attention_attribution` and `margin_attribution` overlap. Consolidating them + changes the payload keys `build_composite_cluster_heatmap` consumes. +- Region accumulation is retained. It is throughput mechanism rather than + method, but removing it changes optimizer-step semantics for streamed bags. diff --git a/docs/method_fidelity.md b/docs/method_fidelity.md new file mode 100644 index 0000000..3a4a2d7 --- /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** | `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 cluster-level attribution derived from it is invalid. + +| Quantity | Definition | Implementation | Test | +|---|---|---|---| +| Instance contribution | `M[i,c] = A[i,c]·(ℓ[i,c] − βc)` | `interpretability/core/attribution.py` | `test_bias_removal_shifts_the_margin_by_exactly_the_bias_term` | +| Additive identity | `Σᵢ M[i,c] = L_c` | same, `identity_residual` | `test_contributions_sum_exactly_to_the_bag_logit` | +| Attention lift | `λ = (Σ A) / prevalence`, 1 = neutral | same | `test_attention_lift_is_neutral_when_attention_matches_abundance` | +| 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_cluster_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 cluster 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 a2d5638..8d666d1 100644 --- a/docs/model_component_contracts.md +++ b/docs/model_component_contracts.md @@ -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. From 5e70cc78fcb0ab19eb27d922e565ba39c2de9bb3 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 10:26:51 +0200 Subject: [PATCH 15/25] Generalise margin attribution to multi-class with a one-vs-rest contrast The margin was computed as contribution_c1 - contribution_c0 for ANY class count, so with three or more classes every class beyond the first two was silently ignored. - instance_ovr_margins gives m[i,c] = M[i,c] - mean over the other classes. With C=2 this is exactly M[i,1] - M[i,0], so binary results are unchanged; with C=1 (regression, or Cox on the log-hazard scale) the contribution is already the margin. - Cluster summaries are now indexed by (cluster_label, class_index). Binary heads summarise the positive class only, since the class-0 margin is its exact negation; multi-class heads summarise every class. Selectable via focus_classes. - Each class's summary uses that class's own attention channel, so the attention lift reported for class c reflects A[:, c] rather than a pooled weight. Covered by a test where one class attends a single instance and another is uniform. The contributions were already the bias-free attention-weighted logits M[i,c] = A[i,c] * (l[i,c] - beta_c); that is unchanged, and the additive identity is now asserted for a 4-class head as well. Co-Authored-By: Claude Opus 5 (1M context) --- docs/method_fidelity.md | 3 +- .../interpretability/plugins/default.yaml | 4 + .../interpretability/core/attribution.py | 100 +++++++++--- .../interpretability/plugins/builtin.py | 1 + tests/interpretability/test_attribution.py | 147 +++++++++++++++++- .../test_plugins_and_tier2.py | 3 +- 6 files changed, 237 insertions(+), 21 deletions(-) diff --git a/docs/method_fidelity.md b/docs/method_fidelity.md index 3a4a2d7..0fe9207 100644 --- a/docs/method_fidelity.md +++ b/docs/method_fidelity.md @@ -31,7 +31,8 @@ The interpretability claim rests on one identity. If `identity_residual` is not | Quantity | Definition | Implementation | Test | |---|---|---|---| -| Instance contribution | `M[i,c] = A[i,c]·(ℓ[i,c] − βc)` | `interpretability/core/attribution.py` | `test_bias_removal_shifts_the_margin_by_exactly_the_bias_term` | +| 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` | | Attention lift | `λ = (Σ A) / prevalence`, 1 = neutral | same | `test_attention_lift_is_neutral_when_attention_matches_abundance` | | Uncertainty | 200-resample percentile bootstrap **over regions** | `percentile_bootstrap_ci` | `test_bootstrap_ci_brackets_the_point_estimate` | diff --git a/src/grass_mil/configs/interpretability/plugins/default.yaml b/src/grass_mil/configs/interpretability/plugins/default.yaml index 7d7b955..1783506 100644 --- a/src/grass_mil/configs/interpretability/plugins/default.yaml +++ b/src/grass_mil/configs/interpretability/plugins/default.yaml @@ -50,6 +50,10 @@ params: 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 diff --git a/src/grass_mil/interpretability/core/attribution.py b/src/grass_mil/interpretability/core/attribution.py index 710c14e..fbf4da0 100644 --- a/src/grass_mil/interpretability/core/attribution.py +++ b/src/grass_mil/interpretability/core/attribution.py @@ -11,11 +11,14 @@ model. :func:`cluster_attribution_summary` reports ``identity_residual`` so the assumption is checked rather than trusted. -For binary tasks the reported margin removes the head's per-class bias, so it -reflects the instance-driven part of the decision: +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`). + Cluster-level summaries aggregate within cluster 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 @@ -25,7 +28,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Dict, Optional, Sequence, Tuple +from typing import Dict, List, Optional, Sequence, Tuple import numpy as np import pandas as pd @@ -34,6 +37,7 @@ "AttributionResult", "cluster_attribution_summary", "instance_margin_contributions", + "instance_ovr_margins", "percentile_bootstrap_ci", ] @@ -142,11 +146,49 @@ def instance_margin_contributions( return attention * centred -def _binary_margin(contributions: np.ndarray) -> np.ndarray: - """Class-1-vs-0 margin per instance, or the single column when C == 1.""" - if contributions.shape[1] == 1: - return contributions[:, 0] - return contributions[:, 1] - contributions[:, 0] +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 cluster-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 cluster_attribution_summary( @@ -161,6 +203,7 @@ def cluster_attribution_summary( 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 cluster's contribution to the bag decisions. @@ -173,6 +216,9 @@ def cluster_attribution_summary( 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: @@ -193,22 +239,31 @@ def cluster_attribution_summary( "cluster_label": np.asarray(cluster_labels), } ) - for index in range(contributions.shape[1]): + 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["margin"] = _binary_margin(contributions) + per_instance[f"margin_c{index}"] = margins[:, index] identity_residual = _identity_residual( per_instance[bag_id_column].to_numpy(), raw_contributions, bag_logits ) - summary = _summarise_clusters( - per_instance, - attention=attention, - bag_id_column=bag_id_column, - n_bootstrap=n_bootstrap, - random_state=random_state, - margin_eps=margin_eps, - ) + summaries = [] + for class_index in _resolve_focus_classes(n_classes, focus_classes): + frame = per_instance.copy() + frame["margin"] = margins[:, class_index] + part = _summarise_clusters( + frame, + attention=_attention_for_class(attention, class_index), + 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_cluster=summary, @@ -216,6 +271,15 @@ def cluster_attribution_summary( ) +def _attention_for_class(attention: np.ndarray, class_index: int) -> np.ndarray: + """Attention weights driving one class, or the shared channel.""" + if attention.ndim == 1: + return attention[:, None] + if attention.shape[1] == 1: + return attention + return attention[:, class_index : class_index + 1] + + def _identity_residual( bag_ids: np.ndarray, contributions: np.ndarray, diff --git a/src/grass_mil/interpretability/plugins/builtin.py b/src/grass_mil/interpretability/plugins/builtin.py index 4d073e2..859642b 100644 --- a/src/grass_mil/interpretability/plugins/builtin.py +++ b/src/grass_mil/interpretability/plugins/builtin.py @@ -486,6 +486,7 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: 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_cluster": result.per_cluster, diff --git a/tests/interpretability/test_attribution.py b/tests/interpretability/test_attribution.py index ce3be01..30628c5 100644 --- a/tests/interpretability/test_attribution.py +++ b/tests/interpretability/test_attribution.py @@ -14,6 +14,7 @@ from grass_mil.interpretability.core.attribution import ( cluster_attribution_summary, instance_margin_contributions, + instance_ovr_margins, percentile_bootstrap_ci, ) @@ -108,7 +109,9 @@ def test_cluster_summaries_are_bounded_and_carry_intervals() -> None: ) summary = result.per_cluster - assert set(summary.index) == {0, 1, 2} + # Indexed by (cluster_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 cluster pushing hard in one direction can exceed 1 # when another cluster opposes it. The real invariant is that a full @@ -179,3 +182,145 @@ def test_bootstrap_ci_handles_degenerate_inputs() -> None: # 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 = cluster_attribution_summary( + table, + labels, + attention_columns=attention_columns, + logit_columns=logit_columns, + n_bootstrap=20, + ) + # One row per (cluster, class): no class is silently ignored. + assert set(result.per_cluster.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 = cluster_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 = cluster_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_cluster.index} == {2} + + with pytest.raises(ValueError, match="out of range"): + cluster_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_multiclass_attention_lift_uses_the_matching_class_channel() -> None: + """Each class's summary must use that class's own attention column.""" + per_bag = 4 + rows = [] + for bag in range(3): + for i in range(per_bag): + rows.append( + { + "bag_id": f"bag_{bag}", + # Class 0 attends instance 0 almost exclusively; class 1 is uniform. + "attention_c0": 0.97 if i == 0 else 0.01, + "attention_c1": 1.0 / per_bag, + "attention_c2": 1.0 / per_bag, + "logit_0": float(i), + "logit_1": float(i), + "logit_2": float(i), + } + ) + table = pd.DataFrame(rows) + labels = np.array([0 if i % per_bag == 0 else 1 for i in range(len(table))]) + + result = cluster_attribution_summary( + table, + labels, + attention_columns=["attention_c0", "attention_c1", "attention_c2"], + logit_columns=["logit_0", "logit_1", "logit_2"], + n_bootstrap=0, + ) + lift = result.per_cluster["attention_lift"] + # Cluster 0 is the singled-out instance: strongly lifted for class 0 only. + assert lift.loc[(0, 0)] > 3.0 + assert lift.loc[(0, 1)] == pytest.approx(1.0, abs=1e-6) diff --git a/tests/interpretability/test_plugins_and_tier2.py b/tests/interpretability/test_plugins_and_tier2.py index 340e6d7..9c43926 100644 --- a/tests/interpretability/test_plugins_and_tier2.py +++ b/tests/interpretability/test_plugins_and_tier2.py @@ -914,7 +914,8 @@ def test_margin_attribution_plugin_reports_per_cluster_intervals() -> None: ) assert out.name == "margin_attribution" summary = out.payload["per_cluster"] - assert set(summary.index) == {0, 1} + # (cluster_label, class_index); a binary head summarises the positive class. + assert set(summary.index) == {(0, 1), (1, 1)} for column in ("margin_signed_share", "attention_lift", "prevalence"): assert column in summary.columns assert f"{column}_lo" in summary.columns From c5dbe2126bef7303b2f630d563e95607c68868f5 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 10:35:02 +0200 Subject: [PATCH 16/25] Use niche vocabulary, name Background, order niches by composition Interpretability output is read by biologists, but it spoke the algorithm's language. - "Cluster" becomes "niche" throughout the interpretability layer: the `niche_label` column, `NicheSummary`, `niche_attribution_summary`, `niche_profiles`/`niche_agreement`/`per_niche_cell_type_enrichment` plugins, figure titles and axis labels. The *algorithm* is still clustering, so `run_clustering`, `ClusteringResult` and the sklearn parameters (`n_clusters`, `min_cluster_size`) keep their names -- those are the method, niches are the result. - Label -1 is named **Background** rather than "noise", including the `skip_background` parameter and the "excluding Background" figure variants. Density-based clustering puts transitional and sparsely populated tissue there; it is a meaningful category, not a failure. - Niches are ordered by a dendrogram over their composition profiles, so compositionally similar niches are adjacent in every table and heatmap instead of appearing in arbitrary label order. Correlation distance is the default, so equal makeup at different densities reads as similar. Background is pinned last rather than joining the tree, since a catch-all would distort it. Ordering flows from the composition summary into enrichment, counts and the report figures, and is disableable. Figures name niches "Niche 3", keeping the numeric id so a reader can join a figure back to the table without a lookup. Co-Authored-By: Claude Opus 5 (1M context) --- docs/interpretability_extension_guide.md | 16 +- docs/interpretability_suite.md | 51 +++-- docs/method_fidelity.md | 6 +- .../interpretability/plugins/default.yaml | 26 +-- .../interpretability/plugins/full.yaml | 6 +- src/grass_mil/interpretability/__init__.py | 12 +- src/grass_mil/interpretability/contracts.py | 12 +- .../interpretability/core/__init__.py | 12 +- .../interpretability/core/agreement.py | 22 +- .../interpretability/core/attribution.py | 52 +++-- .../interpretability/core/biomarkers.py | 74 ++++--- src/grass_mil/interpretability/core/niches.py | 147 +++++++++++++ .../interpretability/core/transfer.py | 54 +++-- src/grass_mil/interpretability/pipeline.py | 24 +-- .../interpretability/plugins/builtin.py | 200 +++++++++--------- .../reporting/plotly_builders.py | 116 +++++----- .../interpretability/reporting/sections.py | 10 +- .../interpretability/tier2/autocorrelation.py | 2 +- .../interpretability/tier2/filtration.py | 26 +-- .../interpretability/tier2/neighborhood.py | 4 +- .../interpretability/tier2/tissue_graph.py | 10 +- tests/interpretability/test_attribution.py | 36 ++-- tests/interpretability/test_core.py | 18 +- tests/interpretability/test_niches.py | 151 +++++++++++++ .../test_pipeline_and_report.py | 18 +- .../test_plugins_and_tier2.py | 188 ++++++++-------- .../test_spatial_statistics.py | 36 ++-- tests/interpretability/test_tissue_graph.py | 8 +- tests/interpretability/test_transfer.py | 18 +- 29 files changed, 847 insertions(+), 508 deletions(-) create mode 100644 src/grass_mil/interpretability/core/niches.py create mode 100644 tests/interpretability/test_niches.py diff --git a/docs/interpretability_extension_guide.md b/docs/interpretability_extension_guide.md index 0090e2d..051ad65 100644 --- a/docs/interpretability_extension_guide.md +++ b/docs/interpretability_extension_guide.md @@ -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}.") @@ -69,7 +69,7 @@ Add plugin class in: 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 ``` diff --git a/docs/interpretability_suite.md b/docs/interpretability_suite.md index d4a6616..7f74a85 100644 --- a/docs/interpretability_suite.md +++ b/docs/interpretability_suite.md @@ -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` and `attention_attribution` require `comp_*` columns. There is no fallback to one-hot `cell_type` composition. `variance_estimator` convention (biomarker z-scoring): @@ -245,7 +245,7 @@ Common parameters: Returns a `ClusteringResult` dataclass with: -- cluster labels +- niche labels - resolved parameters - fitted object @@ -271,7 +271,7 @@ Plugin registration is per pipeline run (fresh `PluginRegistry`). Built-in plugins: -- `cluster_profiles` +- `niche_profiles` - `attention_attribution` - `neighborhood_enrichment` - `diff_neighborhood_enrichment` @@ -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/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. @@ -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() @@ -448,10 +448,10 @@ With spatial plugins: 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,attention_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 grass-mil-report \ @@ -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/method_fidelity.md b/docs/method_fidelity.md index 0fe9207..807aa88 100644 --- a/docs/method_fidelity.md +++ b/docs/method_fidelity.md @@ -27,7 +27,7 @@ listed at the bottom. ## Attribution The interpretability claim rests on one identity. If `identity_residual` is not -~0, every cluster-level attribution derived from it is invalid. +~0, every niche-level attribution derived from it is invalid. | Quantity | Definition | Implementation | Test | |---|---|---|---| @@ -44,7 +44,7 @@ The interpretability claim rests on one identity. If `identity_residual` is not | 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_cluster_agreement_*` | +| Cross-space agreement (ARI/AMI/NMI) | `interpretability/core/agreement.py` | `test_niche_agreement_*` | ## Known divergences from `methods.tex` @@ -58,7 +58,7 @@ The manuscript is stale on these points. Each is deliberate. 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 cluster can + `[-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 diff --git a/src/grass_mil/configs/interpretability/plugins/default.yaml b/src/grass_mil/configs/interpretability/plugins/default.yaml index 1783506..da234ff 100644 --- a/src/grass_mil/configs/interpretability/plugins/default.yaml +++ b/src/grass_mil/configs/interpretability/plugins/default.yaml @@ -1,22 +1,22 @@ enabled: - - cluster_profiles + - niche_profiles - attention_attribution params: - cluster_profiles: + niche_profiles: variance_estimator: unbiased attention_attribution: attention_column: attention score_column: score variance_estimator: unbiased neighborhood_enrichment: - label_column: cluster_label + 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: cluster_label + label_column: niche_label condition_column: condition permutation_group_column: sample_id # One of: zscore, obs-exp, log2fc. @@ -28,12 +28,12 @@ params: threshold_start: 0.0 threshold_stop: 55.0 threshold_count: 500 - per_cluster_cell_type_enrichment: + per_niche_cell_type_enrichment: cell_type_column: cell_type - cluster_column: cluster_label - # Clusters with fewer than 2 instances, and HDBSCAN noise (-1) when - # skip_noise is true, are omitted. - skip_noise: true + 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). @@ -59,10 +59,10 @@ params: n_perms: 100 random_state: 0 min_nodes: 3 - # Contrast every cluster against this one; -1 is the HDBSCAN noise label. + # Contrast every niche against this one; -1 is the Background label. reference_group: -1 ripley: - label_column: cluster_label + label_column: niche_label group_column: null x_column: center_x y_column: center_y @@ -71,10 +71,10 @@ params: min_count: 5 # median | max | first. `first` reproduces the legacy order-dependent grid. radius_source: median - cluster_agreement: + niche_agreement: active_name: active # Columns holding partitions from other representation spaces, e.g. joined - # from a previous run's cluster_labels.csv. + # from a previous run's niche_labels.csv. label_columns: [] metrics: [ari, ami, nmi] tissue_graph: diff --git a/src/grass_mil/configs/interpretability/plugins/full.yaml b/src/grass_mil/configs/interpretability/plugins/full.yaml index 266b93b..e611520 100644 --- a/src/grass_mil/configs/interpretability/plugins/full.yaml +++ b/src/grass_mil/configs/interpretability/plugins/full.yaml @@ -1,15 +1,15 @@ # Every builtin analysis enabled. Slower than the default preset: -# Moran's I and Ripley scale with cluster size and permutation count. +# Moran's I and Ripley scale with niche size and permutation count. defaults: - default - _self_ enabled: - - cluster_profiles + - niche_profiles - attention_attribution - margin_attribution - neighborhood_enrichment - - per_cluster_cell_type_enrichment + - per_niche_cell_type_enrichment - filtration_curves - morans_i - ripley diff --git a/src/grass_mil/interpretability/__init__.py b/src/grass_mil/interpretability/__init__.py index 1decef7..3f8b937 100644 --- a/src/grass_mil/interpretability/__init__.py +++ b/src/grass_mil/interpretability/__init__.py @@ -1,7 +1,7 @@ from grass_mil.interpretability.core.biomarkers import ( - cluster_attention_summary, - cluster_biomarker_summary, - cluster_survival_attention_summary, + niche_attention_summary, + niche_composition_summary, + niche_survival_attention_summary, ) from grass_mil.interpretability.core.clustering import run_clustering from grass_mil.interpretability.core.reduction import run_reduction @@ -26,9 +26,9 @@ __all__ = [ "run_reduction", "run_clustering", - "cluster_biomarker_summary", - "cluster_attention_summary", - "cluster_survival_attention_summary", + "niche_composition_summary", + "niche_attention_summary", + "niche_survival_attention_summary", "run_neighborhood_enrichment", "run_diff_neighborhood_enrichment", "compute_filtration_curves", diff --git a/src/grass_mil/interpretability/contracts.py b/src/grass_mil/interpretability/contracts.py index e133083..d503b67 100644 --- a/src/grass_mil/interpretability/contracts.py +++ b/src/grass_mil/interpretability/contracts.py @@ -43,11 +43,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 + niche_counts: pd.Series weighted_scores: Optional[pd.Series] = None mean_scores: Optional[pd.Series] = None attention_present: Optional[pd.Series] = None @@ -85,7 +85,7 @@ class PluginResult: @dataclass(frozen=True) -class ClusterTransferBundle: +class NicheTransferBundle: embedding_columns: List[str] id_column: str cluster_on_pca: bool @@ -99,9 +99,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/grass_mil/interpretability/core/__init__.py b/src/grass_mil/interpretability/core/__init__.py index 60c422a..5d66594 100644 --- a/src/grass_mil/interpretability/core/__init__.py +++ b/src/grass_mil/interpretability/core/__init__.py @@ -1,7 +1,7 @@ from grass_mil.interpretability.core.biomarkers import ( - cluster_attention_summary, - cluster_biomarker_summary, - cluster_survival_attention_summary, + niche_attention_summary, + niche_composition_summary, + niche_survival_attention_summary, ) from grass_mil.interpretability.core.clustering import run_clustering from grass_mil.interpretability.core.data import ( @@ -19,9 +19,9 @@ __all__ = [ "run_reduction", "run_clustering", - "cluster_biomarker_summary", - "cluster_attention_summary", - "cluster_survival_attention_summary", + "niche_composition_summary", + "niche_attention_summary", + "niche_survival_attention_summary", "load_interpretability_dataset", "extract_embedding_set", "fit_cluster_transfer_from_report_bundle", diff --git a/src/grass_mil/interpretability/core/agreement.py b/src/grass_mil/interpretability/core/agreement.py index 062f9d7..bbb030b 100644 --- a/src/grass_mil/interpretability/core/agreement.py +++ b/src/grass_mil/interpretability/core/agreement.py @@ -1,6 +1,6 @@ """Agreement between partitions of the same instances. -The suite clusters several representation spaces -- the encoder embedding, the +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 @@ -19,26 +19,26 @@ import numpy as np import pandas as pd -__all__ = ["ClusterAgreementResult", "cluster_jaccard_matrix", "compute_cluster_agreement"] +__all__ = ["NicheAgreementResult", "niche_jaccard_matrix", "compute_niche_agreement"] _METRICS = ("ari", "ami", "nmi") @dataclass(frozen=True) -class ClusterAgreementResult: +class NicheAgreementResult: """Pairwise agreement between labelings. Attributes: pairwise_metrics: One row per space pair, columns ``ari``/``ami``/``nmi``. - jaccard: ``(space_a, space_b) -> cluster x cluster`` Jaccard overlap. + jaccard: ``(space_a, space_b) -> niche x niche`` Jaccard overlap. """ pairwise_metrics: pd.DataFrame jaccard: Dict[Tuple[str, str], pd.DataFrame] -def cluster_jaccard_matrix(labels_a: Sequence, labels_b: Sequence) -> pd.DataFrame: - """Jaccard overlap between every pair of clusters from two labelings.""" +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: @@ -56,12 +56,12 @@ def cluster_jaccard_matrix(labels_a: Sequence, labels_b: Sequence) -> pd.DataFra return pd.DataFrame(matrix, index=clusters_a, columns=clusters_b) -def compute_cluster_agreement( +def compute_niche_agreement( labelings: Mapping[str, Sequence], *, metrics: Sequence[str] = _METRICS, include_jaccard: bool = True, -) -> ClusterAgreementResult: +) -> NicheAgreementResult: """Chance-corrected agreement between every pair of labelings.""" from sklearn.metrics import ( adjusted_mutual_info_score, @@ -75,7 +75,7 @@ def compute_cluster_agreement( f"Unsupported agreement metric(s): {sorted(unknown)}. Expected a subset of {_METRICS}." ) if len(labelings) < 2: - raise ValueError("Cluster agreement needs at least two labelings to compare.") + 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()} @@ -97,7 +97,7 @@ def compute_cluster_agreement( row[metric] = float(functions[metric](a, b)) rows.append(row) if include_jaccard: - jaccard[(name_a, name_b)] = cluster_jaccard_matrix(a, b) + jaccard[(name_a, name_b)] = niche_jaccard_matrix(a, b) frame = pd.DataFrame(rows).set_index(["space_a", "space_b"]) - return ClusterAgreementResult(pairwise_metrics=frame, jaccard=jaccard) + 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 index fbf4da0..0981a33 100644 --- a/src/grass_mil/interpretability/core/attribution.py +++ b/src/grass_mil/interpretability/core/attribution.py @@ -8,7 +8,7 @@ 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:`cluster_attribution_summary` reports ``identity_residual`` so the +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 @@ -19,7 +19,7 @@ and generalises to any number of classes (see :func:`instance_ovr_margins`). -Cluster-level summaries aggregate within cluster x region cells and are +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. @@ -35,7 +35,7 @@ __all__ = [ "AttributionResult", - "cluster_attribution_summary", + "niche_attribution_summary", "instance_margin_contributions", "instance_ovr_margins", "percentile_bootstrap_ci", @@ -44,11 +44,11 @@ @dataclass(frozen=True) class AttributionResult: - """Per-instance contributions and their cluster-level summaries. + """Per-instance contributions and their niche-level summaries. Attributes: per_instance: One row per instance with its contribution columns. - per_cluster: One row per cluster; each statistic carries ``_lo``/``_hi`` + per_cluster: 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. @@ -174,7 +174,7 @@ def instance_ovr_margins(contributions: np.ndarray) -> np.ndarray: def _resolve_focus_classes(n_classes: int, focus_classes) -> List[int]: - """Which classes get a cluster-level summary. + """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 @@ -191,9 +191,9 @@ def _resolve_focus_classes(n_classes: int, focus_classes) -> List[int]: return list(range(n_classes)) -def cluster_attribution_summary( +def niche_attribution_summary( table: pd.DataFrame, - cluster_labels: np.ndarray, + niche_labels: np.ndarray, *, attention_columns: Sequence[str], logit_columns: Sequence[str], @@ -205,11 +205,11 @@ def cluster_attribution_summary( bag_logits: Optional[Dict[str, Sequence[float]]] = None, focus_classes: Optional[Sequence[int]] = None, ) -> AttributionResult: - """Summarise each cluster's contribution to the bag decisions. + """Summarise each niche's contribution to the bag decisions. Args: table: Instance table; must carry the attention, logit and bag columns. - cluster_labels: One cluster id per row of ``table``. + 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. @@ -223,10 +223,8 @@ 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(cluster_labels) != len(table): - raise ValueError( - f"cluster_labels has {len(cluster_labels)} entries for {len(table)} rows." - ) + 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) @@ -236,7 +234,7 @@ class for multi-class tasks. per_instance = pd.DataFrame( { bag_id_column: table[bag_id_column].astype(str).to_numpy(), - "cluster_label": np.asarray(cluster_labels), + "niche_label": np.asarray(niche_labels), } ) margins = instance_ovr_margins(contributions) @@ -319,16 +317,16 @@ def _summarise_clusters( bag_margin = frame.groupby(bag_id_column)["margin"].sum() bag_size = frame.groupby(bag_id_column).size() - cell_margin = frame.groupby([bag_id_column, "cluster_label"])["margin"].sum() - cell_count = frame.groupby([bag_id_column, "cluster_label"]).size() - cell_attention = frame.groupby([bag_id_column, "cluster_label"])["__attention"].sum() + cell_margin = frame.groupby([bag_id_column, "niche_label"])["margin"].sum() + cell_count = frame.groupby([bag_id_column, "niche_label"]).size() + cell_attention = frame.groupby([bag_id_column, "niche_label"])["__attention"].sum() rng = np.random.default_rng(random_state) rows = [] - for cluster in sorted(frame["cluster_label"].unique()): - cluster_margin = cell_margin.xs(cluster, level="cluster_label") - cluster_count = cell_count.xs(cluster, level="cluster_label") - cluster_attn = cell_attention.xs(cluster, level="cluster_label") + for niche in sorted(frame["niche_label"].unique()): + cluster_margin = cell_margin.xs(niche, level="niche_label") + cluster_count = cell_count.xs(niche, level="niche_label") + cluster_attn = cell_attention.xs(niche, level="niche_label") bags = cluster_margin.index totals = bag_margin.loc[bags].to_numpy(dtype=float) @@ -339,19 +337,19 @@ def _summarise_clusters( prevalence = counts / np.maximum(sizes, 1.0) # Abundance-corrected attention lift: 1 is neutral, >1 preferentially - # attended relative to how common the cluster is in that bag. + # attended relative to how common the niche is in that bag. lift = np.divide(attn, prevalence, out=np.full_like(attn, np.nan), where=prevalence > 0) # Shares are fractions of the bag margin and sum to 1 across a full - # partition, but an individual cluster is NOT bounded to [-1, 1]: it can - # exceed 1 when another cluster pushes the opposite way. Regions whose + # 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] = {"cluster_label": cluster, "n_regions": int(len(bags))} + 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)), @@ -368,4 +366,4 @@ def _summarise_clusters( row["n_instances"] = int(counts.sum()) rows.append(row) - return pd.DataFrame(rows).set_index("cluster_label").sort_index() + 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 index b5c9dcb..b728e96 100644 --- a/src/grass_mil/interpretability/core/biomarkers.py +++ b/src/grass_mil/interpretability/core/biomarkers.py @@ -3,7 +3,11 @@ import numpy as np import pandas as pd -from grass_mil.interpretability.contracts import ClusterSummary +from grass_mil.interpretability.contracts import NicheSummary +from grass_mil.interpretability.core.niches import ( + order_niches_by_composition, + sort_by_niche_order, +) def _resolve_variance_estimator(variance_estimator: str) -> int: @@ -30,18 +34,26 @@ def _infer_composition_matrix( ) -def cluster_biomarker_summary( +def niche_composition_summary( table: pd.DataFrame, - cluster_labels: np.ndarray, + niche_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.") + order_by_composition: bool = True, +) -> NicheSummary: + """Mean composition and z-scored enrichment per niche. + + Niches are ordered by a dendrogram over their composition profiles, so + compositionally similar niches are adjacent in every downstream table and + heatmap. Background is pinned last. Set ``order_by_composition=False`` for + plain numeric order. + """ + 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(cluster_labels, name="cluster_label") + labels = pd.Series(niche_labels, name="niche_label") ddof = _resolve_variance_estimator(variance_estimator) comp_mean = comp.mean(axis=0) @@ -52,17 +64,23 @@ def cluster_biomarker_summary( enrichment = z_comp.groupby(labels).mean().sort_index() counts = labels.value_counts().sort_index() - return ClusterSummary( - cluster_labels=np.asarray(cluster_labels), + if order_by_composition: + order = order_niches_by_composition(composition) + composition = sort_by_niche_order(composition, order) + enrichment = sort_by_niche_order(enrichment, order) + counts = counts.reindex(composition.index) + + return NicheSummary( + niche_labels=np.asarray(niche_labels), composition=composition, enrichment=enrichment, - cluster_counts=counts, + niche_counts=counts, ) -def cluster_attention_summary( +def niche_attention_summary( table: pd.DataFrame, - cluster_labels: np.ndarray, + niche_labels: np.ndarray, *, attention_column: str = "attention", score_column: str = "score", @@ -70,15 +88,15 @@ def cluster_attention_summary( composition_prefix: str = "comp_", variance_estimator: str = "unbiased", eps: float = 1e-8, -) -> ClusterSummary: - summary = cluster_biomarker_summary( +) -> NicheSummary: + summary = niche_composition_summary( table, - cluster_labels, + niche_labels, composition_prefix=composition_prefix, variance_estimator=variance_estimator, ) frame = table.copy() - frame["cluster_label"] = cluster_labels + frame["niche_label"] = niche_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.") @@ -87,19 +105,19 @@ def cluster_attention_summary( weighted_frame = frame.assign(weighted=frame[attention_column] * frame[score_column]) weighted = ( - weighted_frame.groupby("cluster_label") + weighted_frame.groupby("niche_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() + mean_scores = frame.groupby("niche_label")[score_column].mean().sort_index() cluster_attn_in_bag = ( - frame.groupby([bag_id_column, "cluster_label"])[attention_column] + frame.groupby([bag_id_column, "niche_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) + frame.groupby([bag_id_column, "niche_label"]).size().unstack(fill_value=0) ) present_mask = cluster_count_in_bag > 0 @@ -110,11 +128,11 @@ def cluster_attention_summary( 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, + return NicheSummary( + niche_labels=summary.niche_labels, composition=summary.composition, enrichment=summary.enrichment, - cluster_counts=summary.cluster_counts, + niche_counts=summary.niche_counts, weighted_scores=weighted, mean_scores=mean_scores, attention_present=attn_present, @@ -122,9 +140,9 @@ def cluster_attention_summary( ) -def cluster_survival_attention_summary( +def niche_survival_attention_summary( table: pd.DataFrame, - cluster_labels: np.ndarray, + niche_labels: np.ndarray, *, attention_column: str = "attention", hazard_column: str = "hazard", @@ -132,14 +150,14 @@ def cluster_survival_attention_summary( composition_prefix: str = "comp_", variance_estimator: str = "unbiased", eps: float = 1e-8, -) -> ClusterSummary: +) -> NicheSummary: 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( + return niche_attention_summary( renamed, - cluster_labels, + niche_labels, attention_column=attention_column, score_column="score", bag_id_column=bag_id_column, diff --git a/src/grass_mil/interpretability/core/niches.py b/src/grass_mil/interpretability/core/niches.py new file mode 100644 index 0000000..c5e0051 --- /dev/null +++ b/src/grass_mil/interpretability/core/niches.py @@ -0,0 +1,147 @@ +"""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", + "sort_by_niche_order", +] + +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 sort_by_niche_order( + frame: pd.DataFrame, + order: Sequence, + *, + axis: int = 0, + level: Optional[str] = None, +) -> pd.DataFrame: + """Reindex a frame onto a niche order, keeping anything not listed. + + Args: + level: Name of the niche level when the frame carries a MultiIndex. + """ + if frame is None or frame.empty or not len(order): + return frame + + labels = frame.index if axis == 0 else frame.columns + if level is not None and isinstance(labels, pd.MultiIndex): + niche_values = labels.get_level_values(level) + rank = {value: position for position, value in enumerate(order)} + sort_key = [rank.get(value, len(order)) for value in niche_values] + return frame.iloc[np.argsort(sort_key, kind="stable")] + + known = [value for value in order if value in labels] + remainder = [value for value in labels if value not in set(known)] + ordered = known + remainder + return frame.reindex(index=ordered) if axis == 0 else frame.reindex(columns=ordered) diff --git a/src/grass_mil/interpretability/core/transfer.py b/src/grass_mil/interpretability/core/transfer.py index 9d93cde..6df3a44 100644 --- a/src/grass_mil/interpretability/core/transfer.py +++ b/src/grass_mil/interpretability/core/transfer.py @@ -9,7 +9,7 @@ from sklearn.neighbors import KNeighborsClassifier from grass_mil.interpretability.contracts import ( - ClusterTransferBundle, + NicheTransferBundle, InterpretabilityDataset, 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, @@ -102,7 +100,7 @@ def fit_cluster_transfer_from_report_bundle( def apply_cluster_transfer( - bundle: ClusterTransferBundle, + 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_cluster_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_cluster_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/grass_mil/interpretability/pipeline.py b/src/grass_mil/interpretability/pipeline.py index 27af971..7aef90d 100644 --- a/src/grass_mil/interpretability/pipeline.py +++ b/src/grass_mil/interpretability/pipeline.py @@ -8,12 +8,12 @@ import pandas as pd from grass_mil.interpretability.contracts import ( - ClusterSummary, + NicheSummary, InterpretabilityDataset, PluginResult, ReportBundle, ) -from grass_mil.interpretability.core.biomarkers import cluster_biomarker_summary +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 from grass_mil.interpretability.core.reduction import run_reduction @@ -131,7 +131,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 +147,9 @@ def run_interpretability_pipeline( clustering_result = run_clustering(method, clustering_input, **params) labels = clustering_result.labels - cluster_summary: Optional[ClusterSummary] = None + 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 +160,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,7 +183,7 @@ def run_interpretability_pipeline( # Persist core artifacts base_instance = dataset.instance_table.copy() if labels is not None: - base_instance["cluster_label"] = labels + base_instance["niche_label"] = labels _write_table(base_instance, artifacts_dir / "instance_with_clusters.csv") if reduction_result is not None: red_df = pd.DataFrame(reduction_result.embedding) @@ -196,10 +196,10 @@ def run_interpretability_pipeline( cluster_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(cluster_df, artifacts_dir / "niche_labels.csv") summary_payload: Dict[str, Any] = { "rows": int(len(dataset.instance_table)), @@ -214,9 +214,9 @@ def run_interpretability_pipeline( 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/builtin.py b/src/grass_mil/interpretability/plugins/builtin.py index 859642b..86e548c 100644 --- a/src/grass_mil/interpretability/plugins/builtin.py +++ b/src/grass_mil/interpretability/plugins/builtin.py @@ -7,15 +7,15 @@ import pandas as pd from grass_mil.interpretability.contracts import ( - ClusterSummary, + NicheSummary, PluginResult, ReportSection, ) -from grass_mil.interpretability.core.agreement import compute_cluster_agreement -from grass_mil.interpretability.core.attribution import cluster_attribution_summary +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 ( - cluster_attention_summary, - cluster_biomarker_summary, + niche_attention_summary, + niche_composition_summary, ) from grass_mil.interpretability.plugins.base import InterpretabilityPlugin, PluginContext from grass_mil.interpretability.plugins.registry import PluginRegistry, create_plugin_registry @@ -52,29 +52,29 @@ def _context_cluster_summary( context: PluginContext, *, labels: np.ndarray, -) -> ClusterSummary | None: - summary = context.state.get("cluster_summary") - if not isinstance(summary, ClusterSummary): +) -> NicheSummary | None: + summary = context.state.get("niche_summary") + if not isinstance(summary, NicheSummary): return None - if summary.cluster_labels.shape != labels.shape: + if summary.niche_labels.shape != labels.shape: return None - if not np.array_equal(summary.cluster_labels, labels): + if not np.array_equal(summary.niche_labels, labels): return None return summary @dataclass -class ClusterProfilesPlugin(InterpretabilityPlugin): - name: str = "cluster_profiles" +class NicheProfilesPlugin(InterpretabilityPlugin): + name: str = "niche_profiles" def required_inputs(self) -> List[str]: - return ["cluster_labels"] + return ["niche_labels"] def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: - labels = np.asarray(context.state["cluster_labels"]) + labels = np.asarray(context.state["niche_labels"]) summary = _context_cluster_summary(context, labels=labels) if summary is None: - summary = cluster_biomarker_summary( + summary = niche_composition_summary( dataset.instance_table, labels, composition_prefix=str(params.get("composition_prefix", "comp_")), @@ -83,12 +83,12 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: payload = { "composition": summary.composition, "enrichment": summary.enrichment, - "counts": summary.cluster_counts, + "counts": summary.niche_counts, } sections = [ ReportSection( - title="Cluster Profiles", - description="Cluster-level composition and enrichment summary.", + title="Niche Profiles", + description="Niche-level composition and enrichment summary.", tables={ "composition": summary.composition, "enrichment": summary.enrichment, @@ -103,11 +103,11 @@ class AttentionAttributionPlugin(InterpretabilityPlugin): name: str = "attention_attribution" def required_inputs(self) -> List[str]: - return ["cluster_labels"] + return ["niche_labels"] def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: - labels = np.asarray(context.state["cluster_labels"]) - attn_summary = cluster_attention_summary( + labels = np.asarray(context.state["niche_labels"]) + attn_summary = niche_attention_summary( dataset.instance_table, labels, attention_column=str(params.get("attention_column", "attention")), @@ -118,11 +118,11 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: ) base_summary = _context_cluster_summary(context, labels=labels) if base_summary is not None: - summary = ClusterSummary( - cluster_labels=base_summary.cluster_labels, + summary = NicheSummary( + niche_labels=base_summary.niche_labels, composition=base_summary.composition, enrichment=base_summary.enrichment, - cluster_counts=base_summary.cluster_counts, + niche_counts=base_summary.niche_counts, weighted_scores=attn_summary.weighted_scores, mean_scores=attn_summary.mean_scores, attention_present=attn_summary.attention_present, @@ -143,7 +143,7 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: sections = [ ReportSection( title="Attention Attribution", - description="Attention-weighted and abundance-corrected cluster scores.", + description="Attention-weighted and abundance-corrected niche scores.", tables=section_tables, ) ] @@ -155,17 +155,17 @@ class NeighborhoodEnrichmentPlugin(InterpretabilityPlugin): name: str = "neighborhood_enrichment" def required_inputs(self) -> List[str]: - return ["spatial_table", "cluster_labels"] + 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["cluster_label"] = np.asarray(context.state["cluster_labels"]) + table["niche_label"] = np.asarray(context.state["niche_labels"]) result = run_neighborhood_enrichment( table, dataset.spatial_table, - label_column=str(params.get("label_column", "cluster_label")), + 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)), @@ -188,7 +188,7 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: sections = [ ReportSection( title="Neighborhood Enrichment", - description="Cluster adjacency enrichment relative to expected connectivity.", + description="Niche adjacency enrichment relative to expected connectivity.", tables=section_tables, ) ] @@ -200,20 +200,20 @@ class DiffNeighborhoodEnrichmentPlugin(InterpretabilityPlugin): name: str = "diff_neighborhood_enrichment" def required_inputs(self) -> List[str]: - return ["spatial_table", "cluster_labels"] + 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["cluster_label"] = np.asarray(context.state["cluster_labels"]) + 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", "cluster_label")), + 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)), @@ -257,19 +257,19 @@ class FiltrationCurvesPlugin(InterpretabilityPlugin): name: str = "filtration_curves" def required_inputs(self) -> List[str]: - return ["spatial_table", "cluster_labels"] + 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["cluster_label"] = np.asarray(context.state["cluster_labels"]) + 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, - cluster_column=str(params.get("cluster_column", "cluster_label")), + 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")), @@ -279,7 +279,7 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: sections = [ ReportSection( title="Filtration Curves", - description="Distance-threshold accumulation by cluster and cell type.", + description="Distance-threshold accumulation by niche and cell type.", metadata={"thresholds": result.thresholds.tolist()}, ) ] @@ -291,7 +291,7 @@ class TissueGraphPlugin(InterpretabilityPlugin): name: str = "tissue_graph" def required_inputs(self) -> List[str]: - return ["spatial_table", "cluster_labels"] + return ["spatial_table", "niche_labels"] def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: if dataset.spatial_table is None: @@ -301,12 +301,12 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: 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"]) + 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", "cluster_label")) + label_column = str(params.get("label_column", "niche_label")) view = prepare_tissue_graph_view( table, @@ -321,7 +321,7 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: include_edge_distances=True, ) meta = dict(view.metadata) - cluster_counts = meta.get("cluster_counts", {}) + niche_counts = meta.get("niche_counts", {}) section_summary = pd.DataFrame( [ { @@ -329,9 +329,7 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: "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, + "cluster_count": len(niche_counts) if isinstance(niche_counts, dict) else 0, } ] ) @@ -352,7 +350,7 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: sections = [ ReportSection( title=f"Tissue Graph ({sample_column}={sample_value})", - description="Single tissue/sample graph with cluster-colored cells.", + description="Single tissue/sample graph with niche-colored cells.", tables={"summary": section_summary}, metadata=meta, ) @@ -361,58 +359,58 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: @dataclass -class PerClusterCellTypeEnrichmentPlugin(InterpretabilityPlugin): - """Compute cell-type neighborhood enrichment within each cluster separately. +class PerNicheCellTypeEnrichmentPlugin(InterpretabilityPlugin): + """Compute cell-type neighborhood enrichment within each niche separately. - For each cluster, subsets the node table and spatial table to nodes belonging - to that cluster, then runs neighborhood enrichment with label_column=cell_type. - This answers: "within cluster C, which cell types are spatially co-located?" + 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_cluster_cell_type_enrichment" + name: str = "per_niche_cell_type_enrichment" def required_inputs(self) -> List[str]: - return ["spatial_table", "cluster_labels"] + 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-cluster cell-type enrichment.") + 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") ) - cluster_column = str(params.get("cluster_column", "cluster_label")) + 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[cluster_column] = np.asarray(context.state["cluster_labels"]) + table[niche_column] = np.asarray(context.state["niche_labels"]) - unique_clusters = sorted(set(table[cluster_column])) - skip_noise = bool(params.get("skip_noise", True)) + unique_niches = sorted(set(table[niche_column])) + skip_background = bool(params.get("skip_background", True)) - enrichment_by_cluster: Dict[str, pd.DataFrame] = {} + enrichment_by_niche: Dict[str, pd.DataFrame] = {} sections: List[ReportSection] = [] - for cluster_id in unique_clusters: - if skip_noise and cluster_id == -1: + for niche_id in unique_niches: + if skip_background and niche_id == -1: continue - cluster_mask = table[cluster_column] == cluster_id - cluster_table = table[cluster_mask].copy() - if len(cluster_table) < 2: + niche_mask = table[niche_column] == niche_id + niche_table = table[niche_mask].copy() + if len(niche_table) < 2: continue - cluster_node_ids = set(cluster_table[id_column].astype(str)) + niche_node_ids = set(niche_table[id_column].astype(str)) spatial = dataset.spatial_table.copy() - cluster_spatial = spatial[ - spatial["source_id"].astype(str).isin(cluster_node_ids) - & spatial["target_id"].astype(str).isin(cluster_node_ids) + niche_spatial = spatial[ + spatial["source_id"].astype(str).isin(niche_node_ids) + & spatial["target_id"].astype(str).isin(niche_node_ids) ] - if cluster_spatial.empty: + if niche_spatial.empty: continue try: result = run_neighborhood_enrichment( - cluster_table, - cluster_spatial, + niche_table, + niche_spatial, label_column=cell_type_column, id_column=id_column, n_perms=int(params.get("n_perms", 0)), @@ -420,12 +418,12 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: undirected=bool(params.get("undirected", False)), enrichment_mode=str(params.get("enrichment_mode", "obs-exp")), ) - enrichment_by_cluster[str(cluster_id)] = result.enrichment + enrichment_by_niche[str(niche_id)] = result.enrichment sections.append( ReportSection( - title=f"Cell-Type Enrichment — Cluster {cluster_id}", + title=f"Cell-Type Enrichment — Niche {niche_id}", description=( - f"Cell-type neighborhood enrichment within cluster {cluster_id}." + f"Cell-type neighborhood enrichment within niche {niche_id}." ), tables={"enrichment": result.enrichment}, ) @@ -433,7 +431,7 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: except (ValueError, KeyError): continue - payload: Dict[str, Any] = {"enrichment_by_cluster": enrichment_by_cluster} + payload: Dict[str, Any] = {"enrichment_by_niche": enrichment_by_niche} return PluginResult(name=self.name, payload=payload, sections=sections) @@ -445,7 +443,7 @@ def _resolve_columns(table: pd.DataFrame, prefix: str, explicit) -> List[str]: @dataclass class MarginAttributionPlugin(InterpretabilityPlugin): - """Exact additive attribution of bag decisions to instance clusters. + """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. @@ -454,11 +452,11 @@ class MarginAttributionPlugin(InterpretabilityPlugin): name: str = "margin_attribution" def required_inputs(self) -> List[str]: - return ["cluster_labels"] + return ["niche_labels"] def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: table = dataset.instance_table - labels = np.asarray(context.state["cluster_labels"]) + labels = np.asarray(context.state["niche_labels"]) attention_columns = _resolve_columns( table, @@ -476,7 +474,7 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: "Enable the interpretability export in the predict step." ) - result = cluster_attribution_summary( + result = niche_attribution_summary( table, labels, attention_columns=attention_columns, @@ -497,7 +495,7 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: title="Margin Attribution", description=( "Exact additive contributions M[i,c] = A[i,c] * l[i,c], summarised per " - "cluster with percentile bootstrap intervals over regions." + "niche with percentile bootstrap intervals over regions." ), tables={"per_cluster": result.per_cluster}, ) @@ -507,19 +505,19 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: @dataclass class MoransIPlugin(InterpretabilityPlugin): - """Spatial autocorrelation per cluster, with a permutation null.""" + """Spatial autocorrelation per niche, with a permutation null.""" name: str = "morans_i" def required_inputs(self) -> List[str]: - return ["spatial_table", "cluster_labels"] + 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() - cluster_column = str(params.get("cluster_column", "cluster_label")) - table[cluster_column] = np.asarray(context.state["cluster_labels"]) + niche_column = str(params.get("niche_column", "niche_label")) + table[niche_column] = np.asarray(context.state["niche_labels"]) feature_columns = _resolve_columns( table, @@ -533,7 +531,7 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: table, dataset.spatial_table, feature_columns=feature_columns, - group_column=cluster_column, + 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)), @@ -557,7 +555,7 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: ReportSection( title="Spatial Autocorrelation (Moran's I)", description=( - "Global Moran's I per cluster over the instance graph, with a " + "Global Moran's I per niche over the instance graph, with a " "permutation null and BH-FDR adjustment." ), tables=tables, @@ -573,13 +571,13 @@ class RipleyPlugin(InterpretabilityPlugin): name: str = "ripley" def required_inputs(self) -> List[str]: - return ["cluster_labels"] + return ["niche_labels"] def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: table = dataset.instance_table.copy() - label_column = str(params.get("label_column", "cluster_label")) - if label_column == "cluster_label": - table[label_column] = np.asarray(context.state["cluster_labels"]) + 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: @@ -617,37 +615,37 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: @dataclass -class ClusterAgreementPlugin(InterpretabilityPlugin): +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 cluster_labels.csv rather than re-running clustering. + joins a prior run's niche_labels.csv rather than re-running clustering. """ - name: str = "cluster_agreement" + name: str = "niche_agreement" def required_inputs(self) -> List[str]: - return ["cluster_labels"] + 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["cluster_labels"]) + 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"cluster_agreement label column {column!r} is not in the instance table." + 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( - "cluster_agreement needs at least one additional labeling; set " - "params.cluster_agreement.label_columns." + "niche_agreement needs at least one additional labeling; set " + "params.niche_agreement.label_columns." ) - result = compute_cluster_agreement( + result = compute_niche_agreement( labelings, metrics=tuple(params.get("metrics", ("ari", "ami", "nmi"))), include_jaccard=bool(params.get("include_jaccard", True)), @@ -658,7 +656,7 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: } sections = [ ReportSection( - title="Cross-Space Cluster Agreement", + title="Cross-Space Niche Agreement", description=( "Chance-corrected agreement between partitions of the same instances." ), @@ -669,17 +667,17 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: def register_builtin_plugins(registry: PluginRegistry) -> None: - registry.register(ClusterProfilesPlugin()) + registry.register(NicheProfilesPlugin()) registry.register(AttentionAttributionPlugin()) registry.register(NeighborhoodEnrichmentPlugin()) registry.register(DiffNeighborhoodEnrichmentPlugin()) registry.register(FiltrationCurvesPlugin()) registry.register(TissueGraphPlugin()) - registry.register(PerClusterCellTypeEnrichmentPlugin()) + registry.register(PerNicheCellTypeEnrichmentPlugin()) registry.register(MarginAttributionPlugin()) registry.register(MoransIPlugin()) registry.register(RipleyPlugin()) - registry.register(ClusterAgreementPlugin()) + registry.register(NicheAgreementPlugin()) def create_builtin_registry() -> PluginRegistry: diff --git a/src/grass_mil/interpretability/reporting/plotly_builders.py b/src/grass_mil/interpretability/reporting/plotly_builders.py index 7644c5d..099a072 100644 --- a/src/grass_mil/interpretability/reporting/plotly_builders.py +++ b/src/grass_mil/interpretability/reporting/plotly_builders.py @@ -6,6 +6,11 @@ import pandas as pd from grass_mil.interpretability.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, @@ -69,26 +74,27 @@ def build_reduction_scatter(bundle: ReportBundle) -> "go.Figure": return fig -def build_cluster_enrichment_heatmap(bundle: ReportBundle) -> "go.Figure": +def build_niche_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 + 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], - y=[str(i) for i in enr.index], + # 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="Cluster Enrichment Heatmap", + title="Niche Enrichment Heatmap", xaxis_title="Features", - yaxis_title="Cluster", + yaxis_title="Niche", template="plotly_white", ) return fig @@ -102,14 +108,14 @@ def build_plugin_figures(bundle: ReportBundle) -> Dict[str, "go.Figure"]: thresholds = np.asarray(result.payload["thresholds"]) curves = result.payload["curves"] fig = go.Figure() - for cluster, ct_map in curves.items(): + 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"{cluster}:{cell_type}", + name=f"{niche}:{cell_type}", ) ) fig.update_layout( @@ -158,10 +164,10 @@ def build_plugin_figures(bundle: ReportBundle) -> Dict[str, "go.Figure"]: template="plotly_white", ) figures[f"{name}_{pair}"] = fig - elif "enrichment_by_cluster" in result.payload and isinstance( - result.payload["enrichment_by_cluster"], dict + elif "enrichment_by_niche" in result.payload and isinstance( + result.payload["enrichment_by_niche"], dict ): - for cluster_id, enr in result.payload["enrichment_by_cluster"].items(): + for niche_id, enr in result.payload["enrichment_by_niche"].items(): if not isinstance(enr, pd.DataFrame): continue fig = go.Figure( @@ -176,10 +182,10 @@ def build_plugin_figures(bundle: ReportBundle) -> Dict[str, "go.Figure"]: ] ) fig.update_layout( - title=f"{name} — Cluster {cluster_id}", + title=f"{name} — {niche_display_name(niche_id)}", template="plotly_white", ) - figures[f"{name}_cluster_{_safe_slug(cluster_id)}"] = fig + figures[f"{name}_niche_{_safe_slug(niche_id)}"] = fig elif "tissue_graph_view" in result.payload and isinstance( result.payload["tissue_graph_view"], TissueGraphView ): @@ -235,12 +241,12 @@ def build_multi_attribute_scatters( ) -> List[tuple[str, "go.Figure"]]: """Build scatter plots from reduction embeddings colored by multiple attributes. - Creates scatters for both PCA (from cluster_feature_reduction) and the main - reduction (e.g. UMAP), each colored by cluster labels and any extra columns + 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 cluster_feature_reduction. + bundle: ReportBundle with reduction and optionally niche_feature_reduction. extra_columns: mapping of {column_name: colorscale}. Defaults to common attributes: condition, score, attention. @@ -267,13 +273,14 @@ def build_multi_attribute_scatters( else np.zeros((len(table),), dtype=int) ) - noise_mask = labels != -1 if -1 in labels else None + # 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.cluster_feature_reduction is not None: - emb = np.asarray(bundle.cluster_feature_reduction.embedding) + if bundle.niche_feature_reduction is not None: + emb = np.asarray(bundle.niche_feature_reduction.embedding) if emb.shape[1] >= 2: - fitted = bundle.cluster_feature_reduction.fitted_object + 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)" @@ -295,7 +302,7 @@ def build_multi_attribute_scatters( ) for red_name, emb, x_label, y_label in reductions: - # Cluster scatter (all points) + # Niche scatter (all points) fig = go.Figure( data=[ go.Scattergl( @@ -312,24 +319,24 @@ def build_multi_attribute_scatters( ] ) fig.update_layout( - title=f"{red_name.upper()} — Cluster Labels", + title=f"{red_name.upper()} — Niche Labels", xaxis_title=x_label, yaxis_title=y_label, template="plotly_white", ) - figures.append((f"{red_name}_cluster_all", fig)) + figures.append((f"{red_name}_niche_all", fig)) - # Cluster scatter (no noise) - if noise_mask is not None: + # Niche scatter (excluding Background) + if background_mask is not None: fig = go.Figure( data=[ go.Scattergl( - x=emb[noise_mask, 0], - y=emb[noise_mask, 1], + x=emb[background_mask, 0], + y=emb[background_mask, 1], mode="markers", marker={ "size": 3, - "color": labels[noise_mask], + "color": labels[background_mask], "colorscale": "Viridis", "showscale": True, }, @@ -337,12 +344,12 @@ def build_multi_attribute_scatters( ] ) fig.update_layout( - title=f"{red_name.upper()} — Cluster Labels (no noise)", + 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}_cluster_filtered", fig)) + figures.append((f"{red_name}_niche_filtered", fig)) # Extra column scatters for col, cscale in extra_columns.items(): @@ -372,16 +379,16 @@ def build_multi_attribute_scatters( ) figures.append((f"{red_name}_{_safe_slug(col)}", fig)) - if noise_mask is not None: + if background_mask is not None: fig = go.Figure( data=[ go.Scattergl( - x=emb[noise_mask, 0], - y=emb[noise_mask, 1], + x=emb[background_mask, 0], + y=emb[background_mask, 1], mode="markers", marker={ "size": 3, - "color": values[noise_mask], + "color": values[background_mask], "colorscale": cscale, "showscale": True, "colorbar": colorbar, @@ -390,7 +397,7 @@ def build_multi_attribute_scatters( ] ) fig.update_layout( - title=f"{red_name.upper()} — {col} (no noise)", + title=f"{red_name.upper()} — {col} (excluding Background)", xaxis_title=x_label, yaxis_title=y_label, template="plotly_white", @@ -400,7 +407,7 @@ def build_multi_attribute_scatters( return figures -def build_composite_cluster_heatmap(bundle: ReportBundle) -> "go.Figure": +def build_composite_niche_heatmap(bundle: ReportBundle) -> "go.Figure": """Build a composite heatmap with enrichment + predictions + attention lift + abundance. Replicates the notebook's 4-panel biomarker summary figure using Plotly subplots. @@ -408,12 +415,13 @@ def build_composite_cluster_heatmap(bundle: ReportBundle) -> "go.Figure": go = _require_plotly() from plotly.subplots import make_subplots - summary = bundle.cluster_summary + summary = bundle.niche_summary if summary is None: - raise ValueError("Cluster summary is missing.") + raise ValueError("Niche summary is missing.") enr = summary.enrichment - cluster_ids = [str(i) for i in enr.index] + # Dendrogram order comes from the niche summary; Background is named. + niche_ids = niche_display_names(enr.index) feature_names = [str(c) for c in enr.columns] has_attention = summary.mean_scores is not None and summary.attention_lift_present is not None @@ -434,7 +442,7 @@ def build_composite_cluster_heatmap(bundle: ReportBundle) -> "go.Figure": go.Heatmap( z=enr.to_numpy(), x=feature_names, - y=cluster_ids, + y=niche_ids, colorscale="RdBu", zmid=0.0, colorbar={"title": "Enrichment", "x": 0.65, "len": 0.9}, @@ -447,12 +455,12 @@ def build_composite_cluster_heatmap(bundle: ReportBundle) -> "go.Figure": if summary.mean_scores is not None: preds = summary.mean_scores.reindex(enr.index).to_numpy().reshape(-1, 1) else: - preds = summary.cluster_counts.reindex(enr.index).to_numpy().reshape(-1, 1).astype(float) + preds = summary.niche_counts.reindex(enr.index).to_numpy().reshape(-1, 1).astype(float) fig.add_trace( go.Heatmap( z=preds, x=["Predictions"], - y=cluster_ids, + y=niche_ids, colorscale="RdBu_r", colorbar={"title": "Pred", "x": 0.78, "len": 0.9}, ), @@ -467,7 +475,7 @@ def build_composite_cluster_heatmap(bundle: ReportBundle) -> "go.Figure": go.Heatmap( z=lift, x=["Attn Lift"], - y=cluster_ids, + y=niche_ids, colorscale="RdGy", zmid=1.0, zmin=0.7, @@ -479,13 +487,13 @@ def build_composite_cluster_heatmap(bundle: ReportBundle) -> "go.Figure": ) # Panel 4: Log abundance - counts = summary.cluster_counts.reindex(enr.index).to_numpy().astype(float) + 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=cluster_ids, + y=niche_ids, colorscale="Viridis", colorbar={"title": "Log N", "x": 1.02, "len": 0.9}, ), @@ -494,9 +502,9 @@ def build_composite_cluster_heatmap(bundle: ReportBundle) -> "go.Figure": ) fig.update_layout( - title="Cluster Summary: Enrichment + Predictions + Attention + Abundance", + title="Niche Summary: Enrichment + Predictions + Attention + Abundance", template="plotly_white", - height=max(400, 50 * len(cluster_ids) + 200), + height=max(400, 50 * len(niche_ids) + 200), ) return fig @@ -505,17 +513,17 @@ 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))) + 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.cluster_feature_reduction is not None: + if bundle.reduction is not None or bundle.niche_feature_reduction is not None: figures.extend(build_multi_attribute_scatters(bundle)) - # Composite cluster heatmap (enrichment + predictions + attention + abundance) - if bundle.cluster_summary is not None: + # Composite niche heatmap (enrichment + predictions + attention + abundance) + if bundle.niche_summary is not None: try: - figures.append(("composite_cluster_summary", build_composite_cluster_heatmap(bundle))) + figures.append(("composite_niche_summary", build_composite_niche_heatmap(bundle))) except (ValueError, KeyError): pass diff --git a/src/grass_mil/interpretability/reporting/sections.py b/src/grass_mil/interpretability/reporting/sections.py index 2f8bab1..35b30bc 100644 --- a/src/grass_mil/interpretability/reporting/sections.py +++ b/src/grass_mil/interpretability/reporting/sections.py @@ -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/grass_mil/interpretability/tier2/autocorrelation.py b/src/grass_mil/interpretability/tier2/autocorrelation.py index d16c876..296dd7c 100644 --- a/src/grass_mil/interpretability/tier2/autocorrelation.py +++ b/src/grass_mil/interpretability/tier2/autocorrelation.py @@ -221,7 +221,7 @@ def diff_morans_i_vs_reference( """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 HDBSCAN noise cluster. + variances. The reference is typically the Background niche. """ if reference_group not in result.statistic.index: raise ValueError( diff --git a/src/grass_mil/interpretability/tier2/filtration.py b/src/grass_mil/interpretability/tier2/filtration.py index 6130b9b..70e9000 100644 --- a/src/grass_mil/interpretability/tier2/filtration.py +++ b/src/grass_mil/interpretability/tier2/filtration.py @@ -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/grass_mil/interpretability/tier2/neighborhood.py b/src/grass_mil/interpretability/tier2/neighborhood.py index 1c0cdfe..74af8d6 100644 --- a/src/grass_mil/interpretability/tier2/neighborhood.py +++ b/src/grass_mil/interpretability/tier2/neighborhood.py @@ -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/tissue_graph.py b/src/grass_mil/interpretability/tier2/tissue_graph.py index 1f6abdd..c29659e 100644 --- a/src/grass_mil/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/tests/interpretability/test_attribution.py b/tests/interpretability/test_attribution.py index 30628c5..64c89bb 100644 --- a/tests/interpretability/test_attribution.py +++ b/tests/interpretability/test_attribution.py @@ -1,7 +1,7 @@ """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 cluster-level attribution built on it is invalid, +does not hold exactly, every niche-level attribution built on it is invalid, so it is asserted directly rather than inferred. """ @@ -12,7 +12,7 @@ import pytest from grass_mil.interpretability.core.attribution import ( - cluster_attribution_summary, + niche_attribution_summary, instance_margin_contributions, instance_ovr_margins, percentile_bootstrap_ci, @@ -50,7 +50,7 @@ def test_contributions_sum_exactly_to_the_bag_logit() -> None: logits = group[["logit_0", "logit_1"]].to_numpy() bag_logits[bag_id] = (attention * logits).sum(axis=0) - result = cluster_attribution_summary( + result = niche_attribution_summary( table, labels, attention_columns=["attention_c0", "attention_c1"], @@ -66,7 +66,7 @@ def test_identity_residual_detects_a_broken_bag_logit() -> None: 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 = cluster_attribution_summary( + result = niche_attribution_summary( table, labels, attention_columns=["attention_c0", "attention_c1"], @@ -100,7 +100,7 @@ def test_margin_contributions_reject_mismatched_shapes() -> None: def test_cluster_summaries_are_bounded_and_carry_intervals() -> None: table, labels = _instance_table(n_bags=6, per_bag=8) - result = cluster_attribution_summary( + result = niche_attribution_summary( table, labels, attention_columns=["attention_c0", "attention_c1"], @@ -109,12 +109,12 @@ def test_cluster_summaries_are_bounded_and_carry_intervals() -> None: ) summary = result.per_cluster - # Indexed by (cluster_label, class_index). A binary head summarises the + # 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 cluster pushing hard in one direction can exceed 1 - # when another cluster opposes it. The real invariant is that a full + # 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() @@ -126,7 +126,7 @@ def test_cluster_summaries_are_bounded_and_carry_intervals() -> None: def test_attention_lift_is_neutral_when_attention_matches_abundance() -> None: - """Uniform attention over a bag gives every cluster a lift of 1.""" + """Uniform attention over a bag gives every niche a lift of 1.""" per_bag = 6 rows = [] for bag in range(4): @@ -141,7 +141,7 @@ def test_attention_lift_is_neutral_when_attention_matches_abundance() -> None: table = pd.DataFrame(rows) labels = np.array([i % 2 for i in range(len(table))]) - result = cluster_attribution_summary( + result = niche_attribution_summary( table, labels, attention_columns=["attention_c0"], @@ -154,7 +154,7 @@ def test_attention_lift_is_neutral_when_attention_matches_abundance() -> None: 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 = cluster_attribution_summary( + result = niche_attribution_summary( table, labels, attention_columns=["attention_c0", "attention_c1"], @@ -234,14 +234,14 @@ def test_multiclass_attribution_summarises_every_class() -> None: attention_columns = [f"attention_c{c}" for c in range(3)] logit_columns = [f"logit_{c}" for c in range(3)] - result = cluster_attribution_summary( + result = niche_attribution_summary( table, labels, attention_columns=attention_columns, logit_columns=logit_columns, n_bootstrap=20, ) - # One row per (cluster, class): no class is silently ignored. + # One row per (niche, class): no class is silently ignored. assert set(result.per_cluster.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 @@ -259,7 +259,7 @@ def test_multiclass_identity_holds_for_every_class() -> None: logits = group[logit_columns].to_numpy() bag_logits[bag_id] = (attention * logits).sum(axis=0) - result = cluster_attribution_summary( + result = niche_attribution_summary( table, labels, attention_columns=attention_columns, @@ -272,7 +272,7 @@ def test_multiclass_identity_holds_for_every_class() -> None: def test_focus_classes_can_be_selected_explicitly() -> None: table, labels = _multiclass_table(n_classes=3) - result = cluster_attribution_summary( + result = niche_attribution_summary( table, labels, attention_columns=[f"attention_c{c}" for c in range(3)], @@ -283,7 +283,7 @@ def test_focus_classes_can_be_selected_explicitly() -> None: assert {k for _, k in result.per_cluster.index} == {2} with pytest.raises(ValueError, match="out of range"): - cluster_attribution_summary( + niche_attribution_summary( table, labels, attention_columns=[f"attention_c{c}" for c in range(3)], @@ -313,7 +313,7 @@ def test_multiclass_attention_lift_uses_the_matching_class_channel() -> None: table = pd.DataFrame(rows) labels = np.array([0 if i % per_bag == 0 else 1 for i in range(len(table))]) - result = cluster_attribution_summary( + result = niche_attribution_summary( table, labels, attention_columns=["attention_c0", "attention_c1", "attention_c2"], @@ -321,6 +321,6 @@ def test_multiclass_attention_lift_uses_the_matching_class_channel() -> None: n_bootstrap=0, ) lift = result.per_cluster["attention_lift"] - # Cluster 0 is the singled-out instance: strongly lifted for class 0 only. + # Niche 0 is the singled-out instance: strongly lifted for class 0 only. assert lift.loc[(0, 0)] > 3.0 assert lift.loc[(0, 1)] == pytest.approx(1.0, abs=1e-6) diff --git a/tests/interpretability/test_core.py b/tests/interpretability/test_core.py index dd5fdce..ffed789 100644 --- a/tests/interpretability/test_core.py +++ b/tests/interpretability/test_core.py @@ -5,8 +5,8 @@ import pytest from grass_mil.interpretability.core.biomarkers import ( - cluster_attention_summary, - cluster_biomarker_summary, + niche_attention_summary, + niche_composition_summary, ) from grass_mil.interpretability.core.clustering import run_clustering from grass_mil.interpretability.core.reduction import run_reduction @@ -37,11 +37,11 @@ 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) + assert bio.niche_counts.sum() == len(table) - att = cluster_attention_summary( + att = niche_attention_summary( table, labels, attention_column="attention", @@ -61,7 +61,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 +74,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 +92,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..ac8cbea --- /dev/null +++ b/tests/interpretability/test_niches.py @@ -0,0 +1,151 @@ +"""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, + sort_by_niche_order, +) + + +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_sort_by_niche_order_keeps_unlisted_rows() -> None: + frame = pd.DataFrame({"value": [1, 2, 3]}, index=[0, 1, 2]) + ordered = sort_by_niche_order(frame, [2, 0]) + assert list(ordered.index) == [2, 0, 1] + + +def test_composition_summary_returns_niches_in_dendrogram_order() -> None: + rng = np.random.default_rng(0) + rows = [] + labels = [] + # Niches 0 and 2 are tumour-rich; 1 and 3 stroma-rich; -1 is Background. + profiles = {0: (0.9, 0.1), 1: (0.1, 0.9), 2: (0.85, 0.15), 3: (0.15, 0.85), -1: (0.5, 0.5)} + for niche, (tumour, stroma) in profiles.items(): + for _ in range(8): + jitter = rng.normal(scale=0.01) + rows.append({"comp_tumour": tumour + jitter, "comp_stroma": stroma - jitter}) + labels.append(niche) + + summary = niche_composition_summary(pd.DataFrame(rows), np.array(labels)) + + order = list(summary.composition.index) + assert order[-1] == BACKGROUND_NICHE_ID + positions = {niche: position for position, niche in enumerate(order)} + assert abs(positions[0] - positions[2]) == 1 + assert abs(positions[1] - positions[3]) == 1 + + # Enrichment and counts follow the same order, so tables line up. + assert list(summary.enrichment.index) == order + assert list(summary.niche_counts.index) == order + + +def test_composition_ordering_can_be_disabled() -> None: + rows = [{"comp_a": 0.9, "comp_b": 0.1}] * 4 + [{"comp_a": 0.1, "comp_b": 0.9}] * 4 + labels = np.array([2, 2, 0, 0, 1, 1, 3, 3]) + summary = niche_composition_summary(pd.DataFrame(rows), labels, order_by_composition=False) + assert list(summary.composition.index) == [0, 1, 2, 3] + + +def test_report_heatmap_labels_niches_by_name() -> None: + pytest.importorskip("plotly.graph_objects") + from grass_mil.interpretability.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 diff --git a/tests/interpretability/test_pipeline_and_report.py b/tests/interpretability/test_pipeline_and_report.py index 41157e0..05c7c42 100644 --- a/tests/interpretability/test_pipeline_and_report.py +++ b/tests/interpretability/test_pipeline_and_report.py @@ -69,7 +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"], + "enabled": ["niche_profiles", "attention_attribution"], "params": { "attention_attribution": {"attention_column": "attention", "score_column": "score"} }, @@ -116,7 +116,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": {}, }, } @@ -248,9 +248,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") @@ -316,7 +316,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( @@ -341,9 +341,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( @@ -367,7 +367,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", + "niche_profiles", "attention_attribution", "neighborhood_enrichment", "diff_neighborhood_enrichment", diff --git a/tests/interpretability/test_plugins_and_tier2.py b/tests/interpretability/test_plugins_and_tier2.py index 9c43926..e0f4e13 100644 --- a/tests/interpretability/test_plugins_and_tier2.py +++ b/tests/interpretability/test_plugins_and_tier2.py @@ -6,7 +6,7 @@ import pandas as pd import pytest -from grass_mil.interpretability.contracts import ClusterSummary, InterpretabilityDataset +from grass_mil.interpretability.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 @@ -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,23 +147,23 @@ 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( - "grass_mil.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) @@ -173,17 +173,17 @@ 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(), + base_summary = NicheSummary( + niche_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}), + niche_counts=pd.Series({0: 3, 1: 3}), ) - attn_summary = ClusterSummary( - cluster_labels=labels.copy(), + attn_summary = NicheSummary( + niche_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}), + niche_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}), @@ -191,13 +191,13 @@ def test_attention_attribution_reuses_context_base_summary(monkeypatch) -> None: ) monkeypatch.setattr( - "grass_mil.interpretability.plugins.builtin.cluster_attention_summary", + "grass_mil.interpretability.plugins.builtin.niche_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}), + PluginContext(state={"niche_labels": labels, "niche_summary": base_summary}), ) assert out.payload["weighted_scores"].equals(attn_summary.weighted_scores) assert out.payload["mean_scores"].equals(attn_summary.mean_scores) @@ -277,11 +277,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 +291,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 +303,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 +318,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 +337,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 +350,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 +359,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 +375,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 +387,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 +401,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 +414,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 +428,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 +440,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 +450,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 +460,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 +476,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 +489,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 +501,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 +516,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, @@ -556,7 +556,7 @@ def _fake_run_neighborhood(*args, **kwargs): # type: ignore[no-untyped-def] 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 +565,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 +580,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 +593,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 +604,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 +620,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 +634,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 +645,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 +660,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 +672,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 +685,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 +695,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 +709,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 +721,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 +758,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 +769,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 +787,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 +796,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 +814,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 +832,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", @@ -846,16 +846,16 @@ def test_builtin_per_cluster_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_cluster_cell_type_enrichment") + plugin = registry.get("per_niche_cell_type_enrichment") out = plugin.run( ds, - PluginContext(state={"cluster_labels": labels}), + PluginContext(state={"niche_labels": labels}), n_perms=0, enrichment_mode="obs-exp", ) - assert out.name == "per_cluster_cell_type_enrichment" - enrichment = out.payload["enrichment_by_cluster"] - # Both clusters are internally connected in the sample spatial table. + 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. @@ -867,23 +867,23 @@ def test_per_cluster_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_cluster_cell_type_enrichment") + plugin = registry.get("per_niche_cell_type_enrichment") skipped = plugin.run( ds, - PluginContext(state={"cluster_labels": labels}), - skip_noise=True, + PluginContext(state={"niche_labels": labels}), + skip_background=True, n_perms=0, ) - assert set(skipped.payload["enrichment_by_cluster"]) == {"1"} + assert set(skipped.payload["enrichment_by_niche"]) == {"1"} retained = plugin.run( ds, - PluginContext(state={"cluster_labels": labels}), - skip_noise=False, + PluginContext(state={"niche_labels": labels}), + skip_background=False, n_perms=0, ) - assert set(retained.payload["enrichment_by_cluster"]) == {"-1", "1"} + assert set(retained.payload["enrichment_by_niche"]) == {"-1", "1"} def _attribution_dataset() -> InterpretabilityDataset: @@ -910,11 +910,11 @@ def test_margin_attribution_plugin_reports_per_cluster_intervals() -> None: labels = np.array([0, 0, 1, 1, 0, 1]) out = registry.get("margin_attribution").run( - ds, PluginContext(state={"cluster_labels": labels}), n_bootstrap=25 + ds, PluginContext(state={"niche_labels": labels}), n_bootstrap=25 ) assert out.name == "margin_attribution" summary = out.payload["per_cluster"] - # (cluster_label, class_index); a binary head summarises the positive class. + # (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", "attention_lift", "prevalence"): assert column in summary.columns @@ -926,7 +926,7 @@ def test_margin_attribution_requires_instance_logits() -> None: ds = _sample_dataset() # has no logit_* columns with pytest.raises(ValueError, match="instance logit columns"): registry.get("margin_attribution").run( - ds, PluginContext(state={"cluster_labels": np.zeros(6, dtype=int)}) + ds, PluginContext(state={"niche_labels": np.zeros(6, dtype=int)}) ) @@ -936,7 +936,7 @@ def test_morans_i_plugin_runs_over_the_instance_graph() -> None: labels = np.array([0, 0, 0, 1, 1, 1]) out = registry.get("morans_i").run( - ds, PluginContext(state={"cluster_labels": labels}), n_perms=10, min_nodes=2 + 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} @@ -950,7 +950,7 @@ def test_ripley_plugin_produces_centred_curves() -> None: labels = np.array([0, 0, 1, 1, 0, 1]) out = registry.get("ripley").run( - ds, PluginContext(state={"cluster_labels": labels}), n_radii=6, min_count=2 + ds, PluginContext(state={"niche_labels": labels}), n_radii=6, min_count=2 ) assert out.name == "ripley" curves = out.payload["curves"] @@ -971,9 +971,9 @@ def test_cluster_agreement_plugin_compares_stored_labelings() -> None: cell_type_column=ds.cell_type_column, ) - out = registry.get("cluster_agreement").run( + out = registry.get("niche_agreement").run( ds, - PluginContext(state={"cluster_labels": np.array([0, 0, 0, 1, 1, 1])}), + PluginContext(state={"niche_labels": np.array([0, 0, 0, 1, 1, 1])}), label_columns=["prior_labels"], ) metrics = out.payload["pairwise_metrics"] @@ -984,6 +984,6 @@ def test_cluster_agreement_plugin_compares_stored_labelings() -> None: 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("cluster_agreement").run( - _sample_dataset(), PluginContext(state={"cluster_labels": np.zeros(6, dtype=int)}) + 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 index a0fa32d..dadf4c3 100644 --- a/tests/interpretability/test_spatial_statistics.py +++ b/tests/interpretability/test_spatial_statistics.py @@ -11,8 +11,8 @@ import pytest from grass_mil.interpretability.core.agreement import ( - cluster_jaccard_matrix, - compute_cluster_agreement, + niche_jaccard_matrix, + compute_niche_agreement, ) from grass_mil.interpretability.tier2.autocorrelation import ( diff_morans_i_vs_reference, @@ -65,7 +65,7 @@ def _autocorrelation_tables(n_per_group: int = 30): 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, "cluster_label": group, "feature": value}) + 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) @@ -77,7 +77,7 @@ def test_run_morans_i_separates_smooth_from_alternating_groups() -> None: nodes, edges, feature_columns=["feature"], - group_column="cluster_label", + group_column="niche_label", n_perms=50, random_state=0, ) @@ -93,25 +93,25 @@ def test_run_morans_i_skips_groups_below_the_size_threshold() -> None: nodes = pd.DataFrame( { "instance_id": ["n0", "n1", "n2"], - "cluster_label": ["small", "small", "other"], + "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="cluster_label", n_perms=5 + 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["cluster_label"] = nodes["cluster_label"].map({"a": 0, "b": -1}) + nodes["niche_label"] = nodes["niche_label"].map({"a": 0, "b": -1}) result = run_morans_i( nodes, edges, feature_columns=["feature"], - group_column="cluster_label", + group_column="niche_label", n_perms=50, random_state=0, ) @@ -227,36 +227,36 @@ def test_cluster_agreement_is_one_for_identical_and_near_zero_for_independent() labels = rng.integers(0, 4, size=300) independent = rng.integers(0, 4, size=300) - result = compute_cluster_agreement( + 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 clusters are numbered. - relabelled = compute_cluster_agreement({"a": labels, "b": (labels + 1) % 4}) + # 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_cluster_agreement({"only": [0, 1, 2]}) + compute_niche_agreement({"only": [0, 1, 2]}) with pytest.raises(ValueError, match="same instances"): - compute_cluster_agreement({"a": [0, 1, 2], "b": [0, 1]}) + compute_niche_agreement({"a": [0, 1, 2], "b": [0, 1]}) with pytest.raises(ValueError, match="Unsupported agreement metric"): - compute_cluster_agreement({"a": [0, 1], "b": [1, 0]}, metrics=["nonsense"]) + 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 = cluster_jaccard_matrix(a, b) + matrix = niche_jaccard_matrix(a, b) - # Cluster 0 of `a` = {0,1}; cluster 1 of `b` = {1,2,3}; overlap {1}, union 4. + # 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) - # Cluster 1 of `a` = {2,3} is fully inside cluster 1 of `b`. + # 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"): - cluster_jaccard_matrix(np.array([0, 1]), np.array([0, 1, 2])) + 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 ed0ce5c..11150f7 100644 --- a/tests/interpretability/test_tissue_graph.py +++ b/tests/interpretability/test_tissue_graph.py @@ -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 cfa1397..b1ecb26 100644 --- a/tests/interpretability/test_transfer.py +++ b/tests/interpretability/test_transfer.py @@ -51,9 +51,9 @@ 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}, @@ -78,8 +78,8 @@ def test_cluster_transfer_fit_and_apply_matches_direct_pca_knn_reference(tmp_pat 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) @@ -167,9 +167,9 @@ 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}, @@ -180,14 +180,14 @@ def test_cluster_transfer_fit_fails_fast_for_missing_inputs(tmp_path: Path) -> N 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"): + with pytest.raises(ValueError, match="missing niche_feature_reduction"): fit_cluster_transfer_from_report_bundle(dataset, missing_cluster_pca, n_neighbors=3) with pytest.raises(ValueError, match="n_neighbors cannot exceed"): From a026f3d9c788ce50f3b302d81a3be44a4fa6ebbf Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 10:36:01 +0200 Subject: [PATCH 17/25] Finish the niche rename in transfer helpers and output artifacts Identifiers that embedded "cluster" inside a longer name were missed by the word-boundary pass: the cluster-transfer helpers, AttributionResult.per_cluster, and the instance_with_clusters.csv artifact. All now use niche vocabulary. The clustering algorithm's own names are deliberately left alone (run_clustering, ClusteringResult, n_clusters, min_cluster_size, and the internal PCA feeding the clustering step). Co-Authored-By: Claude Opus 5 (1M context) --- docs/interpretability_suite.md | 30 +++++----- src/grass_mil/interpretability/__init__.py | 16 +++--- .../interpretability/core/__init__.py | 16 +++--- .../interpretability/core/attribution.py | 6 +- .../interpretability/core/transfer.py | 8 +-- src/grass_mil/interpretability/pipeline.py | 6 +- .../interpretability/plugins/builtin.py | 4 +- tests/interpretability/test_attribution.py | 12 ++-- .../test_plugins_and_tier2.py | 8 +-- tests/interpretability/test_transfer.py | 56 +++++++++---------- 10 files changed, 81 insertions(+), 81 deletions(-) diff --git a/docs/interpretability_suite.md b/docs/interpretability_suite.md index 7f74a85..22e321d 100644 --- a/docs/interpretability_suite.md +++ b/docs/interpretability_suite.md @@ -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, ) @@ -305,7 +305,7 @@ Workflow: 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/niche_labels.csv` (if enabled) - `artifacts/pipeline_summary.json` @@ -340,10 +340,10 @@ import pandas as pd from grass_mil.interpretability.core.data import load_interpretability_dataset from grass_mil.interpretability.core.transfer import ( - apply_cluster_transfer, - fit_cluster_transfer_from_report_bundle, - load_cluster_transfer_bundle, - save_cluster_transfer_bundle, + 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 @@ -363,12 +363,12 @@ 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: diff --git a/src/grass_mil/interpretability/__init__.py b/src/grass_mil/interpretability/__init__.py index 3f8b937..a0bacbb 100644 --- a/src/grass_mil/interpretability/__init__.py +++ b/src/grass_mil/interpretability/__init__.py @@ -6,10 +6,10 @@ 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_cluster_transfer, - fit_cluster_transfer_from_report_bundle, - load_cluster_transfer_bundle, - save_cluster_transfer_bundle, + 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 @@ -34,10 +34,10 @@ "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", ] diff --git a/src/grass_mil/interpretability/core/__init__.py b/src/grass_mil/interpretability/core/__init__.py index 5d66594..15edec0 100644 --- a/src/grass_mil/interpretability/core/__init__.py +++ b/src/grass_mil/interpretability/core/__init__.py @@ -10,10 +10,10 @@ ) from grass_mil.interpretability.core.reduction import run_reduction from grass_mil.interpretability.core.transfer import ( - apply_cluster_transfer, - fit_cluster_transfer_from_report_bundle, - load_cluster_transfer_bundle, - save_cluster_transfer_bundle, + apply_niche_transfer, + fit_niche_transfer_from_report_bundle, + load_niche_transfer_bundle, + save_niche_transfer_bundle, ) __all__ = [ @@ -24,8 +24,8 @@ "niche_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", + "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/attribution.py b/src/grass_mil/interpretability/core/attribution.py index 0981a33..1d5baa6 100644 --- a/src/grass_mil/interpretability/core/attribution.py +++ b/src/grass_mil/interpretability/core/attribution.py @@ -48,14 +48,14 @@ class AttributionResult: Attributes: per_instance: One row per instance with its contribution columns. - per_cluster: One row per niche; each statistic carries ``_lo``/``_hi`` + 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_cluster: pd.DataFrame + per_niche: pd.DataFrame identity_residual: float @@ -264,7 +264,7 @@ class for multi-class tasks. summary = pd.concat(summaries).sort_index() return AttributionResult( per_instance=per_instance, - per_cluster=summary, + per_niche=summary, identity_residual=identity_residual, ) diff --git a/src/grass_mil/interpretability/core/transfer.py b/src/grass_mil/interpretability/core/transfer.py index 6df3a44..0b5f617 100644 --- a/src/grass_mil/interpretability/core/transfer.py +++ b/src/grass_mil/interpretability/core/transfer.py @@ -18,7 +18,7 @@ _TRANSFER_BUNDLE_SCHEMA_VERSION = 1 -def fit_cluster_transfer_from_report_bundle( +def fit_niche_transfer_from_report_bundle( dataset: InterpretabilityDataset, report_bundle: ReportBundle, *, @@ -99,7 +99,7 @@ def fit_cluster_transfer_from_report_bundle( ) -def apply_cluster_transfer( +def apply_niche_transfer( bundle: NicheTransferBundle, instance_table: pd.DataFrame, *, @@ -155,7 +155,7 @@ def apply_cluster_transfer( ) -def save_cluster_transfer_bundle(bundle: NicheTransferBundle, 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 = { @@ -171,7 +171,7 @@ def save_cluster_transfer_bundle(bundle: NicheTransferBundle, path: str | Path) joblib.dump(payload, output_path) -def load_cluster_transfer_bundle(path: str | Path) -> NicheTransferBundle: +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.") diff --git a/src/grass_mil/interpretability/pipeline.py b/src/grass_mil/interpretability/pipeline.py index 7aef90d..1a364df 100644 --- a/src/grass_mil/interpretability/pipeline.py +++ b/src/grass_mil/interpretability/pipeline.py @@ -184,7 +184,7 @@ def run_interpretability_pipeline( base_instance = dataset.instance_table.copy() if labels is not None: base_instance["niche_label"] = labels - _write_table(base_instance, artifacts_dir / "instance_with_clusters.csv") + _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 +193,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(), "niche_label": clustering_result.labels, } ) - _write_table(cluster_df, artifacts_dir / "niche_labels.csv") + _write_table(niche_df, artifacts_dir / "niche_labels.csv") summary_payload: Dict[str, Any] = { "rows": int(len(dataset.instance_table)), diff --git a/src/grass_mil/interpretability/plugins/builtin.py b/src/grass_mil/interpretability/plugins/builtin.py index 86e548c..8c99bd5 100644 --- a/src/grass_mil/interpretability/plugins/builtin.py +++ b/src/grass_mil/interpretability/plugins/builtin.py @@ -487,7 +487,7 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: focus_classes=params.get("focus_classes"), ) payload = { - "per_cluster": result.per_cluster, + "per_niche": result.per_niche, "identity_residual": result.identity_residual, } sections = [ @@ -497,7 +497,7 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: "Exact additive contributions M[i,c] = A[i,c] * l[i,c], summarised per " "niche with percentile bootstrap intervals over regions." ), - tables={"per_cluster": result.per_cluster}, + tables={"per_niche": result.per_niche}, ) ] return PluginResult(name=self.name, payload=payload, sections=sections) diff --git a/tests/interpretability/test_attribution.py b/tests/interpretability/test_attribution.py index 64c89bb..3db38ea 100644 --- a/tests/interpretability/test_attribution.py +++ b/tests/interpretability/test_attribution.py @@ -107,7 +107,7 @@ def test_cluster_summaries_are_bounded_and_carry_intervals() -> None: logit_columns=["logit_0", "logit_1"], n_bootstrap=50, ) - summary = result.per_cluster + 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. @@ -148,7 +148,7 @@ def test_attention_lift_is_neutral_when_attention_matches_abundance() -> None: logit_columns=["logit_0"], n_bootstrap=20, ) - np.testing.assert_allclose(result.per_cluster["attention_lift"].to_numpy(), 1.0, atol=1e-9) + np.testing.assert_allclose(result.per_niche["attention_lift"].to_numpy(), 1.0, atol=1e-9) def test_signed_share_across_clusters_sums_to_one_per_region() -> None: @@ -161,7 +161,7 @@ def test_signed_share_across_clusters_sums_to_one_per_region() -> None: logit_columns=["logit_0", "logit_1"], n_bootstrap=0, ) - assert result.per_cluster["margin_signed_share"].sum() == pytest.approx(1.0, abs=1e-9) + assert result.per_niche["margin_signed_share"].sum() == pytest.approx(1.0, abs=1e-9) def test_bootstrap_ci_brackets_the_point_estimate() -> None: @@ -242,7 +242,7 @@ def test_multiclass_attribution_summarises_every_class() -> None: n_bootstrap=20, ) # One row per (niche, class): no class is silently ignored. - assert set(result.per_cluster.index) == {(c, k) for c in (0, 1, 2) for k in (0, 1, 2)} + 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 @@ -280,7 +280,7 @@ def test_focus_classes_can_be_selected_explicitly() -> None: focus_classes=[2], n_bootstrap=10, ) - assert {k for _, k in result.per_cluster.index} == {2} + assert {k for _, k in result.per_niche.index} == {2} with pytest.raises(ValueError, match="out of range"): niche_attribution_summary( @@ -320,7 +320,7 @@ def test_multiclass_attention_lift_uses_the_matching_class_channel() -> None: logit_columns=["logit_0", "logit_1", "logit_2"], n_bootstrap=0, ) - lift = result.per_cluster["attention_lift"] + lift = result.per_niche["attention_lift"] # Niche 0 is the singled-out instance: strongly lifted for class 0 only. assert lift.loc[(0, 0)] > 3.0 assert lift.loc[(0, 1)] == pytest.approx(1.0, abs=1e-6) diff --git a/tests/interpretability/test_plugins_and_tier2.py b/tests/interpretability/test_plugins_and_tier2.py index e0f4e13..d87a5be 100644 --- a/tests/interpretability/test_plugins_and_tier2.py +++ b/tests/interpretability/test_plugins_and_tier2.py @@ -842,7 +842,7 @@ def test_filtration_runtime_sanity() -> None: assert elapsed < 8.0 -def test_builtin_per_cluster_cell_type_enrichment_plugin_execution() -> None: +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]) @@ -863,7 +863,7 @@ def test_builtin_per_cluster_cell_type_enrichment_plugin_execution() -> None: assert all(section.title.startswith("Cell-Type Enrichment") for section in out.sections) -def test_per_cluster_cell_type_enrichment_skips_noise_cluster() -> None: +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]) @@ -904,7 +904,7 @@ def _attribution_dataset() -> InterpretabilityDataset: ) -def test_margin_attribution_plugin_reports_per_cluster_intervals() -> None: +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]) @@ -913,7 +913,7 @@ def test_margin_attribution_plugin_reports_per_cluster_intervals() -> None: ds, PluginContext(state={"niche_labels": labels}), n_bootstrap=25 ) assert out.name == "margin_attribution" - summary = out.payload["per_cluster"] + 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", "attention_lift", "prevalence"): diff --git a/tests/interpretability/test_transfer.py b/tests/interpretability/test_transfer.py index b1ecb26..5d3d2e1 100644 --- a/tests/interpretability/test_transfer.py +++ b/tests/interpretability/test_transfer.py @@ -14,10 +14,10 @@ ) from grass_mil.interpretability.core.reduction import run_reduction from grass_mil.interpretability.core.transfer import ( - apply_cluster_transfer, - fit_cluster_transfer_from_report_bundle, - load_cluster_transfer_bundle, - save_cluster_transfer_bundle, + apply_niche_transfer, + fit_niche_transfer_from_report_bundle, + load_niche_transfer_bundle, + save_niche_transfer_bundle, ) @@ -60,7 +60,7 @@ def _build_report_bundle( ) -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) @@ -72,9 +72,9 @@ 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") @@ -89,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) @@ -99,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) @@ -118,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) @@ -140,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) @@ -175,7 +175,7 @@ def test_cluster_transfer_fit_fails_fast_for_missing_inputs(tmp_path: Path) -> N 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, @@ -188,13 +188,13 @@ def test_cluster_transfer_fit_fails_fast_for_missing_inputs(tmp_path: Path) -> N metadata={"cluster_on_pca": True}, ) with pytest.raises(ValueError, match="missing niche_feature_reduction"): - fit_cluster_transfer_from_report_bundle(dataset, missing_cluster_pca, n_neighbors=3) + 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) @@ -204,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) From 24781d7821e0c3a0f6cea52eea6ddd2a5f29f4f5 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 10:58:31 +0200 Subject: [PATCH 18/25] Require per-class attention and merge the contracts modules Attention was allowed to emit a single channel broadcast across every class logit. That cannot express a niche which pushes toward one class and away from the other, and it makes the per-class attribution identity vacuous. A binary head now needs two attention channels, a scalar head (regression, Cox log-hazard) one; anything else raises with the required width named. Also merges grass_mil/interpretability/contracts.py into grass_mil/contracts.py. Two modules with the same name, one holding a single validator and the other the interpretability payload dataclasses, was a needless ambiguity; the single module stays dependency-light so any layer can import it without a cycle. Co-Authored-By: Claude Opus 5 (1M context) --- .../configs/model/attention/gated.yaml | 2 + .../model/attention/gated_projected.yaml | 2 + src/grass_mil/configs/model/components.yaml | 3 +- .../configs/model/supervised_module.yaml | 1 + src/grass_mil/contracts.py | 135 ++++++++++++++++-- src/grass_mil/interpretability/contracts.py | 107 -------------- .../interpretability/core/biomarkers.py | 2 +- .../interpretability/core/clustering.py | 2 +- src/grass_mil/interpretability/core/data.py | 2 +- .../interpretability/core/reduction.py | 2 +- .../interpretability/core/transfer.py | 2 +- src/grass_mil/interpretability/pipeline.py | 2 +- .../interpretability/plugins/base.py | 2 +- .../interpretability/plugins/builtin.py | 2 +- .../interpretability/reporting/html.py | 2 +- .../reporting/plotly_builders.py | 2 +- .../interpretability/reporting/render.py | 2 +- .../interpretability/reporting/sections.py | 2 +- .../interpretability/tier2/filtration.py | 2 +- .../interpretability/tier2/neighborhood.py | 2 +- tests/eval/test_aggregation.py | 15 +- tests/interpretability/test_niches.py | 2 +- .../test_pipeline_and_report.py | 2 +- .../test_plugins_and_tier2.py | 2 +- tests/interpretability/test_transfer.py | 2 +- tests/models/test_bagging.py | 29 ++-- 26 files changed, 171 insertions(+), 159 deletions(-) delete mode 100644 src/grass_mil/interpretability/contracts.py diff --git a/src/grass_mil/configs/model/attention/gated.yaml b/src/grass_mil/configs/model/attention/gated.yaml index 922a481..db41112 100644 --- a/src/grass_mil/configs/model/attention/gated.yaml +++ b/src/grass_mil/configs/model/attention/gated.yaml @@ -2,4 +2,6 @@ _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 index 38a187a..85478c6 100644 --- a/src/grass_mil/configs/model/attention/gated_projected.yaml +++ b/src/grass_mil/configs/model/attention/gated_projected.yaml @@ -3,4 +3,6 @@ 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/src/grass_mil/configs/model/components.yaml b/src/grass_mil/configs/model/components.yaml index 2db45cd..2116e22 100644 --- a/src/grass_mil/configs/model/components.yaml +++ b/src/grass_mil/configs/model/components.yaml @@ -34,7 +34,8 @@ model: projection_dim: 64 hidden_dim: 16 dropout: false - n_classes: 1 + # Must equal graph_head.output_dim: attention is per-class. + n_classes: 1 ssl: method: bgrl diff --git a/src/grass_mil/configs/model/supervised_module.yaml b/src/grass_mil/configs/model/supervised_module.yaml index e973ac3..e74132f 100644 --- a/src/grass_mil/configs/model/supervised_module.yaml +++ b/src/grass_mil/configs/model/supervised_module.yaml @@ -33,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/src/grass_mil/contracts.py b/src/grass_mil/contracts.py index 837de9c..dea055a 100644 --- a/src/grass_mil/contracts.py +++ b/src/grass_mil/contracts.py @@ -1,27 +1,136 @@ -"""Invariants shared across the model and inference layers. +"""Cross-cutting contracts. -Kept dependency-free and outside both subpackages so either side can import it -without creating a cycle (``grass_mil.models`` and ``grass_mil.inference`` -already import each other through the supervised module). +Invariants and payload shapes shared between the model, inference and +interpretability layers. Kept in one dependency-light module so any layer can +import it without creating a cycle -- ``grass_mil.models`` and +``grass_mil.inference`` already import each other through the supervised +module. """ from __future__ import annotations +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Dict, List, Optional + +import numpy as np +import pandas as pd + __all__ = ["validate_attention_width"] def validate_attention_width(*, attention_width: int, num_classes: int) -> None: - """Reject attention widths that cannot be broadcast against class logits. + """Require one attention channel per class. - Bag pooling computes ``L_c = sum_i A[i,c] * l[i,c]``. Only two widths are - meaningful: one shared attention channel broadcast across all classes, or - one channel per class (the published form, and the one that makes the - attribution identity exact). Any other width would silently broadcast into - the wrong shape. + 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 not in (1, num_classes): + if attention_width != num_classes: raise ValueError( f"Attention emits {attention_width} channel(s) but the head emits " - f"{num_classes} class logit(s). Set model.attention.n_classes to 1 " - f"(one shared channel) or to {num_classes} (one per class)." + f"{num_classes} class logit(s). Attention is per-class: set " + f"model.attention.n_classes to {num_classes}." ) + + +@dataclass(frozen=True) +class EmbeddingSet: + matrix: np.ndarray + feature_columns: List[str] + id_column: str = "instance_id" + + +@dataclass(frozen=True) +class InterpretabilityDataset: + instance_table: pd.DataFrame + bag_table: Optional[pd.DataFrame] = None + spatial_table: Optional[pd.DataFrame] = None + id_column: str = "instance_id" + bag_id_column: str = "bag_id" + cell_type_column: Optional[str] = None + condition_column: Optional[str] = None + + +@dataclass(frozen=True) +class ReductionResult: + method: str + embedding: np.ndarray + params: Dict[str, Any] + fitted_object: Any = None + + +@dataclass(frozen=True) +class ClusteringResult: + method: str + labels: np.ndarray + params: Dict[str, Any] + fitted_object: Any = None + + +@dataclass(frozen=True) +class NicheSummary: + niche_labels: np.ndarray + composition: pd.DataFrame + enrichment: pd.DataFrame + niche_counts: pd.Series + 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 + + +@dataclass(frozen=True) +class NeighborhoodEnrichmentResult: + enrichment: pd.DataFrame + observed: pd.DataFrame + expected: pd.DataFrame + pvalues: Optional[pd.DataFrame] = None + + +@dataclass(frozen=True) +class FiltrationCurvesResult: + thresholds: np.ndarray + curves: Dict[str, Dict[str, np.ndarray]] + + +@dataclass(frozen=True) +class ReportSection: + title: str + description: str + figure_html: Optional[str] = None + tables: Dict[str, pd.DataFrame] = field(default_factory=dict) + metadata: Dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class PluginResult: + name: str + payload: Dict[str, Any] + sections: List[ReportSection] + + +@dataclass(frozen=True) +class NicheTransferBundle: + embedding_columns: List[str] + id_column: str + cluster_on_pca: bool + pca_model: Any = None + knn_model: Any = None + label_values: np.ndarray = field(default_factory=lambda: np.asarray([])) + metadata: Dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class ReportBundle: + dataset: InterpretabilityDataset + reduction: Optional[ReductionResult] + niche_feature_reduction: Optional[ReductionResult] + clustering: Optional[ClusteringResult] + niche_summary: Optional[NicheSummary] + plugin_results: Dict[str, PluginResult] + artifacts_dir: Path + metadata: Dict[str, Any] = field(default_factory=dict) diff --git a/src/grass_mil/interpretability/contracts.py b/src/grass_mil/interpretability/contracts.py deleted file mode 100644 index d503b67..0000000 --- a/src/grass_mil/interpretability/contracts.py +++ /dev/null @@ -1,107 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Dict, List, Optional - -import numpy as np -import pandas as pd - - -@dataclass(frozen=True) -class EmbeddingSet: - matrix: np.ndarray - feature_columns: List[str] - id_column: str = "instance_id" - - -@dataclass(frozen=True) -class InterpretabilityDataset: - instance_table: pd.DataFrame - bag_table: Optional[pd.DataFrame] = None - spatial_table: Optional[pd.DataFrame] = None - id_column: str = "instance_id" - bag_id_column: str = "bag_id" - cell_type_column: Optional[str] = None - condition_column: Optional[str] = None - - -@dataclass(frozen=True) -class ReductionResult: - method: str - embedding: np.ndarray - params: Dict[str, Any] - fitted_object: Any = None - - -@dataclass(frozen=True) -class ClusteringResult: - method: str - labels: np.ndarray - params: Dict[str, Any] - fitted_object: Any = None - - -@dataclass(frozen=True) -class NicheSummary: - niche_labels: np.ndarray - composition: pd.DataFrame - enrichment: pd.DataFrame - niche_counts: pd.Series - 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 - - -@dataclass(frozen=True) -class NeighborhoodEnrichmentResult: - enrichment: pd.DataFrame - observed: pd.DataFrame - expected: pd.DataFrame - pvalues: Optional[pd.DataFrame] = None - - -@dataclass(frozen=True) -class FiltrationCurvesResult: - thresholds: np.ndarray - curves: Dict[str, Dict[str, np.ndarray]] - - -@dataclass(frozen=True) -class ReportSection: - title: str - description: str - figure_html: Optional[str] = None - tables: Dict[str, pd.DataFrame] = field(default_factory=dict) - metadata: Dict[str, Any] = field(default_factory=dict) - - -@dataclass(frozen=True) -class PluginResult: - name: str - payload: Dict[str, Any] - sections: List[ReportSection] - - -@dataclass(frozen=True) -class NicheTransferBundle: - embedding_columns: List[str] - id_column: str - cluster_on_pca: bool - pca_model: Any = None - knn_model: Any = None - label_values: np.ndarray = field(default_factory=lambda: np.asarray([])) - metadata: Dict[str, Any] = field(default_factory=dict) - - -@dataclass(frozen=True) -class ReportBundle: - dataset: InterpretabilityDataset - reduction: Optional[ReductionResult] - niche_feature_reduction: Optional[ReductionResult] - clustering: Optional[ClusteringResult] - niche_summary: Optional[NicheSummary] - plugin_results: Dict[str, PluginResult] - artifacts_dir: Path - metadata: Dict[str, Any] = field(default_factory=dict) diff --git a/src/grass_mil/interpretability/core/biomarkers.py b/src/grass_mil/interpretability/core/biomarkers.py index b728e96..1b061fe 100644 --- a/src/grass_mil/interpretability/core/biomarkers.py +++ b/src/grass_mil/interpretability/core/biomarkers.py @@ -3,7 +3,7 @@ import numpy as np import pandas as pd -from grass_mil.interpretability.contracts import NicheSummary +from grass_mil.contracts import NicheSummary from grass_mil.interpretability.core.niches import ( order_niches_by_composition, sort_by_niche_order, diff --git a/src/grass_mil/interpretability/core/clustering.py b/src/grass_mil/interpretability/core/clustering.py index 61ea0f6..ac10e78 100644 --- a/src/grass_mil/interpretability/core/clustering.py +++ b/src/grass_mil/interpretability/core/clustering.py @@ -11,7 +11,7 @@ SpectralClustering, ) -from grass_mil.interpretability.contracts import ClusteringResult +from grass_mil.contracts import ClusteringResult def run_clustering( diff --git a/src/grass_mil/interpretability/core/data.py b/src/grass_mil/interpretability/core/data.py index 013e20b..a39278d 100644 --- a/src/grass_mil/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 grass_mil.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/reduction.py b/src/grass_mil/interpretability/core/reduction.py index 6306ff5..c0e889a 100644 --- a/src/grass_mil/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 grass_mil.interpretability.contracts import ReductionResult +from grass_mil.contracts import ReductionResult def run_reduction( diff --git a/src/grass_mil/interpretability/core/transfer.py b/src/grass_mil/interpretability/core/transfer.py index 0b5f617..cdc8ebe 100644 --- a/src/grass_mil/interpretability/core/transfer.py +++ b/src/grass_mil/interpretability/core/transfer.py @@ -8,7 +8,7 @@ import pandas as pd from sklearn.neighbors import KNeighborsClassifier -from grass_mil.interpretability.contracts import ( +from grass_mil.contracts import ( NicheTransferBundle, InterpretabilityDataset, ReportBundle, diff --git a/src/grass_mil/interpretability/pipeline.py b/src/grass_mil/interpretability/pipeline.py index 1a364df..f2dc120 100644 --- a/src/grass_mil/interpretability/pipeline.py +++ b/src/grass_mil/interpretability/pipeline.py @@ -7,7 +7,7 @@ import numpy as np import pandas as pd -from grass_mil.interpretability.contracts import ( +from grass_mil.contracts import ( NicheSummary, InterpretabilityDataset, PluginResult, diff --git a/src/grass_mil/interpretability/plugins/base.py b/src/grass_mil/interpretability/plugins/base.py index 04ea2ce..d778118 100644 --- a/src/grass_mil/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 grass_mil.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 index 8c99bd5..46a9939 100644 --- a/src/grass_mil/interpretability/plugins/builtin.py +++ b/src/grass_mil/interpretability/plugins/builtin.py @@ -6,7 +6,7 @@ import numpy as np import pandas as pd -from grass_mil.interpretability.contracts import ( +from grass_mil.contracts import ( NicheSummary, PluginResult, ReportSection, diff --git a/src/grass_mil/interpretability/reporting/html.py b/src/grass_mil/interpretability/reporting/html.py index cfd2e7c..4613a91 100644 --- a/src/grass_mil/interpretability/reporting/html.py +++ b/src/grass_mil/interpretability/reporting/html.py @@ -6,7 +6,7 @@ import pandas as pd -from grass_mil.interpretability.contracts import ReportSection +from grass_mil.contracts import ReportSection def _render_table(title: str, table: pd.DataFrame) -> str: diff --git a/src/grass_mil/interpretability/reporting/plotly_builders.py b/src/grass_mil/interpretability/reporting/plotly_builders.py index 099a072..9edf292 100644 --- a/src/grass_mil/interpretability/reporting/plotly_builders.py +++ b/src/grass_mil/interpretability/reporting/plotly_builders.py @@ -5,7 +5,7 @@ import numpy as np import pandas as pd -from grass_mil.interpretability.contracts import ReportBundle +from grass_mil.contracts import ReportBundle from grass_mil.interpretability.core.niches import ( BACKGROUND_NICHE_ID, niche_display_name, diff --git a/src/grass_mil/interpretability/reporting/render.py b/src/grass_mil/interpretability/reporting/render.py index 8cce6d2..065a207 100644 --- a/src/grass_mil/interpretability/reporting/render.py +++ b/src/grass_mil/interpretability/reporting/render.py @@ -4,7 +4,7 @@ from pathlib import Path from typing import Any, Dict -from grass_mil.interpretability.contracts import ReportBundle +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 diff --git a/src/grass_mil/interpretability/reporting/sections.py b/src/grass_mil/interpretability/reporting/sections.py index 35b30bc..9797347 100644 --- a/src/grass_mil/interpretability/reporting/sections.py +++ b/src/grass_mil/interpretability/reporting/sections.py @@ -2,7 +2,7 @@ from typing import List -from grass_mil.interpretability.contracts import ReportBundle, ReportSection +from grass_mil.contracts import ReportBundle, ReportSection def build_report_sections(bundle: ReportBundle) -> List[ReportSection]: diff --git a/src/grass_mil/interpretability/tier2/filtration.py b/src/grass_mil/interpretability/tier2/filtration.py index 70e9000..65f8132 100644 --- a/src/grass_mil/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 grass_mil.interpretability.contracts import FiltrationCurvesResult +from grass_mil.contracts import FiltrationCurvesResult def compute_filtration_curves( diff --git a/src/grass_mil/interpretability/tier2/neighborhood.py b/src/grass_mil/interpretability/tier2/neighborhood.py index 74af8d6..c8cb6f2 100644 --- a/src/grass_mil/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 grass_mil.interpretability.contracts import NeighborhoodEnrichmentResult +from grass_mil.contracts import NeighborhoodEnrichmentResult def _coerce_edges(spatial_table: pd.DataFrame) -> pd.DataFrame: diff --git a/tests/eval/test_aggregation.py b/tests/eval/test_aggregation.py index b24a2f7..366ed5d 100644 --- a/tests/eval/test_aggregation.py +++ b/tests/eval/test_aggregation.py @@ -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,9 +72,13 @@ 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 - # Attention is [n_instances, n_channels]; a shared channel keeps one column. - assert out.bag_attention[0].shape == (2, 1) - 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: diff --git a/tests/interpretability/test_niches.py b/tests/interpretability/test_niches.py index ac8cbea..ee2791c 100644 --- a/tests/interpretability/test_niches.py +++ b/tests/interpretability/test_niches.py @@ -115,7 +115,7 @@ def test_composition_ordering_can_be_disabled() -> None: def test_report_heatmap_labels_niches_by_name() -> None: pytest.importorskip("plotly.graph_objects") - from grass_mil.interpretability.contracts import NicheSummary, ReportBundle + 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, diff --git a/tests/interpretability/test_pipeline_and_report.py b/tests/interpretability/test_pipeline_and_report.py index 05c7c42..e55993c 100644 --- a/tests/interpretability/test_pipeline_and_report.py +++ b/tests/interpretability/test_pipeline_and_report.py @@ -8,7 +8,7 @@ from hydra.core.hydra_config import HydraConfig from omegaconf import DictConfig, open_dict -from grass_mil.interpretability.contracts import ClusteringResult +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 diff --git a/tests/interpretability/test_plugins_and_tier2.py b/tests/interpretability/test_plugins_and_tier2.py index d87a5be..c54d50a 100644 --- a/tests/interpretability/test_plugins_and_tier2.py +++ b/tests/interpretability/test_plugins_and_tier2.py @@ -6,7 +6,7 @@ import pandas as pd import pytest -from grass_mil.interpretability.contracts import NicheSummary, InterpretabilityDataset +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 diff --git a/tests/interpretability/test_transfer.py b/tests/interpretability/test_transfer.py index 5d3d2e1..7466bff 100644 --- a/tests/interpretability/test_transfer.py +++ b/tests/interpretability/test_transfer.py @@ -7,7 +7,7 @@ import pytest from sklearn.neighbors import KNeighborsClassifier -from grass_mil.interpretability.contracts import ( +from grass_mil.contracts import ( ClusteringResult, InterpretabilityDataset, ReportBundle, diff --git a/tests/models/test_bagging.py b/tests/models/test_bagging.py index b40f6a4..0f7934d 100644 --- a/tests/models/test_bagging.py +++ b/tests/models/test_bagging.py @@ -78,30 +78,29 @@ def test_per_class_attention_can_select_different_instances_per_class() -> None: assert bag_logits[0, 1] == pytest.approx(200.0, abs=1e-2) -def test_shared_attention_channel_broadcasts_across_classes() -> None: - """n_classes=1 keeps the previous single-channel behaviour.""" - logits = torch.tensor([[1.0, 10.0], [3.0, 30.0]]) - attn_logits = torch.tensor([[1.0], [9.0]]) +def test_pooling_rejects_a_shared_attention_channel() -> None: + """One channel broadcast across classes cannot express opposing evidence. - bag_logits, _, bag_attention, _ = _pool(logits=logits, attn_logits=attn_logits) - - attention = bag_attention["bag"] - assert attention.shape == (2, 1) - expected = (logits * attention).sum(dim=0, keepdim=True) - torch.testing.assert_close(bag_logits, expected, rtol=0, atol=1e-12) + 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_is_neither_shared_nor_per_class() -> None: +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)) -@pytest.mark.parametrize("width", [1, 4]) -def test_validate_attention_width_accepts_shared_and_per_class(width: int) -> None: - validate_attention_width(attention_width=width, num_classes=4) +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, 2, 5]) +@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) From 4c1a322e754a045ee2e908d5e808abbdc15714f6 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 11:01:14 +0200 Subject: [PATCH 19/25] Make the augmentation generator functional and drop the duplicate path drop_edges accepted a torch.Generator and never used it: it delegated to torch_geometric's dropout_edge, which takes no generator and draws from the global RNG. The module advertised purity while BGRL additionally passed a literal None, so pretraining views were not reproducible under a fixed seed and the existing determinism test passed only because it exercised a different function. - drop_edges now takes the draw itself, so the generator actually governs it. - BGRL seeds one generator from task.augmentation_seed and advances it across views, so the two views of a batch differ while the run stays reproducible and independent of unrelated global RNG consumption. - Deleted ssl_runtime.augment_graph. It duplicated drop_edges/drop_features with different semantics -- elementwise feature masking rather than dropping whole columns -- so the same-looking config behaved differently depending on which path ran. The column-wise version is kept, matching the legacy DropFeatures it was ported from. The uniform preset now goes through build_augmentation like every other mode. Tests assert that the generator drives the draw, that the global RNG cannot influence it, and that BGRL's two views repeat across runs while differing from each other. Co-Authored-By: Claude Opus 5 (1M context) --- src/grass_mil/configs/task/pretrain_bgrl.yaml | 3 + src/grass_mil/models/bgrl_module.py | 71 ++++++++++++------- src/grass_mil/models/training/__init__.py | 3 +- .../models/training/augmentations.py | 41 ++++++----- src/grass_mil/models/training/ssl_runtime.py | 40 ----------- tests/models/test_augmentations.py | 70 ++++++++++++++++++ 6 files changed, 140 insertions(+), 88 deletions(-) diff --git a/src/grass_mil/configs/task/pretrain_bgrl.yaml b/src/grass_mil/configs/task/pretrain_bgrl.yaml index e199f0d..fbbc2c3 100644 --- a/src/grass_mil/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/src/grass_mil/models/bgrl_module.py b/src/grass_mil/models/bgrl_module.py index 5aa8cef..a2f1d51 100644 --- a/src/grass_mil/models/bgrl_module.py +++ b/src/grass_mil/models/bgrl_module.py @@ -9,7 +9,6 @@ from .components import EncoderConfig, build_encoder, build_ssl from .training import ( CosineWarmup, - augment_graph, build_augmentation, infer_encoder_input_dim, resolve_encoder_cfg, @@ -58,10 +57,21 @@ 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)) @@ -71,6 +81,7 @@ def __init__( # 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: @@ -79,16 +90,18 @@ def setup(self, stage: Optional[str] = None) -> None: 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: - input_proj = getattr(self.encoder, "input_proj", None) - categorical = getattr(input_proj, "categorical", None) - self._augmentation = build_augmentation( - self._augmentation_cfg, - 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 - ), + 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: @@ -127,21 +140,25 @@ def _encode_tuple(self, batch): 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, None), self._augmentation(batch, None) - return ( - augment_graph( - batch, - drop_edge_p=self._drop_edge_p1, - drop_feat_p=self._drop_feat_p1, - ), - augment_graph( - batch, - drop_edge_p=self._drop_edge_p2, - drop_feat_p=self._drop_feat_p2, - ), - ) + return self._augmentation(batch, generator), self._augmentation(batch, generator) + first, second = self._fallback_augmentations + return first(batch, generator), second(batch, generator) def _step(self, batch) -> torch.Tensor: view1, view2 = self._make_views(batch) diff --git a/src/grass_mil/models/training/__init__.py b/src/grass_mil/models/training/__init__.py index 4e56a4c..978c0a3 100644 --- a/src/grass_mil/models/training/__init__.py +++ b/src/grass_mil/models/training/__init__.py @@ -36,7 +36,7 @@ instantiate_scheduler, instantiate_scheduler_with_warmup, ) -from .ssl_runtime import CosineWarmup, augment_graph +from .ssl_runtime import CosineWarmup __all__ = [ "aggregate_bag_logits_attention", @@ -66,5 +66,4 @@ "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 index 30392cf..15b7650 100644 --- a/src/grass_mil/models/training/augmentations.py +++ b/src/grass_mil/models/training/augmentations.py @@ -145,31 +145,34 @@ def drop_edges( ): """Uniformly drop edges. - ``force_undirected`` defaults to True: the legacy finetune-time transform - defaulted it to False while BGRL pretraining used True, so the two regimes - silently augmented differently. + ``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. """ - from torch_geometric.utils import dropout_edge - 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, edge_mask = dropout_edge( - out.edge_index, - p=float(p), - force_undirected=force_undirected, - training=True, - ) - out.edge_index = edge_index + 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: - if force_undirected: - # dropout_edge returns the mask over the original edges and then - # re-symmetrises, so the attributes are duplicated to match. - kept = edge_attr[edge_mask] - out.edge_attr = torch.cat([kept, kept], dim=0)[: edge_index.size(1)] - else: - out.edge_attr = edge_attr[edge_mask] + out.edge_attr = edge_attr[keep] return out diff --git a/src/grass_mil/models/training/ssl_runtime.py b/src/grass_mil/models/training/ssl_runtime.py index a690e19..f1bad8c 100644 --- a/src/grass_mil/models/training/ssl_runtime.py +++ b/src/grass_mil/models/training/ssl_runtime.py @@ -3,14 +3,6 @@ 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: @@ -28,35 +20,3 @@ def value(self, step: int) -> float: 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, - force_undirected: bool = True, -) -> Data: - """Elementwise feature masking plus uniform edge dropping. - - ``force_undirected`` defaults to True. The legacy transform left it False - here while BGRL pretraining used True, so finetuning silently augmented an - asymmetric graph and message passing became direction-dependent. - """ - 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, force_undirected=force_undirected - ) - aug.edge_index = edge_index - if hasattr(aug, "edge_attr") and aug.edge_attr is not None: - kept = aug.edge_attr[edge_mask] - if force_undirected: - # dropout_edge re-symmetrises after masking, so attributes are - # duplicated to match the restored reverse edges. - kept = torch.cat([kept, kept], dim=0)[: edge_index.size(1)] - aug.edge_attr = kept - return aug diff --git a/tests/models/test_augmentations.py b/tests/models/test_augmentations.py index dad3a0c..3ef132c 100644 --- a/tests/models/test_augmentations.py +++ b/tests/models/test_augmentations.py @@ -171,3 +171,73 @@ def test_build_augmentation_dispatches_on_mode() -> None: 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)) From 57d787a91ace4b998dfdc7cbb8e705db88ce9100 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 11:03:17 +0200 Subject: [PATCH 20/25] Read the head bias from the checkpoint instead of asking for it by hand The config told users to read logit_bias from model_summary.json, but nothing ever wrote that file, so bias-free margins were unreachable in practice: the value could only be transcribed by hand. read_graph_head_bias pulls the output bias straight from a training checkpoint -- the highest-numbered net..bias under the head prefix, which is the final Linear of the head MLP -- and margin_attribution takes logit_bias_checkpoint. An explicit logit_bias still overrides it. An encoder-only checkpoint raises rather than silently attributing with the bias left in. Co-Authored-By: Claude Opus 5 (1M context) --- .../interpretability/plugins/default.yaml | 8 ++- .../interpretability/plugins/builtin.py | 32 +++++++++- src/grass_mil/models/training/__init__.py | 2 + .../models/training/checkpoint_init.py | 38 ++++++++++++ tests/interpretability/test_attribution.py | 59 +++++++++++++++++++ 5 files changed, 135 insertions(+), 4 deletions(-) diff --git a/src/grass_mil/configs/interpretability/plugins/default.yaml b/src/grass_mil/configs/interpretability/plugins/default.yaml index da234ff..dddd311 100644 --- a/src/grass_mil/configs/interpretability/plugins/default.yaml +++ b/src/grass_mil/configs/interpretability/plugins/default.yaml @@ -45,8 +45,12 @@ params: # channel per class. Falls back to the single `attention` column. attention_prefix: attention_c logit_prefix: logit_ - # Head output bias per class, removed so the margin reflects the - # instance-driven part of the decision. Read it from model_summary.json. + # 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 diff --git a/src/grass_mil/interpretability/plugins/builtin.py b/src/grass_mil/interpretability/plugins/builtin.py index 46a9939..9e88b1c 100644 --- a/src/grass_mil/interpretability/plugins/builtin.py +++ b/src/grass_mil/interpretability/plugins/builtin.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional import numpy as np import pandas as pd @@ -441,6 +441,34 @@ def _resolve_columns(table: pd.DataFrame, prefix: str, explicit) -> List[str]: 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. @@ -480,7 +508,7 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: attention_columns=attention_columns, logit_columns=logit_columns, bag_id_column=str(params.get("bag_id_column", dataset.bag_id_column)), - logit_bias=params.get("logit_bias"), + 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)), diff --git a/src/grass_mil/models/training/__init__.py b/src/grass_mil/models/training/__init__.py index 978c0a3..e847a6b 100644 --- a/src/grass_mil/models/training/__init__.py +++ b/src/grass_mil/models/training/__init__.py @@ -23,6 +23,7 @@ validate_task_config, ) from .checkpoint_init import ( + read_graph_head_bias, load_state_dict_with_optional_mapping, remap_encoder_keys, ) @@ -58,6 +59,7 @@ "resolve_encoder_cfg", "validate_task_config", "load_state_dict_with_optional_mapping", + "read_graph_head_bias", "remap_encoder_keys", "compute_binary_accuracy", "compute_categorical_accuracy", diff --git a/src/grass_mil/models/training/checkpoint_init.py b/src/grass_mil/models/training/checkpoint_init.py index 94b052b..e555018 100644 --- a/src/grass_mil/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/tests/interpretability/test_attribution.py b/tests/interpretability/test_attribution.py index 3db38ea..bf565d4 100644 --- a/tests/interpretability/test_attribution.py +++ b/tests/interpretability/test_attribution.py @@ -324,3 +324,62 @@ def test_multiclass_attention_lift_uses_the_matching_class_channel() -> None: # Niche 0 is the singled-out instance: strongly lifted for class 0 only. assert lift.loc[(0, 0)] > 3.0 assert lift.loc[(0, 1)] == pytest.approx(1.0, abs=1e-6) + + +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)) From 77cf8966e57716b60374d765e8b717c3a6acb36d Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 11:05:40 +0200 Subject: [PATCH 21/25] Renumber niches by dendrogram order once, after clustering Niche ordering was re-derived inside the composition summary and reached only that one table, so the attribution and spatial-statistics outputs in the same report still came out in arbitrary label order. Niches are now renumbered once, immediately after clustering: Niche 1 is first in dendrogram order, Niche 2 second, and so on. The id carries the ordering, so every downstream table, figure and export gets composition order by simply sorting on it, and no ordering state has to be threaded through the pipeline. Background keeps its own label and takes no position in the tree. The mapping is recorded in pipeline_summary.json. Consequently niche_composition_summary loses its order_by_composition flag, and sort_by_niche_order is deleted -- it existed only to re-apply an ordering that the ids now carry, and its axis/level parameters were never exercised. Co-Authored-By: Claude Opus 5 (1M context) --- .../interpretability/core/biomarkers.py | 18 +-- src/grass_mil/interpretability/core/niches.py | 68 ++++++--- src/grass_mil/interpretability/pipeline.py | 16 ++ tests/interpretability/test_niches.py | 137 ++++++++++++++---- 4 files changed, 174 insertions(+), 65 deletions(-) diff --git a/src/grass_mil/interpretability/core/biomarkers.py b/src/grass_mil/interpretability/core/biomarkers.py index 1b061fe..999f517 100644 --- a/src/grass_mil/interpretability/core/biomarkers.py +++ b/src/grass_mil/interpretability/core/biomarkers.py @@ -4,10 +4,6 @@ import pandas as pd from grass_mil.contracts import NicheSummary -from grass_mil.interpretability.core.niches import ( - order_niches_by_composition, - sort_by_niche_order, -) def _resolve_variance_estimator(variance_estimator: str) -> int: @@ -40,14 +36,12 @@ def niche_composition_summary( *, composition_prefix: str = "comp_", variance_estimator: str = "unbiased", - order_by_composition: bool = True, ) -> NicheSummary: """Mean composition and z-scored enrichment per niche. - Niches are ordered by a dendrogram over their composition profiles, so - compositionally similar niches are adjacent in every downstream table and - heatmap. Background is pinned last. Set ``order_by_composition=False`` for - plain numeric order. + 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.") @@ -64,12 +58,6 @@ def niche_composition_summary( enrichment = z_comp.groupby(labels).mean().sort_index() counts = labels.value_counts().sort_index() - if order_by_composition: - order = order_niches_by_composition(composition) - composition = sort_by_niche_order(composition, order) - enrichment = sort_by_niche_order(enrichment, order) - counts = counts.reindex(composition.index) - return NicheSummary( niche_labels=np.asarray(niche_labels), composition=composition, diff --git a/src/grass_mil/interpretability/core/niches.py b/src/grass_mil/interpretability/core/niches.py index c5e0051..9933ed8 100644 --- a/src/grass_mil/interpretability/core/niches.py +++ b/src/grass_mil/interpretability/core/niches.py @@ -30,7 +30,7 @@ "niche_display_name", "niche_display_names", "order_niches_by_composition", - "sort_by_niche_order", + "relabel_niches_by_composition", ] BACKGROUND_NICHE_ID = -1 @@ -119,29 +119,49 @@ def order_niches_by_composition( return [ordered_ids[int(position)] for position in order] + background -def sort_by_niche_order( - frame: pd.DataFrame, - order: Sequence, +def relabel_niches_by_composition( + table: pd.DataFrame, + labels: np.ndarray, *, - axis: int = 0, - level: Optional[str] = None, -) -> pd.DataFrame: - """Reindex a frame onto a niche order, keeping anything not listed. + 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. - Args: - level: Name of the niche level when the frame carries a MultiIndex. + 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. """ - if frame is None or frame.empty or not len(order): - return frame - - labels = frame.index if axis == 0 else frame.columns - if level is not None and isinstance(labels, pd.MultiIndex): - niche_values = labels.get_level_values(level) - rank = {value: position for position, value in enumerate(order)} - sort_key = [rank.get(value, len(order)) for value in niche_values] - return frame.iloc[np.argsort(sort_key, kind="stable")] - - known = [value for value in order if value in labels] - remainder = [value for value in labels if value not in set(known)] - ordered = known + remainder - return frame.reindex(index=ordered) if axis == 0 else frame.reindex(columns=ordered) + labels = np.asarray(labels) + composition_columns = [c for c in table.columns if str(c).startswith(composition_prefix)] + if not composition_columns: + # Nothing to order by; leave the labels untouched rather than guessing. + return labels, {value: value for value in np.unique(labels)} + + composition = ( + table.loc[:, composition_columns] + .astype(float) + .groupby(pd.Series(labels, name="niche_label")) + .mean() + .sort_index() + ) + order = order_niches_by_composition(composition, method=method, metric=metric) + + mapping = {} + next_id = int(start_at) + for niche_id in order: + if is_background(niche_id): + mapping[niche_id] = BACKGROUND_NICHE_ID + continue + mapping[niche_id] = next_id + next_id += 1 + + relabelled = np.array([mapping.get(value, value) for value in labels]) + return relabelled, mapping diff --git a/src/grass_mil/interpretability/pipeline.py b/src/grass_mil/interpretability/pipeline.py index f2dc120..0e6e31e 100644 --- a/src/grass_mil/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 @@ -14,6 +16,7 @@ ReportBundle, ) from grass_mil.interpretability.core.biomarkers import niche_composition_summary +from grass_mil.interpretability.core.niches import relabel_niches_by_composition from grass_mil.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 @@ -147,6 +150,18 @@ def run_interpretability_pipeline( clustering_result = run_clustering(method, clustering_input, **params) labels = clustering_result.labels + niche_relabel_map: Dict[Any, Any] = {} + if labels is not None and bool(config.get("relabel_niches_by_composition", True)): + # Renumber once so the niche id carries the dendrogram order. Everything + # downstream then gets composition order by sorting on the id. + labels, niche_relabel_map = relabel_niches_by_composition( + dataset.instance_table, + labels, + composition_prefix=str(config.get("composition_prefix", "comp_")), + ) + if clustering_result is not None: + clustering_result = replace(clustering_result, labels=labels) + niche_summary: Optional[NicheSummary] = None if labels is not None: niche_summary = niche_composition_summary( @@ -208,6 +223,7 @@ 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") diff --git a/tests/interpretability/test_niches.py b/tests/interpretability/test_niches.py index ee2791c..9ef213e 100644 --- a/tests/interpretability/test_niches.py +++ b/tests/interpretability/test_niches.py @@ -14,7 +14,7 @@ niche_display_name, niche_display_names, order_niches_by_composition, - sort_by_niche_order, + relabel_niches_by_composition, ) @@ -75,44 +75,66 @@ def test_ordering_degrades_gracefully_on_tiny_or_empty_inputs() -> None: assert sorted(order_niches_by_composition(constant)) == [0, 1, 2, 3] -def test_sort_by_niche_order_keeps_unlisted_rows() -> None: - frame = pd.DataFrame({"value": [1, 2, 3]}, index=[0, 1, 2]) - ordered = sort_by_niche_order(frame, [2, 0]) - assert list(ordered.index) == [2, 0, 1] +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)) -def test_composition_summary_returns_niches_in_dendrogram_order() -> None: - rng = np.random.default_rng(0) + # 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 = [] - # Niches 0 and 2 are tumour-rich; 1 and 3 stroma-rich; -1 is Background. - profiles = {0: (0.9, 0.1), 1: (0.1, 0.9), 2: (0.85, 0.15), 3: (0.15, 0.85), -1: (0.5, 0.5)} - for niche, (tumour, stroma) in profiles.items(): - for _ in range(8): - jitter = rng.normal(scale=0.01) - rows.append({"comp_tumour": tumour + jitter, "comp_stroma": stroma - jitter}) + 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[-1] == BACKGROUND_NICHE_ID - positions = {niche: position for position, niche in enumerate(order)} - assert abs(positions[0] - positions[2]) == 1 - assert abs(positions[1] - positions[3]) == 1 - - # Enrichment and counts follow the same order, so tables line up. + assert order == sorted(order) assert list(summary.enrichment.index) == order assert list(summary.niche_counts.index) == order -def test_composition_ordering_can_be_disabled() -> None: - rows = [{"comp_a": 0.9, "comp_b": 0.1}] * 4 + [{"comp_a": 0.1, "comp_b": 0.9}] * 4 - labels = np.array([2, 2, 0, 0, 1, 1, 3, 3]) - summary = niche_composition_summary(pd.DataFrame(rows), labels, order_by_composition=False) - assert list(summary.composition.index) == [0, 1, 2, 3] - - def test_report_heatmap_labels_niches_by_name() -> None: pytest.importorskip("plotly.graph_objects") from grass_mil.contracts import NicheSummary, ReportBundle @@ -149,3 +171,66 @@ def test_report_heatmap_labels_niches_by_name() -> 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] From d45fe40c5f75133c522838b3cc58be67b6fadd86 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 11:09:10 +0200 Subject: [PATCH 22/25] Interpret through weighted logit margins only Attention weights on their own say what the model looked at, not which way that pushed the decision, so reporting them alongside margins invited reading attention as evidence. Interpretation is now expressed solely through the weighted logit margins M[i,c] = A[i,c] * (l[i,c] - beta_c). Removed: the attention_attribution plugin, niche_attention_summary and niche_survival_attention_summary, the attention-derived NicheSummary fields (weighted_scores, mean_scores, attention_present, attention_lift_present), and the attention_lift column from margin attribution. The composite report figure drops its predictions and attention-lift panels, keeping composition enrichment beside niche abundance. Attention itself is unchanged -- it still drives bag pooling and still appears in the instance table. What is gone is the reporting of attention as an interpretation in its own right. Co-Authored-By: Claude Opus 5 (1M context) --- .../interpretability/plugins/default.yaml | 5 -- .../interpretability/plugins/full.yaml | 1 - src/grass_mil/contracts.py | 4 - src/grass_mil/interpretability/__init__.py | 4 - .../interpretability/core/__init__.py | 4 - .../interpretability/core/attribution.py | 36 ++------ .../interpretability/core/biomarkers.py | 89 ------------------- .../interpretability/plugins/builtin.py | 58 +----------- .../reporting/plotly_builders.py | 72 ++++----------- tests/interpretability/test_attribution.py | 62 +------------ tests/interpretability/test_core.py | 11 --- .../test_pipeline_and_report.py | 6 +- .../test_plugins_and_tier2.py | 36 +------- 13 files changed, 27 insertions(+), 361 deletions(-) diff --git a/src/grass_mil/configs/interpretability/plugins/default.yaml b/src/grass_mil/configs/interpretability/plugins/default.yaml index dddd311..979321a 100644 --- a/src/grass_mil/configs/interpretability/plugins/default.yaml +++ b/src/grass_mil/configs/interpretability/plugins/default.yaml @@ -1,13 +1,8 @@ enabled: - niche_profiles - - attention_attribution params: niche_profiles: variance_estimator: unbiased - attention_attribution: - attention_column: attention - score_column: score - variance_estimator: unbiased neighborhood_enrichment: label_column: niche_label # One of: zscore, obs-exp, log2fc. diff --git a/src/grass_mil/configs/interpretability/plugins/full.yaml b/src/grass_mil/configs/interpretability/plugins/full.yaml index e611520..fd1e66c 100644 --- a/src/grass_mil/configs/interpretability/plugins/full.yaml +++ b/src/grass_mil/configs/interpretability/plugins/full.yaml @@ -6,7 +6,6 @@ defaults: enabled: - niche_profiles - - attention_attribution - margin_attribution - neighborhood_enrichment - per_niche_cell_type_enrichment diff --git a/src/grass_mil/contracts.py b/src/grass_mil/contracts.py index dea055a..f4b4063 100644 --- a/src/grass_mil/contracts.py +++ b/src/grass_mil/contracts.py @@ -77,10 +77,6 @@ class NicheSummary: composition: pd.DataFrame enrichment: pd.DataFrame niche_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 @dataclass(frozen=True) diff --git a/src/grass_mil/interpretability/__init__.py b/src/grass_mil/interpretability/__init__.py index a0bacbb..9a59eef 100644 --- a/src/grass_mil/interpretability/__init__.py +++ b/src/grass_mil/interpretability/__init__.py @@ -1,7 +1,5 @@ from grass_mil.interpretability.core.biomarkers import ( - niche_attention_summary, niche_composition_summary, - niche_survival_attention_summary, ) from grass_mil.interpretability.core.clustering import run_clustering from grass_mil.interpretability.core.reduction import run_reduction @@ -27,8 +25,6 @@ "run_reduction", "run_clustering", "niche_composition_summary", - "niche_attention_summary", - "niche_survival_attention_summary", "run_neighborhood_enrichment", "run_diff_neighborhood_enrichment", "compute_filtration_curves", diff --git a/src/grass_mil/interpretability/core/__init__.py b/src/grass_mil/interpretability/core/__init__.py index 15edec0..2bf512d 100644 --- a/src/grass_mil/interpretability/core/__init__.py +++ b/src/grass_mil/interpretability/core/__init__.py @@ -1,7 +1,5 @@ from grass_mil.interpretability.core.biomarkers import ( - niche_attention_summary, niche_composition_summary, - niche_survival_attention_summary, ) from grass_mil.interpretability.core.clustering import run_clustering from grass_mil.interpretability.core.data import ( @@ -20,8 +18,6 @@ "run_reduction", "run_clustering", "niche_composition_summary", - "niche_attention_summary", - "niche_survival_attention_summary", "load_interpretability_dataset", "extract_embedding_set", "fit_niche_transfer_from_report_bundle", diff --git a/src/grass_mil/interpretability/core/attribution.py b/src/grass_mil/interpretability/core/attribution.py index 1d5baa6..7566099 100644 --- a/src/grass_mil/interpretability/core/attribution.py +++ b/src/grass_mil/interpretability/core/attribution.py @@ -251,9 +251,8 @@ class for multi-class tasks. for class_index in _resolve_focus_classes(n_classes, focus_classes): frame = per_instance.copy() frame["margin"] = margins[:, class_index] - part = _summarise_clusters( + part = _summarise_niches( frame, - attention=_attention_for_class(attention, class_index), bag_id_column=bag_id_column, n_bootstrap=n_bootstrap, random_state=random_state, @@ -269,15 +268,6 @@ class for multi-class tasks. ) -def _attention_for_class(attention: np.ndarray, class_index: int) -> np.ndarray: - """Attention weights driving one class, or the shared channel.""" - if attention.ndim == 1: - return attention[:, None] - if attention.shape[1] == 1: - return attention - return attention[:, class_index : class_index + 1] - - def _identity_residual( bag_ids: np.ndarray, contributions: np.ndarray, @@ -300,45 +290,34 @@ def _identity_residual( return residual -def _summarise_clusters( +def _summarise_niches( per_instance: pd.DataFrame, *, - attention: np.ndarray, bag_id_column: str, n_bootstrap: int, random_state: int, margin_eps: float, ) -> pd.DataFrame: frame = per_instance.copy() - # A shared attention channel is broadcast; per-class attention is summed to - # one selection weight per instance. - frame["__attention"] = attention.sum(axis=1) if attention.ndim == 2 else attention - 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() - cell_attention = frame.groupby([bag_id_column, "niche_label"])["__attention"].sum() rng = np.random.default_rng(random_state) rows = [] for niche in sorted(frame["niche_label"].unique()): - cluster_margin = cell_margin.xs(niche, level="niche_label") - cluster_count = cell_count.xs(niche, level="niche_label") - cluster_attn = cell_attention.xs(niche, level="niche_label") - bags = cluster_margin.index + 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 = cluster_margin.to_numpy(dtype=float) - counts = cluster_count.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) - attn = cluster_attn.to_numpy(dtype=float) prevalence = counts / np.maximum(sizes, 1.0) - # Abundance-corrected attention lift: 1 is neutral, >1 preferentially - # attended relative to how common the niche is in that bag. - lift = np.divide(attn, prevalence, out=np.full_like(attn, np.nan), where=prevalence > 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 @@ -354,7 +333,6 @@ def _summarise_clusters( ("margin_signed_share", share, None), ("margin_weighted", margins, np.abs(totals)), ("margin_sign_consistency", consistency, None), - ("attention_lift", lift, None), ("prevalence", prevalence, None), ): point, lo, hi = percentile_bootstrap_ci( diff --git a/src/grass_mil/interpretability/core/biomarkers.py b/src/grass_mil/interpretability/core/biomarkers.py index 999f517..11e7a25 100644 --- a/src/grass_mil/interpretability/core/biomarkers.py +++ b/src/grass_mil/interpretability/core/biomarkers.py @@ -64,92 +64,3 @@ def niche_composition_summary( enrichment=enrichment, niche_counts=counts, ) - - -def niche_attention_summary( - table: pd.DataFrame, - niche_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, -) -> NicheSummary: - summary = niche_composition_summary( - table, - niche_labels, - composition_prefix=composition_prefix, - variance_estimator=variance_estimator, - ) - frame = table.copy() - frame["niche_label"] = niche_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("niche_label") - .apply(lambda g: float(g["weighted"].sum() / max(g[attention_column].sum(), eps))) - .sort_index() - ) - mean_scores = frame.groupby("niche_label")[score_column].mean().sort_index() - - cluster_attn_in_bag = ( - frame.groupby([bag_id_column, "niche_label"])[attention_column] - .sum() - .unstack(fill_value=0.0) - ) - cluster_count_in_bag = ( - frame.groupby([bag_id_column, "niche_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 NicheSummary( - niche_labels=summary.niche_labels, - composition=summary.composition, - enrichment=summary.enrichment, - niche_counts=summary.niche_counts, - weighted_scores=weighted, - mean_scores=mean_scores, - attention_present=attn_present, - attention_lift_present=lift_present, - ) - - -def niche_survival_attention_summary( - table: pd.DataFrame, - niche_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, -) -> NicheSummary: - 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 niche_attention_summary( - renamed, - niche_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/grass_mil/interpretability/plugins/builtin.py b/src/grass_mil/interpretability/plugins/builtin.py index 9e88b1c..335cc78 100644 --- a/src/grass_mil/interpretability/plugins/builtin.py +++ b/src/grass_mil/interpretability/plugins/builtin.py @@ -13,10 +13,7 @@ ) 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_attention_summary, - niche_composition_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 ( @@ -98,58 +95,6 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: return PluginResult(name=self.name, payload=payload, sections=sections) -@dataclass -class AttentionAttributionPlugin(InterpretabilityPlugin): - name: str = "attention_attribution" - - 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"]) - attn_summary = niche_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 = NicheSummary( - niche_labels=base_summary.niche_labels, - composition=base_summary.composition, - enrichment=base_summary.enrichment, - niche_counts=base_summary.niche_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 niche scores.", - tables=section_tables, - ) - ] - return PluginResult(name=self.name, payload=payload, sections=sections) - - @dataclass class NeighborhoodEnrichmentPlugin(InterpretabilityPlugin): name: str = "neighborhood_enrichment" @@ -696,7 +641,6 @@ def run(self, dataset, context: PluginContext, **params: Any) -> PluginResult: def register_builtin_plugins(registry: PluginRegistry) -> None: registry.register(NicheProfilesPlugin()) - registry.register(AttentionAttributionPlugin()) registry.register(NeighborhoodEnrichmentPlugin()) registry.register(DiffNeighborhoodEnrichmentPlugin()) registry.register(FiltrationCurvesPlugin()) diff --git a/src/grass_mil/interpretability/reporting/plotly_builders.py b/src/grass_mil/interpretability/reporting/plotly_builders.py index 9edf292..4079e3a 100644 --- a/src/grass_mil/interpretability/reporting/plotly_builders.py +++ b/src/grass_mil/interpretability/reporting/plotly_builders.py @@ -408,9 +408,12 @@ def build_multi_attribute_scatters( def build_composite_niche_heatmap(bundle: ReportBundle) -> "go.Figure": - """Build a composite heatmap with enrichment + predictions + attention lift + abundance. + """Composition enrichment beside niche abundance. - Replicates the notebook's 4-panel biomarker summary figure using Plotly subplots. + 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 @@ -420,24 +423,17 @@ def build_composite_niche_heatmap(bundle: ReportBundle) -> "go.Figure": raise ValueError("Niche summary is missing.") enr = summary.enrichment - # Dendrogram order comes from the niche summary; Background is named. + # 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] - has_attention = summary.mean_scores is not None and summary.attention_lift_present is not None - - n_cols = 4 if has_attention else 2 - widths = [13, 1, 1, 1] if has_attention else [13, 1] - fig = make_subplots( rows=1, - cols=n_cols, + cols=2, shared_yaxes=True, - column_widths=widths, + column_widths=[13, 1], horizontal_spacing=0.02, ) - - # Panel 1: Enrichment heatmap fig.add_trace( go.Heatmap( z=enr.to_numpy(), @@ -445,64 +441,28 @@ def build_composite_niche_heatmap(bundle: ReportBundle) -> "go.Figure": y=niche_ids, colorscale="RdBu", zmid=0.0, - colorbar={"title": "Enrichment", "x": 0.65, "len": 0.9}, + colorbar={"title": "Enrichment", "x": 0.86, "len": 0.9}, ), row=1, col=1, ) - # Panel 2: Mean predictions - if summary.mean_scores is not None: - preds = summary.mean_scores.reindex(enr.index).to_numpy().reshape(-1, 1) - else: - preds = summary.niche_counts.reindex(enr.index).to_numpy().reshape(-1, 1).astype(float) + 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=preds, - x=["Predictions"], + z=log_counts, + x=["Log Abund."], y=niche_ids, - colorscale="RdBu_r", - colorbar={"title": "Pred", "x": 0.78, "len": 0.9}, + colorscale="Viridis", + colorbar={"title": "Log N", "x": 1.02, "len": 0.9}, ), row=1, col=2, ) - if has_attention: - # Panel 3: Attention lift (presence conditioned) - lift = summary.attention_lift_present.reindex(enr.index).to_numpy().reshape(-1, 1) - fig.add_trace( - go.Heatmap( - z=lift, - x=["Attn Lift"], - y=niche_ids, - colorscale="RdGy", - zmid=1.0, - zmin=0.7, - zmax=1.3, - colorbar={"title": "Attn Lift", "x": 0.90, "len": 0.9}, - ), - row=1, - col=3, - ) - - # Panel 4: Log abundance - 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=4, - ) - fig.update_layout( - title="Niche Summary: Enrichment + Predictions + Attention + Abundance", + title="Niche Composition and Abundance", template="plotly_white", height=max(400, 50 * len(niche_ids) + 200), ) diff --git a/tests/interpretability/test_attribution.py b/tests/interpretability/test_attribution.py index bf565d4..923cfe3 100644 --- a/tests/interpretability/test_attribution.py +++ b/tests/interpretability/test_attribution.py @@ -120,37 +120,11 @@ def test_cluster_summaries_are_bounded_and_carry_intervals() -> None: assert summary["margin_sign_consistency"].between(0.0, 1.0).all() assert summary["prevalence"].between(0.0, 1.0).all() - for column in ("attention_lift", "prevalence"): + 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_attention_lift_is_neutral_when_attention_matches_abundance() -> None: - """Uniform attention over a bag gives every niche a lift of 1.""" - per_bag = 6 - rows = [] - for bag in range(4): - for i in range(per_bag): - rows.append( - { - "bag_id": f"bag_{bag}", - "attention_c0": 1.0 / per_bag, - "logit_0": float(i), - } - ) - table = pd.DataFrame(rows) - labels = np.array([i % 2 for i in range(len(table))]) - - result = niche_attribution_summary( - table, - labels, - attention_columns=["attention_c0"], - logit_columns=["logit_0"], - n_bootstrap=20, - ) - np.testing.assert_allclose(result.per_niche["attention_lift"].to_numpy(), 1.0, atol=1e-9) - - 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) @@ -292,40 +266,6 @@ def test_focus_classes_can_be_selected_explicitly() -> None: ) -def test_multiclass_attention_lift_uses_the_matching_class_channel() -> None: - """Each class's summary must use that class's own attention column.""" - per_bag = 4 - rows = [] - for bag in range(3): - for i in range(per_bag): - rows.append( - { - "bag_id": f"bag_{bag}", - # Class 0 attends instance 0 almost exclusively; class 1 is uniform. - "attention_c0": 0.97 if i == 0 else 0.01, - "attention_c1": 1.0 / per_bag, - "attention_c2": 1.0 / per_bag, - "logit_0": float(i), - "logit_1": float(i), - "logit_2": float(i), - } - ) - table = pd.DataFrame(rows) - labels = np.array([0 if i % per_bag == 0 else 1 for i in range(len(table))]) - - result = niche_attribution_summary( - table, - labels, - attention_columns=["attention_c0", "attention_c1", "attention_c2"], - logit_columns=["logit_0", "logit_1", "logit_2"], - n_bootstrap=0, - ) - lift = result.per_niche["attention_lift"] - # Niche 0 is the singled-out instance: strongly lifted for class 0 only. - assert lift.loc[(0, 0)] > 3.0 - assert lift.loc[(0, 1)] == pytest.approx(1.0, abs=1e-6) - - 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 diff --git a/tests/interpretability/test_core.py b/tests/interpretability/test_core.py index ffed789..0eec984 100644 --- a/tests/interpretability/test_core.py +++ b/tests/interpretability/test_core.py @@ -5,7 +5,6 @@ import pytest from grass_mil.interpretability.core.biomarkers import ( - niche_attention_summary, niche_composition_summary, ) from grass_mil.interpretability.core.clustering import run_clustering @@ -41,16 +40,6 @@ def test_cluster_biomarker_and_attention_summary() -> None: assert "comp_A" in bio.composition.columns assert bio.niche_counts.sum() == len(table) - att = niche_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 - def test_cluster_biomarker_requires_composition_columns() -> None: table = pd.DataFrame( diff --git a/tests/interpretability/test_pipeline_and_report.py b/tests/interpretability/test_pipeline_and_report.py index e55993c..f1149bc 100644 --- a/tests/interpretability/test_pipeline_and_report.py +++ b/tests/interpretability/test_pipeline_and_report.py @@ -69,10 +69,7 @@ def test_pipeline_and_report_render(monkeypatch, tmp_path: Path) -> None: "params": {"n_clusters": 2, "linkage": "ward"}, }, "plugins": { - "enabled": ["niche_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") @@ -368,7 +365,6 @@ def test_report_cli_smoke( cfg_interpret.reduction.params = {"n_components": 2, "random_state": 42} cfg_interpret.plugins.enabled = [ "niche_profiles", - "attention_attribution", "neighborhood_enrichment", "diff_neighborhood_enrichment", "filtration_curves", diff --git a/tests/interpretability/test_plugins_and_tier2.py b/tests/interpretability/test_plugins_and_tier2.py index c54d50a..846e7b5 100644 --- a/tests/interpretability/test_plugins_and_tier2.py +++ b/tests/interpretability/test_plugins_and_tier2.py @@ -169,40 +169,6 @@ def _should_not_run(*args, **kwargs): # type: ignore[no-untyped-def] 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 = NicheSummary( - niche_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]), - niche_counts=pd.Series({0: 3, 1: 3}), - ) - attn_summary = NicheSummary( - niche_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]), - niche_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( - "grass_mil.interpretability.plugins.builtin.niche_attention_summary", - lambda *args, **kwargs: attn_summary, # type: ignore[no-untyped-def] - ) - plugin = registry.get("attention_attribution") - out = plugin.run( - ds, - PluginContext(state={"niche_labels": labels, "niche_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" @@ -916,7 +882,7 @@ def test_margin_attribution_plugin_reports_per_niche_intervals() -> None: 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", "attention_lift", "prevalence"): + for column in ("margin_signed_share", "prevalence"): assert column in summary.columns assert f"{column}_lo" in summary.columns From 2d4d08ae1b1c5ac4021bf0adcf971d5a7c931c6c Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 11:11:06 +0200 Subject: [PATCH 23/25] Cleanups: one config resolver, honest loss names, no slow marker - resolve_encoder_cfg now calls resolve_categorical_embedding_config instead of reimplementing the same inference inline. The shared helper had been reachable only from tests, so the tested implementation was the dead one. - Renamed the region-accumulation locals: total_losses, step_total_loss, mean_total_loss and _region_total_loss_buffer all held the region loss alone once the auxiliary terms were removed, so the names implied a composition that no longer exists. The rename immediately exposed a test that had been setting a stale attribute name. - Dropped the `slow` marker. The six marked tests run in ~13s, so `-k "not slow"` bought nothing and left `make test` and `make test-full` as two commands that ran different subsets for no reason. - EncoderConfig.categorical_embedding is typed Dict[str, Any] like its neighbours instead of a bare dict. Co-Authored-By: Claude Opus 5 (1M context) --- Makefile | 5 +-- pyproject.toml | 3 -- src/grass_mil/models/components/backbones.py | 4 +- src/grass_mil/models/supervised_module.py | 42 ++++++++++---------- src/grass_mil/models/training/builders.py | 29 ++++++-------- tests/eval/test_predict_cli.py | 1 - tests/models/test_runtime.py | 6 +-- tests/test_eval.py | 1 - tests/test_train.py | 4 -- 9 files changed, 39 insertions(+), 56 deletions(-) diff --git a/Makefile b/Makefile index 0ce3e7b..df29c6f 100644 --- a/Makefile +++ b/Makefile @@ -20,10 +20,7 @@ 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 diff --git a/pyproject.toml b/pyproject.toml index 3b70fcf..95b20ff 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,9 +54,6 @@ filterwarnings = [ "ignore::UserWarning", ] log_cli = "True" -markers = [ - "slow: slow tests", -] minversion = "6.0" testpaths = "tests/" diff --git a/src/grass_mil/models/components/backbones.py b/src/grass_mil/models/components/backbones.py index 1667db8..05178b1 100644 --- a/src/grass_mil/models/components/backbones.py +++ b/src/grass_mil/models/components/backbones.py @@ -2,7 +2,7 @@ import warnings from dataclasses import dataclass -from typing import Literal, Optional +from typing import Any, Dict, Literal, Optional import torch from torch import nn @@ -43,7 +43,7 @@ class EncoderConfig: # 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] = None + categorical_embedding: Optional[Dict[str, Any]] = None def _get_activation(name: str) -> nn.Module: diff --git a/src/grass_mil/models/supervised_module.py b/src/grass_mil/models/supervised_module.py index a59dfcb..be447c5 100644 --- a/src/grass_mil/models/supervised_module.py +++ b/src/grass_mil/models/supervised_module.py @@ -100,7 +100,7 @@ def __init__( self.region_accum_hyperbatch_size = 0 self.region_accum_flush_on_epoch_end = False - 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. @@ -757,7 +757,7 @@ def training_step(self, batch, batch_idx): patch_logits = payload["patch_logits"] region_losses: list[torch.Tensor] = [] - total_losses: list[torch.Tensor] = [] + buffered_region_losses: list[torch.Tensor] = [] for idx, bag_id in enumerate(ordered_bag_ids): region_w = bag_weights[idx : idx + 1] if bag_weights is not None else None region_loss = self._compute_losses( @@ -771,15 +771,15 @@ def training_step(self, batch, batch_idx): stage="train", ) region_losses.append(region_loss.detach()) - total_losses.append(region_loss) + 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: @@ -792,22 +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() 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, ) - return mean_total_loss + return mean_region_loss_value def validation_step(self, batch, batch_idx): out = self._shared_step(batch, "val") @@ -900,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() @@ -910,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) @@ -938,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 @@ -948,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/grass_mil/models/training/builders.py b/src/grass_mil/models/training/builders.py index 75bcbe5..57b5158 100644 --- a/src/grass_mil/models/training/builders.py +++ b/src/grass_mil/models/training/builders.py @@ -1,11 +1,13 @@ 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, @@ -132,23 +134,16 @@ def resolve_encoder_cfg( cfg["input_dim"] = max(int(inferred_input_dim), 0) if categorical_cfg is not None: - resolved = dict(categorical_cfg) - label = str(resolved.get("label", "cell_type")) - if module is not None and ( - resolved.get("num_embeddings") is None or resolved.get("column_index") is None - ): - num_embeddings, column_index = infer_categorical_binding(module, label) - if resolved.get("num_embeddings") is None: - resolved["num_embeddings"] = num_embeddings - if resolved.get("column_index") is None: - resolved["column_index"] = column_index - if resolved.get("num_embeddings") is None: - raise ValueError( - f"Could not infer the vocabulary size for categorical label '{label}' from the " - "datamodule. Set model.encoder.categorical_embedding.num_embeddings explicitly, " - "or include the label in data.categorical_features.include_labels." - ) - cfg["categorical_embedding"] = resolved + 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 diff --git a/tests/eval/test_predict_cli.py b/tests/eval/test_predict_cli.py index c775add..05f9384 100644 --- a/tests/eval/test_predict_cli.py +++ b/tests/eval/test_predict_cli.py @@ -24,7 +24,6 @@ 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 diff --git a/tests/models/test_runtime.py b/tests/models/test_runtime.py index 9e8b3bb..570a2f0 100644 --- a/tests/models/test_runtime.py +++ b/tests/models/test_runtime.py @@ -595,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), @@ -604,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/test_eval.py b/tests/test_eval.py index d66be8b..f1c2d9d 100644 --- a/tests/test_eval.py +++ b/tests/test_eval.py @@ -9,7 +9,6 @@ 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_train.py b/tests/test_train.py index 595859b..dfad5b1 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -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. From b5eb829d79002352ff8d6afb586031ad890c470d Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 11:12:49 +0200 Subject: [PATCH 24/25] Update docs to the current code and delete requirements.txt The docs still described removed features, and the hyperparameter reference handed users copy-pasteable commands that now fail: task.node_aux.*, task.entropy_reg.*, task.loss_weights.*, set2set/attention pooling, configs/model/heads/node.yaml, and the attention_attribution plugin. requirements.txt is deleted. It duplicated pyproject.toml's dependency list and was what two docs told users to install, which bypassed the extras and would drift silently. Install instructions now point at the package with its extras. Also records the survival preset, the per-class attention requirement, and the two things deliberately retained: region accumulation (the final NSCLC runs use hyperbatch_size 16) and the clustering algorithm's own parameter names. Co-Authored-By: Claude Opus 5 (1M context) --- docs/capabilities_and_modes.md | 8 +++---- docs/hyperparameter_reference.md | 18 +------------- docs/interpretability_suite.md | 8 +++---- docs/legacy_gap_audit.md | 10 ++++---- docs/method_fidelity.md | 3 +-- docs/repository_map.md | 3 +-- docs/workflows_train_infer.md | 29 ++++------------------ requirements.txt | 41 -------------------------------- 8 files changed, 21 insertions(+), 99 deletions(-) delete mode 100644 requirements.txt diff --git a/docs/capabilities_and_modes.md b/docs/capabilities_and_modes.md index 9edb981..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 diff --git a/docs/hyperparameter_reference.md b/docs/hyperparameter_reference.md index 232f91a..bbb7002 100644 --- a/docs/hyperparameter_reference.md +++ b/docs/hyperparameter_reference.md @@ -242,16 +242,6 @@ Includes all keys above plus MIL-specific controls: | `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` | -| `task.node_aux.enabled` | `false` | bool | `true/false` | MIL + binary training stage | `grass-mil-train task.node_aux.enabled=true` | -| `task.node_aux.target_mode` | `attention_shaped_ti` | str | `graph_label`, `attention_shaped_ti` | node auxiliary targets | `grass-mil-train task.node_aux.target_mode=graph_label` | -| `task.node_aux.loss_mode` | `bce` | str | `bce`, `weighted_bce` | node auxiliary loss | `grass-mil-train task.node_aux.loss_mode=weighted_bce` | -| `task.node_aux.weight` | `1.0` | float | any float | node auxiliary weighting | `grass-mil-train task.node_aux.weight=0.2` | -| `task.entropy_reg.enabled` | `false` | bool | `true/false` | MIL + binary training stage | `grass-mil-train task.entropy_reg.enabled=true` | -| `task.entropy_reg.mode` | `attention_shaped_target` | str | `attention`, `attention_shaped_target` | entropy regularization | `grass-mil-train task.entropy_reg.mode=attention` | -| `task.entropy_reg.weight` | `0.01` | float | any float | entropy weighting | `grass-mil-train task.entropy_reg.weight=0.005` | -| `task.loss_weights.region` | `1.0` | float | any float | MIL total loss composition | `grass-mil-train task.loss_weights.region=1.0` | -| `task.loss_weights.node_aux` | `1.0` | float | any float | MIL total loss composition | `grass-mil-train task.loss_weights.node_aux=0.5` | -| `task.loss_weights.entropy` | `0.0` | float | any float | MIL total loss composition | `grass-mil-train task.loss_weights.entropy=0.1` | ## 3.3 `configs/task/pretrain_bgrl.yaml` @@ -297,8 +287,7 @@ Includes all keys above plus MIL-specific controls: | `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`,`attention`,`set2set` | graph embedding | `grass-mil-train model.encoder.pooling=set2set` | -| `model.encoder.set2set_steps` | `3` | int | `>=1` | `pooling=set2set` | `grass-mil-train model.encoder.set2set_steps=4` | +| `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` | @@ -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: @@ -729,9 +716,6 @@ These constraints are not just documentation conventions; they are enforced in r | `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` | -| `task.node_aux.target_mode` must be `graph_label` or `attention_shaped_ti` | `src/grass_mil/models/training/loss_utils.py` | raises `ValueError` | -| `task.node_aux.loss_mode` must be `bce` or `weighted_bce` | `src/grass_mil/models/training/loss_utils.py` | raises `ValueError` | -| `task.entropy_reg.mode` must be `attention` or `attention_shaped_target` | `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 | diff --git a/docs/interpretability_suite.md b/docs/interpretability_suite.md index 22e321d..6c3e529 100644 --- a/docs/interpretability_suite.md +++ b/docs/interpretability_suite.md @@ -102,7 +102,7 @@ Common optional columns: - `instance_attention_logits` export requires shape `(N, 1)`; multi-column attention logits are rejected as ambiguous. -Note: `niche_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): @@ -272,7 +272,7 @@ Plugin registration is per pipeline run (fresh `PluginRegistry`). Built-in plugins: - `niche_profiles` -- `attention_attribution` +- `margin_attribution` - `neighborhood_enrichment` - `diff_neighborhood_enrichment` - `filtration_curves` @@ -448,7 +448,7 @@ With spatial plugins: grass-mil-report \ data.instance_table=/abs/path/instance_table.csv \ data.spatial_table=/abs/path/spatial_table.csv \ - plugins.enabled="[niche_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 niche labels: @@ -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: diff --git a/docs/legacy_gap_audit.md b/docs/legacy_gap_audit.md index 2c4a343..752a552 100644 --- a/docs/legacy_gap_audit.md +++ b/docs/legacy_gap_audit.md @@ -66,9 +66,9 @@ survived: - Native PyG operators are preferred; adapters stay thin. - Optional components instantiate only when their config is set. -## Still open +## Deliberately retained -- `attention_attribution` and `margin_attribution` overlap. Consolidating them - changes the payload keys `build_composite_cluster_heatmap` consumes. -- Region accumulation is retained. It is throughput mechanism rather than - method, but removing it changes optimizer-step semantics for streamed bags. +- 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 index 807aa88..abe44ab 100644 --- a/docs/method_fidelity.md +++ b/docs/method_fidelity.md @@ -17,7 +17,7 @@ listed at the bottom. | 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** | `models/components/attention.py`, `models/training/bagging.py` | `test_bagging.py::test_attention_is_normalised_per_class_over_instances` | +| 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` | @@ -34,7 +34,6 @@ The interpretability claim rests on one identity. If `identity_residual` is not | 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` | -| Attention lift | `λ = (Σ A) / prevalence`, 1 = neutral | same | `test_attention_lift_is_neutral_when_attention_matches_abundance` | | Uncertainty | 200-resample percentile bootstrap **over regions** | `percentile_bootstrap_ci` | `test_bootstrap_ci_brackets_the_point_estimate` | ## Spatial statistics diff --git a/docs/repository_map.md b/docs/repository_map.md index f2309a5..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: diff --git a/docs/workflows_train_infer.md b/docs/workflows_train_infer.md index 63e9c66..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 ``` @@ -371,8 +371,7 @@ grass-mil-train \ 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 @@ -386,9 +385,10 @@ grass-mil-train task=finetune_mil \ 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 grass-mil-train task=finetune_mil \ @@ -396,25 +396,6 @@ grass-mil-train task=finetune_mil \ task.region_accumulation.hyperbatch_size=8 ``` -Node auxiliary loss: - -```bash -grass-mil-train 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 -grass-mil-train 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 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) From 85f21368b8a883fddf36468e4cc7f036faf75ff1 Mon Sep 17 00:00:00 2001 From: Lovro Rabuzin Date: Thu, 13 Aug 2026 11:34:45 +0200 Subject: [PATCH 25/25] Fix macOS CI: install torch-sparse from conda, not pip The macOS job failed building torch-sparse from source. The compile error is in PyTorch's own header -- c10/util/strong_type.h specializes std::is_arithmetic, which Xcode 26's libc++ marks _LIBCPP_NO_SPECIALIZATIONS and rejects outright -- so on that runner *any* pip-built torch C++ extension fails, not just this one. torch-sparse was the only PyG extension routed through pip; pytorch-scatter was already coming from the pyg conda channel and installed cleanly as a prebuilt osx-arm64 binary. The pyg channel publishes pytorch-sparse 0.6.18 py312_torch_2.4.0_cpu for osx-arm64 too, so it now comes from there as well. Nothing is dropped: macOS keeps full coverage, including the shadow_custom sampler tests that need torch-sparse. Also clears the warnings the run surfaced: actions/checkout and setup-python bumped for the Node 20 runner deprecation, codecov-action to v5, and setup-miniconda's renamed auto-activate input. setup-miniconda@v3 still targets Node 20 upstream; nothing to do on our side. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/code-quality-main.yaml | 4 ++-- .github/workflows/code-quality-pr.yaml | 4 ++-- .github/workflows/test.yml | 14 +++++++------- environment.ci-spatial.yaml | 7 ++++++- 4 files changed, 17 insertions(+), 12 deletions(-) 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 22caa9b..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 @@ -80,7 +80,7 @@ jobs: 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 @@ -89,10 +89,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.11" 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