Refactor: Centralize data management and reorganize test suite

Data Management Framework Refactoring:
- Created default `GeneVocab` class in `data/gene_vocab.py` to act as the single source of truth for gene list handling, replacing duplicated loading logic.
- Centralized path resolution into `data/paths.py` (`resolve_feature_dir`, `resolve_gene_vocab_path`), eliminating hardcoded dataset paths.
- Extracted HEST-specific patient-aware dataset splitting logic from core train.py into `recipes/hest/utils.py::get_train_val_ids()`, decoupling the training script from benchmark-specific data properties.
- Updated predict.py to handle `anndata >= 0.12` strict shape enforcement by dynamically rebuilding the AnnData object in-place when filtering variables.

Test Suite Restructuring & Bug Fixes:
- Consolidated 37 flat, root-level test files into 15 focused module files organized across `tests/models/`, `tests/training/`, `tests/data/`, and `tests/recipes/hest/`.
- Resolved 6 shadowed test name collisions that caused tests to silently overwrite each other during collection (e.g., `test_interaction_output_shape`).
- Fixed broken `download_hest` imports in the `test_io.py` script by correctly resolving the root `scripts/` directory via `sys.path`.
- Fixed `test_pathways.py` Mock failures by explicitly defining `args.pathways = None` to enforce default filtering behavior.
- Total active test suite passes locally (199/199 collected tests).

Author: BenjaminIsaac0111

docs/SVG_PROTOTYPE_RESULTS.md

@@ -0,0 +1,57 @@
+# Spatially Variable Gene (SVG) Selection & Validation Prototype Results
+
+## Overview
+
+This document summarizes the prototype results for integrating spatially variable gene (SVG) scoring into the SpatialTranscriptFormer pipeline. Two major components were built and validated:
+
+1. **SVG-Aware Gene Selection**: Using Moran's I to bias the 1000-gene vocabulary towards genes that exhibit strong spatial autocorrelation, improving the biological relevance of the model's bottleneck.
+2. **Spatial Coherence Validation**: An end-to-end validation metric that calculates Moran's I on the predicted vs. ground-truth expression vectors for the top-50 spatially variable genes, reported as a Pearson correlation score.
+
+This prototype was run on a colorectal/bowel cancer cohort from the HEST-1k dataset (84 human samples) for ~400 epochs.
+
+---
+
+## 1. Training Progress & Stability
+
+Training the `stf_small` model (4 layers, 384 dim, 8 heads) on the bowel cancer subset with SVG weighting (`--svg-weight 0.5`) showed steady convergence and stability.
+
+![Training Loss Landscape](./assets/bowel_svg_loss_curve.png)
+
+* **Learning Schedule**: 10 warm-up epochs followed by cosine annealing.
+* **Overfitting**: The gap between training and validation loss is characteristic of the small dataset size (84 samples), but the model finds a strong minimum.
+* **Best Checkpoint**: Reached a best validation loss of **1.611 at epoch 367**, an improvement over early training stages (~1.835 at epoch 130).
+
+![Validation Metrics](./assets/bowel_svg_val_metrics.png)
+
+* **Mean Absolute Error (MAE)**: Decreased smoothly from ~0.80 down strictly towards 0.60, indicating improved pixel-wise prediction accuracy.
+* **Prediction Variance**: Evaluated as a collapse detector. The variance ramped up from near-zero to ~0.03, confirming the model successfully escaped "mean-prediction" collapse and learned to output highly differentiated spatial patterns. *(Note: A variance of 0.0 indicates a collapsed model predicting the same average value everywhere).*
+
+---
+
+## 2. Pathway Spatial Coherence (Bowel Cancer)
+
+To validate the biological plausibility of the predictions, we visualised clinically relevant Hallmarks pathways for colorectal/bowel cancer using the best epoch checkpoint (epoch 367) on sample `TENX29`.
+
+Pathways selected for their relevance to colorectal cancer progression, invasion, and tumor microenvironment:
+
+* **Wnt/β-Catenin Signaling**: Mutated in >80% of sporadic colorectal cancers (APC mutation).
+* **Epithelial-Mesenchymal Transition (EMT)**: Key marker of invasion and metastasis.
+* **TNF-α Signaling via NF-κB** & **Inflammatory Response**: Chronic inflammation is a hallmark driver of CRC.
+* **KRAS Signaling (Up)**: KRAS mutations occurring in ~40% of CRCs.
+* **Angiogenesis**: Critical for tumor vasculature, the target of anti-VEGF therapies.
+
+### Pathway Predictions vs Ground Truth (Epoch 524)
+
+![Bowel Cancer Pathway Predictions](./assets/TENX29_epoch_524.png)
+
+**Observations:**
+
+1. **Spatial Pattern Matching**: The model successfully reconstructs complex, heterogeneous spatial patterns across the tissue architecture. High-expression regions (yellow/green) in the ground truth strongly correlate with high-expression regions in the predictions.
+2. **Biological Gradients**: Crucially, the model does not just predict average expression but captures the relative spatial *gradients* of these pathways, confirming the health of the non-collapsed prediction variance metric.
+3. **Tumor Microenvironment (TME)**: Inflammatory pathways (TNF-α, Inflammatory Response) show distinct spatial localization differing from tumor-intrinsic pathways (Wnt, KRAS), accurately reflecting TME heterogeneity.
+
+---
+
+## Conclusion
+
+The integration of Moran's I for both **SVG-aware gene selection** and **Spatial Coherence Validation** provides a significantly more robust, biologically grounded training pipeline. The model demonstrates the ability to learn and reconstruct clinically relevant spatial pathway patterns from H&E imaging alone, laying a strong foundation for scaling the model to the full HEST-1k dataset and evaluating patient-level stratification.

docs/TRAINING_GUIDE.md

@@ -145,6 +145,17 @@ python scripts/run_preset.py --preset stf_medium
 python scripts/run_preset.py --preset stf_large
 ```
 
+#### Disease-Specific Priors (Colorectal Cancer)
+
+To learn representations specifically constrained to phenotypes of a target disease, you can explicitly filter the initialization pathway bottleneck using the `--pathways` argument. The `crc` presets demonstrate this by shrinking the dimensionality down from 50 generic hallmarks to 14 CRC-specific pathways (e.g. Wnt/Beta-catenin, EMT, Angiogenesis).
+
+```bash
+# Small CRC Variant (14 explicit pathways)
+python scripts/run_preset.py --preset stf_crc_small
+```
+
+These presets are defined directly in `scripts/run_preset.py`, serving as a template for how you can introduce your own biological priors for other diseases.
+
 ### Choosing Interaction Modes
 
 By default, the model runs in **Full Interaction** mode (`p2p p2h h2p h2h`) where all token types attend to each other. You can selectively disable interactions using the `--interactions` flag for ablation or to enforce specific architectural constraints.

scripts/batch_qc_stats.py

@@ -0,0 +1,59 @@
+import os
+import h5py
+import numpy as np
+import matplotlib.pyplot as plt
+from scipy.sparse import csr_matrix
+
+def calculate_qc_stats(h5ad_path, min_umis=500, min_genes=200, max_mt=0.15):
+    with h5py.File(h5ad_path, "r") as f:
+        barcodes = [b.decode("utf-8") if isinstance(b, bytes) else b for b in f["obs"]["_index"][:]]
+        gene_names = [g.decode("utf-8") if isinstance(g, bytes) else g for g in f["var"]["_index"][:]]
+        X_group = f["X"]
+        if isinstance(X_group, h5py.Group):
+            X = csr_matrix((X_group["data"][:], X_group["indices"][:], X_group["indptr"][:]), shape=(len(barcodes), len(gene_names)))
+        else:
+            X = f["X"][:]
+
+    n_counts = np.array(X.sum(axis=1)).flatten()
+    n_genes = np.array((X > 0).sum(axis=1)).flatten()
+    mt_genes = [i for i, name in enumerate(gene_names) if "mt-" in name.lower() or "mt:" in name.lower()]
+    if mt_genes:
+        mt_counts = np.array(X[:, mt_genes].sum(axis=1)).flatten()
+        pct_counts_mt = mt_counts / (n_counts + 1e-9)
+    else:
+        pct_counts_mt = np.zeros_like(n_counts)
+
+    keep_mask = (n_counts >= min_umis) & (n_genes >= min_genes) & (pct_counts_mt <= max_mt)
+    return len(barcodes), np.sum(keep_mask), np.sum(n_counts < min_umis), np.sum(n_genes < min_genes), np.sum(pct_counts_mt > max_mt)
+
+if __name__ == "__main__":
+    st_dir = r"A:\hest_data\st"
+    
+    # Get all h5ad files
+    samples = [f.replace(".h5ad", "") for f in os.listdir(st_dir) if f.endswith(".h5ad")]
+    samples.sort()
+    
+    print(f"Analyzing {len(samples)} samples...")
+    print(f"{'Sample':<15} | {'Total':<6} | {'Kept':<6} | {'% Kept':<8} | {'Low UMI':<8} | {'Low Gene':<8} | {'High MT':<8}")
+    print("-" * 80)
+    
+    all_results = []
+    for s in samples:
+        path = os.path.join(st_dir, f"{s}.h5ad")
+        try:
+            total, kept, l_umi, l_gene, h_mt = calculate_qc_stats(path)
+            pct = kept / total if total > 0 else 0
+            print(f"{s:<15} | {total:<6} | {kept:<6} | {pct:7.1%} | {l_umi:<8} | {l_gene:<8} | {h_mt:<8}")
+            all_results.append((total, kept, pct))
+        except Exception as e:
+            print(f"Error processing {s}: {e}")
+
+    if all_results:
+        avg_kept = np.mean([r[2] for r in all_results])
+        min_kept = np.min([r[2] for r in all_results])
+        max_kept = np.max([r[2] for r in all_results])
+        print("-" * 80)
+        print(f"GLOBAL SUMMARY ({len(all_results)} samples):")
+        print(f"Average Kept: {avg_kept:.1%}")
+        print(f"Minimum Kept: {min_kept:.1%}")
+        print(f"Maximum Kept: {max_kept:.1%}")

scripts/diagnose_qc.py

@@ -0,0 +1,123 @@
+import os
+import h5py
+import numpy as np
+import matplotlib.pyplot as plt
+from scipy.sparse import csr_matrix
+
+def diagnose_qc(h5ad_path, output_path, min_umis=500, min_genes=200, max_mt=0.15):
+    print(f"Loading {h5ad_path}...")
+    
+    with h5py.File(h5ad_path, "r") as f:
+        # Load barcodes and gene names
+        barcodes = [b.decode("utf-8") if isinstance(b, bytes) else b for b in f["obs"]["_index"][:]]
+        gene_names = [g.decode("utf-8") if isinstance(g, bytes) else g for g in f["var"]["_index"][:]]
+        
+        # Load expression matrix X
+        X_group = f["X"]
+        if isinstance(X_group, h5py.Group):
+            data = X_group["data"][:]
+            indices = X_group["indices"][:]
+            indptr = X_group["indptr"][:]
+            X = csr_matrix((data, indices, indptr), shape=(len(barcodes), len(gene_names)))
+        else:
+            X = f["X"][:]
+            
+        # Spatial coordinates (usually in obsm/spatial)
+        if "obsm" in f and "spatial" in f["obsm"]:
+            coords = f["obsm"]["spatial"][:]
+        else:
+            # Fallback if spatial not found
+            coords = np.zeros((len(barcodes), 2))
+            print("Warning: Spatial coordinates not found in obsm/spatial")
+
+    # Calculate QC metrics
+    n_counts = np.array(X.sum(axis=1)).flatten()
+    n_genes = np.array((X > 0).sum(axis=1)).flatten()
+    
+    # MT fraction
+    # Robust MT detection: check for mt- or MT- anywhere, but prioritize common patterns
+    mt_genes = [i for i, name in enumerate(gene_names) if "mt-" in name.lower() or "mt:" in name.lower()]
+    if mt_genes:
+        print(f"Found {len(mt_genes)} mitochondrial genes.")
+        mt_counts = np.array(X[:, mt_genes].sum(axis=1)).flatten()
+        pct_counts_mt = mt_counts / (n_counts + 1e-9)
+    else:
+        pct_counts_mt = np.zeros_like(n_counts)
+        print("No mitochondrial genes found.")
+
+    # Filter mask
+    pass_umi = n_counts >= min_umis
+    pass_gene = n_genes >= min_genes
+    pass_mt = pct_counts_mt <= max_mt
+    
+    keep_mask = pass_umi & pass_gene & pass_mt
+    
+    total_spots = len(barcodes)
+    kept_spots = np.sum(keep_mask)
+    
+    print(f"Total spots: {total_spots}")
+    print(f"Kept spots: {kept_spots} ({kept_spots/total_spots:.1%})")
+    print(f"Filtered: {total_spots - kept_spots}")
+    print(f"  - Low UMI (<{min_umis}): {np.sum(~pass_umi)}")
+    print(f"  - Low Genes (<{min_genes}): {np.sum(~pass_gene)}")
+    print(f"  - High MT (>{max_mt:.1%}): {np.sum(~pass_mt)}")
+
+    # Plotting
+    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
+    
+    # 1. UMI vs Genes
+    axes[0, 0].scatter(n_counts, n_genes, c=keep_mask, cmap="RdYlGn", alpha=0.5, s=10)
+    axes[0, 0].axvline(min_umis, color="red", linestyle="--", label=f"Min UMI={min_umis}")
+    axes[0, 0].axhline(min_genes, color="blue", linestyle="--", label=f"Min Genes={min_genes}")
+    axes[0, 0].set_xlabel("Total UMI counts")
+    axes[0, 0].set_ylabel("Number of detected genes")
+    axes[0, 0].set_title("QC: UMI vs Genes")
+    axes[0, 0].legend()
+    
+    # 2. MT Fraction distribution
+    axes[0, 1].hist(pct_counts_mt, bins=50, color="gray", alpha=0.7)
+    axes[0, 1].axvline(max_mt, color="red", linestyle="--", label=f"Max MT={max_mt:.0%}")
+    axes[0, 1].set_xlabel("Mitochondrial Fraction")
+    axes[0, 1].set_ylabel("Count")
+    axes[0, 1].set_title("QC: MT Fraction Distribution")
+    axes[0, 1].legend()
+    
+    # 3. Spatial: Before (Total)
+    axes[1, 0].scatter(coords[:, 0], coords[:, 1], c="lightgray", s=15, label="All Spots")
+    axes[1, 0].scatter(coords[keep_mask, 0], coords[keep_mask, 1], c="green", s=15, label="Pass QC")
+    axes[1, 0].set_title("Spatial Distribution: Kept vs Filtered")
+    axes[1, 0].set_aspect("equal")
+    axes[1, 0].legend()
+    
+    # 4. Summary Table (as text)
+    stats_text = (
+        f"Sample: {os.path.basename(h5ad_path)}\n\n"
+        f"Total Spots: {total_spots}\n"
+        f"Kept Spots: {kept_spots} ({kept_spots/total_spots:.1%})\n"
+        f"Filtered: {total_spots - kept_spots}\n\n"
+        f"Thresholds:\n"
+        f"Min UMI: {min_umis}\n"
+        f"Min Genes: {min_genes}\n"
+        f"Max MT: {max_mt:.0%}"
+    )
+    axes[1, 1].text(0.1, 0.5, stats_text, fontsize=14, family="monospace")
+    axes[1, 1].axis("off")
+    
+    plt.tight_layout()
+    plt.savefig(output_path)
+    print(f"Plot saved to {output_path}")
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--sample", type=str, default="MEND29")
+    args = parser.parse_args()
+    
+    sample_id = args.sample
+    h5ad_file = f"A:\\hest_data\\st\\{sample_id}.h5ad"
+    output_file = f"qc_diagnosis_{sample_id}.png"
+    
+    if not os.path.exists(h5ad_file):
+        print(f"Error: {h5ad_file} not found.")
+    else:
+        diagnose_qc(h5ad_file, output_file)

scripts/find_outliers.py

@@ -0,0 +1,32 @@
+import os
+import sys
+sys.path.append(r"z:\Projects\SpatialTranscriptFormer\scripts")
+from batch_qc_stats import calculate_qc_stats
+import numpy as np
+
+st_dir = r"A:\hest_data\st"
+samples = [f.replace(".h5ad", "") for f in os.listdir(st_dir) if f.endswith(".h5ad")]
+samples.sort()
+
+results = []
+for s in samples:
+    try:
+        total, kept, l_umi, l_gene, h_mt = calculate_qc_stats(os.path.join(st_dir, f"{s}.h5ad"))
+        results.append({
+            "sample": s,
+            "total": total,
+            "kept": kept,
+            "pct": kept / total if total > 0 else 0,
+            "low_umi": l_umi,
+            "low_gene": l_gene,
+            "high_mt": h_mt
+        })
+    except Exception as e:
+        print(f"Error {s}: {e}")
+
+results.sort(key=lambda x: x["pct"])
+
+print(f"{'Sample':<15} | {'Kept %':<8} | {'Filtered':<10} | {'Low UMI':<8} | {'Low Gene':<8} | {'High MT':<8}")
+print("-" * 75)
+for r in results[:15]:
+    print(f"{r['sample']:<15} | {r['pct']:7.1%} | {r['total']-r['kept']:<10} | {r['low_umi']:<8} | {r['low_gene']:<8} | {r['high_mt']:<8}")

scripts/run_preset.py

@@ -5,6 +5,24 @@
 
 from spatial_transcript_former.config import get_config
 
+# Curated list of MSigDB Hallmarks with strong evidence of involvement in Colorectal/Bowel Cancer
+CRC_PATHWAYS = [
+    "HALLMARK_WNT_BETA_CATENIN_SIGNALING",
+    "HALLMARK_TGF_BETA_SIGNALING",
+    "HALLMARK_KRAS_SIGNALING_UP",
+    "HALLMARK_KRAS_SIGNALING_DN",
+    "HALLMARK_PI3K_AKT_MTOR_SIGNALING",
+    "HALLMARK_EPITHELIAL_MESENCHYMAL_TRANSITION",
+    "HALLMARK_ANGIOGENESIS",
+    "HALLMARK_APICAL_JUNCTION",
+    "HALLMARK_INFLAMMATORY_RESPONSE",
+    "HALLMARK_IL6_JAK_STAT3_SIGNALING",
+    "HALLMARK_APOPTOSIS",
+    "HALLMARK_P53_PATHWAY",
+    "HALLMARK_DNA_REPAIR",
+    "HALLMARK_HYPOXIA",
+]
+
 
 def make_stf_params(n_layers: int, token_dim: int, n_heads: int, batch_size: int):
     """Helper to create standard SpatialTranscriptFormer parameters."""
@@ -14,6 +32,7 @@ def make_stf_params(n_layers: int, token_dim: int, n_heads: int, batch_size: int
         "precomputed": True,
         "whole-slide": True,
         "pathway-init": True,
+        "pathway-loss-weight": 0.5,
         "use-amp": True,
         "log-transform": True,
         "loss": "mse_pcc",
@@ -55,6 +74,11 @@ def make_stf_params(n_layers: int, token_dim: int, n_heads: int, batch_size: int
     "stf_small": make_stf_params(n_layers=4, token_dim=384, n_heads=8, batch_size=8),
     "stf_medium": make_stf_params(n_layers=6, token_dim=512, n_heads=8, batch_size=8),
     "stf_large": make_stf_params(n_layers=12, token_dim=768, n_heads=12, batch_size=8),
+    # --- Biologically-Prioritized Variants (e.g. Colorectal Cancer) ---
+    "stf_crc_tiny": {**make_stf_params(2, 256, 4, 8), "pathways": CRC_PATHWAYS},
+    "stf_crc_small": {**make_stf_params(4, 384, 8, 8), "pathways": CRC_PATHWAYS},
+    "stf_crc_medium": {**make_stf_params(6, 512, 8, 8), "pathways": CRC_PATHWAYS},
+    "stf_crc_large": {**make_stf_params(12, 768, 12, 8), "pathways": CRC_PATHWAYS},
 }
 
 
@@ -67,6 +91,9 @@ def params_to_args(params_dict):
             args.append(arg_name)
         elif value is False or value is None:
             continue
+        elif isinstance(value, list) or isinstance(value, tuple):
+            args.append(arg_name)
+            args.extend([str(v) for v in value])
         else:
             args.extend([arg_name, str(value)])
     return args

src/spatial_transcript_former/data/__init__.py

@@ -2,10 +2,13 @@
 Data abstractions for SpatialTranscriptFormer.
 
 Core exports:
+    - :class:`GeneVocab` — single source of truth for gene vocabulary
     - :class:`SpatialDataset` — abstract base class implementing the data contract
     - :func:`apply_dihedral_augmentation` — D4 coordinate augmentation
     - :func:`apply_dihedral_to_tensor` — D4 image augmentation
     - :func:`normalize_coordinates` — auto-normalise spatial coordinates
+    - :func:`resolve_feature_dir` — centralised feature directory discovery
+    - :func:`resolve_gene_vocab_path` — find ``global_genes.json``
 
 HEST-specific exports (backward compatibility — prefer ``recipes.hest``):
     - :class:`HEST_Dataset`, :func:`get_hest_dataloader`
@@ -20,3 +23,5 @@
     apply_dihedral_to_tensor,
     normalize_coordinates,
 )
+from .gene_vocab import GeneVocab
+from .paths import resolve_feature_dir, resolve_gene_vocab_path