feat: integrate auxiliary pathway loss and sparsity regularization

Implement key mathematical and architectural improvements to the spatial
transcriptomics pipeline, focusing on pathway-driven interpretability.

Core Changes:
- Add Auxiliary Pathway Loss to ground the model's latent space in biological
  priors using MSigDB pathway membership markers.
- Introduce L1 sparsity regularization (--sparsity-lambda) on reconstruction
  weights to prune redundant pathway-to-gene mappings.
- Implement --log-transform data strategy (log1p) to stabilize MSE against
  high-variance housekeeping gene expression.

Architecture & Refactoring:
- Update dataloader and train/predict scripts to support auxiliary objectives
  and log-transformed targets.
- Refactor visualization modules to improve histology-pathway overlays.
- Remove deprecated Bowel Cancer download script.

Documentation & Compliance:
- Formally document auxiliary loss and training objectives in the IP Statement.
- Expand model documentation and training guides to reflect bottleneck
  transitions and latent discovery capabilities.
- Add Future Directions and Clinical Collaborations to project README.
- Resolve markdownlint error in LICENSE by adding a top-level H1 heading.

Testing:
- Add unit tests for pathway pruning logic and auxiliary loss calculation.
- Strengthen visualization tests to ensure coordinate alignment integrity.

Author: BenjaminIsaac0111

CONTRIBUTING.md

@@ -35,18 +35,23 @@ We use `black` for formatting and `flake8` for linting. Please ensure your code
 
 ```bash
 black .
-flake8 src/
 ```
 
-### 3. Testing
+### 4. AI-Assisted Development
 
-All new features must include unit tests in the `tests/` directory. We use `pytest` for our test suite.
+We welcome contributions developed with the assistance of AI tools (e.g., Copilot, ChatGPT, Claude, or agentic frameworks). However, to ensure the long-term maintainability and integrity of the project:
 
-```bash
-# Run all tests
-.\test.ps1  # Windows
-bash test.sh # Linux
-```
+- **Ownership**: You are ultimately responsible for the code you submit. Do not commit code you do not fully understand.
+- **Explainability**: During the review process, you must be able to explain the logic, design decisions, and any subtle side effects of the AI-suggested changes.
+- **Verification**: AI-generated code must strictly follow our coding standards, naming conventions, and architectural patterns. It must be accompanied by robust tests (see our [Testing Guide](docs/TESTING.md)).
+
+### 3. Testing & Quality Assurance
+
+All new features must be accompanied by relevant tests in the `tests/` directory natively using `pytest`.
+
+We highly encourage rigorous testing approaches such as **Mutation Testing** (via `cosmic-ray`) for critical model components to prevent surviving mutants.
+
+For full details on our testing requirements, how to run the test suites locally, and our guidelines on mutation testing, please read the [Testing Guide](docs/TESTING.md).
 
 ## Pull Request Process
 

Download-BowelCancer.ps1

@@ -1,4 +1,4 @@
-PROPRIETARY SOURCE CODE LICENSE (NON-COMMERCIAL + NEGOTIATED COMMERCIAL)
+# PROPRIETARY SOURCE CODE LICENSE (NON-COMMERCIAL + NEGOTIATED COMMERCIAL)
 
 Copyright (c) 2026 Benjamin Isaac Wilson. All rights reserved.
 

LICENSE

@@ -3,13 +3,14 @@
 > [!WARNING]
 > **Work in Progress**: This project is under active development. Core architectures, CLI flags, and data formats are subject to major changes.
 
-A transformer-based model for spatial transcriptomics that bridges histology and biological pathways.
+**SpatialTranscriptFormer** bridges histology and biological pathways through a high-performance transformer architecture. By modeling the dense interplay between morphological features and gene expression signatures, it provides an interpretable and spatially-coherent mapping of the tissue microenvironment.
 
 ## Key Features
 
 - **Quad-Flow Interaction**: Configurable attention between Pathways and Histology patches (`p2p`, `p2h`, `h2p`, `h2h`).
 - **Pathway Bottleneck**: Interpretable gene expression prediction via 50 MSigDB Hallmark tokens.
 - **Spatial Pattern Coherence**: Optimized using a composite **MSE + PCC (Pearson Correlation) loss** to prevent spatial collapse and ensure accurate morphology-expression mapping.
+- **Foundation Model Ready**: Native support for **CTransPath**, **Phikon**, **Hibou**, and **GigaPath**.
 - **Biologically Informed Initialization**: Gene reconstruction weights derived from known hallmark memberships.
 
 ## License
@@ -35,15 +36,24 @@ This project requires [Conda](https://docs.conda.io/en/latest/).
 
 ## Usage
 
-### Download HEST Data
+### Dataset Access
 
-Download specific subsets using filters or patterns:
+The model uses the **HEST1k** dataset. You can download specific subsets (by organ, technology, etc.) or the entire dataset using the `stf-download` utility:
 
 ```bash
-# Download only the Bowel Cancer subset (including ST data and WSIs)
-stf-download --organ Bowel --disease Cancer --local_dir hest_data
+# List available filtering options
+stf-download --list-options
+
+# Download a specific subset (e.g., Breast Cancer samples from Visium)
+stf-download --organ Breast --disease Cancer --tech Visium --local_dir hest_data
+
+# Download all human samples
+stf-download --species "Homo sapiens" --local_dir hest_data
 ```
 
+> [!NOTE]
+> The HEST dataset is gated on Hugging Face. Ensure you have accepted the terms at [MahmoodLab/hest](https://huggingface.co/datasets/MahmoodLab/hest) and are logged in via `huggingface-cli login`.
+
 ### Train Models
 
 We provide presets for baseline models and scaled versions of the SpatialTranscriptFormer.
@@ -95,6 +105,13 @@ Visualization plots will be saved to the `./results` directory.
 .\test.ps1
 ```
 
+## Future Directions & Clinical Collaborations
+
+A major future direction for **SpatialTranscriptFormer** is to integrate this architecture into an **end-to-end pipeline for patient risk assessment** and prognosis tracking. By leveraging the model's predicted expression and pathway activations, we aim to build a downstream risk prediction module that allows users to directly evaluate how spatially-resolved expression relates to patient survival.
+
+> [!NOTE]
+> **Call for Collaborators:** Rigorous risk assessment models require vast datasets of clinical metadata and survival outcomes, which we currently lack access to. We are open to investigating *any* disease of interest! If you have access to large clinical cohorts and are interested in exploring how spatial pathway activation correlates with patient prognosis, we would love to partner with you.
+
 ## Contributing
 
 We welcome contributions! Please see [CONTRIBUTING.md](CONTRIBUTING.md) for details on our coding standards and the process for submitting pull requests. Note that this project is under a proprietary license; contributions involve an assignment of rights for non-academic use.

README.md

@@ -14,6 +14,18 @@ The primary innovation is the **multimodal bottleneck transformer** designed for
 - **Quadrant-Based Interaction Masking**: The logic used to zero out specific attention quadrants (e.g., $A_{H \to H}$) to optimize memory while maintaining multimodal context.
 - **Biologically-Informed Reconstruction Bottleneck**: The specific matrix decomposition approach where gene expression is reconstructed from a linear combination of pathway activations.
 
+### Proposed Auxiliary Pathway Loss
+
+To prevent bottleneck collapse and provide a direct gradient signal to the pathway tokens, we use the `AuxiliaryPathwayLoss`. This loss compares the model's internal pathway scores against "ground truth" pathway activations computed from the gene expression targets via MSigDB membership.
+
+The total objective becomes:
+$$\mathcal{L} = \mathcal{L}_{gene} + \lambda_{aux} (1 - \text{PCC}(\text{pathway\_scores}, \text{target\_pathways}))$$
+
+The `--log-transform` flag applies `log1p` to targets, mitigating the heavy-tailed gene expression distribution where housekeeping genes dominate MSE.
+
+The full training objective with pathway sparsity regularisation:
+$$\mathcal{L} = \mathcal{L}_{task} + \lambda \|W_{recon}\|_1$$
+
 ## 2. Spatial Context Methodologies
 
 - **Euclidean-Gated Attention**: The implementation of spatial distance-based masking ($M_{spatial}$) to constrain model focus to local morphological regions.

docs/IP_STATEMENT.md

@@ -122,7 +122,7 @@ Together, these ensure the model learns *spatially-varying* pathway activation m
 
 #### Frozen Backbone (Feature Extraction)
 
-Pre-computed features from a pathology foundation model. The backbone is never fine-tuned.
+Pre-computed features from a pathology foundation model. (The backbone is never fine-tuned, though this might change!)
 
 | Backbone | Feature Dim | Source |
 | :--- | :--- | :--- |
@@ -214,7 +214,7 @@ The Zero-Inflated Negative Binomial (ZINB) loss is designed for raw, highly disp
 
 The model outputs these parameters, and the loss computes the negative log-likelihood of the ground truth counts given this distribution.
 
-### Auxiliary Pathway Loss
+### Proposed Auxiliary Pathway Loss
 
 To prevent bottleneck collapse and provide a direct gradient signal to the pathway tokens, we use the `AuxiliaryPathwayLoss`. This loss compares the model's internal pathway scores against "ground truth" pathway activations computed from the gene expression targets via MSigDB membership.
 

docs/MODELS.md

@@ -21,49 +21,55 @@ After the model makes predictions (N spots x G genes), we run a statistical test
 - **Tool**: `gseapy` or a custom mapping script.
 - **Use Case**: Generating a "Pathway Activation Map" from a trained model's output.
 
-### B. Pathway Bottleneck (Model Architecture)
+### B. Interaction Model via Multi-Task Learning (MTL)
 
-The **SpatialTranscriptFormer** replaces the standard linear output head with a two-step projection that can be configured in two modes:
+The **SpatialTranscriptFormer** interaction model inherently represents pathway activations as part of its attention mechanism and output process. Rather than a simple linear bottleneck, it utilizes learnable pathway tokens and Multi-Task Learning (MTL).
 
-#### 1. Informed Projection (Prior Knowledge)
+#### 1. Informed Supervision via Auxiliary Loss
 
-In this mode, the **Gene Reconstruction Matrix** $\mathbf{W}_{recon}$ is guided by established biological databases (MSigDB, KEGG).
+In this mode, the network receives direct supervision on its pathway tokens, guided by established biological databases (e.g., MSigDB):
 
-- **Implementation**: $\mathbf{W}_{recon}$ is initialized as a binary mask $M \in \{0, 1\}^{G \times P}$ where $M_{gk} = 1$ if gene $g$ belongs to pathway $k$.
-- **Benefit**: Predictions are guaranteed to be linear combinations of known biological processes, making them instantly interpretable by clinicians.
+- **Architecture Flow**:
+    1. **Interaction**: Learnable pathway tokens $P$ interact with Histology patch features $H$ via self-attention (e.g., $p2h$, $h2p$).
+    2. **Activation**: Pathway scores $S \in \mathbb{R}^P$ are computed using a learnable temperature-scaled cosine similarity between the pathway tokens and image patch tokens.
+    3. **Gene Reconstruction**: $\hat{y} = S \cdot \mathbf{W}_{recon} + b$, where $\mathbf{W}_{recon}$ is initialized using the binary pathway membership matrix $M$.
+- **MTL Auxiliary Loss**: To prevent standard bottleneck collapse, an explicit auxiliary loss bridges the spatial representations directly to biological data. The pathway scores $S$ are supervised against a pathway ground truth ($Y_{genes} \cdot M^T$) using a Pearson Correlation Coefficient (PCC) loss.
+  $$L_{total} = L_{gene} + \lambda_{pathway} (1 - PCC(S, Y_{genes} \cdot M^T))$$
+- **Benefit**: The model is forced to explicitly align its internal interaction tokens with concrete biological pathways, granting direct interpretability.
 
-#### 2. Data-Driven Projection (Latent Discovery)
+#### 2. Data-Driven Discovery (Latent Projection)
 
-In this mode, the model learns its own "latent pathways" based on morphological patterns.
+In the absence of a biological prior, the model can learn its own "latent pathways".
 
-- **Implementation**: $\mathbf{W}_{recon}$ is randomly initialized and learned via backpropagation.
-- **Sparsity Constraint**: We apply an L1 penalty to force the model to identify "canonical" gene sets: $L_{total} = L_{MSE} + \lambda \|\mathbf{W}_{recon}\|_1$.
+- **Implementation**: $\mathbf{W}_{recon}$ is randomly initialized and the auxiliary pathway loss is disabled.
+- **Sparsity Constraint**: We apply an L1 penalty to force the model to identify "canonical" sparse gene sets: $L_{total} = L_{gene} + \lambda_{sparsity} \|\mathbf{W}_{recon}\|_1$.
 - **Benefit**: Can discover novel spatial-transcriptomic relationships that aren't yet captured in curated databases.
 
-- **Architecture Flow**:
-    1. **Interaction**: Pathway tokens $P$ query the Histology $H$.
-    2. **Activation**: A linear layer reduces $P_{tokens}$ to activation scores $S \in \mathbb{R}^P$.
-    3. **Reconstruction**: $\hat{y} = S \cdot \mathbf{W}_{recon} + b$.
+## 3. Generalizing to HEST1k Tissues
+
+The model supports any dataset within the HEST1k collection (e.g., Breast, Kidney, Lung, Colon). Instead of being bound to a single disease context, users can leverage the `--custom-gmt` flag to map genes to pathways relevant to their specific investigation.
 
-## 3. Clinical Application in Bowel Cancer
+### Example: Profiling the Tumor Microenvironment
 
-For colorectal cancer, we should prioritize monitoring these specific pathways:
+Regardless of the tissue of origin (e.g., Kidney versus Breast), researchers often track core functional states within the tumor microenvironment. A user might define a `.gmt` file to explicitly monitor:
 
-| Pathway | Clinically Relevant Genes | Clinical Significance |
+| Pathway Concept | Hallmarks / Relevant Genes | Interpretive Value across Tissues |
 | :--- | :--- | :--- |
-| **Wnt Signaling** | `CTNNB1`, `MYC`, `AXIN2` | Common driver in CRC (APC mutations) |
-| **MMR / DNA Repair** | `MLH1`, `MSH2`, `MSH6` | MSI vs MSS status (Immunotherapy response) |
-| **EMT** | `SNAI1`, `VIM`, `ZEB1` | Tumor invasion and metastasis risk |
-| **Angiogenesis** | `VEGFA`, `FLT1` | Potential for anti-angiogenic therapy |
+| **Hypoxia & Angiogenesis** | `VEGFA`, `FLT1`, `HIF1A` | Identifies oxygen-deprived or highly vascularized tumor cores. |
+| **Immune Infiltration** | `CD8A`, `GZMB`, `IFNG` | Maps regions of active anti-tumor immune response. |
+| **Stromal / EMT** | `VIM`, `SNAI1`, `ZEB1` | Highlights desmoplastic stroma and invasion fronts. |
+| **Proliferation** | `MKI67`, `PCNA`, `MYC` | Pinpoints highly active, dividing cell populations. |
+
+By supplying these functional groupings via `--custom-gmt`, the model's MTL process explicitly aligns its spatial interaction tokens to monitor these exact states across any whole-slide image in the HEST1k dataset.
 
 ## 4. Implementation Status
 
 ### Implemented
 
 - **MSigDB Hallmarks Initialization** (`--pathway-init` flag): Downloads the GMT file, matches genes against `global_genes.json`, and initializes `gene_reconstructor.weight` with the binary membership matrix. See [`pathways.py`](../src/spatial_transcript_former/data/pathways.py).
-  - 50 Hallmark pathways (fixed when using `--pathway-init`)
-  - ~54% gene coverage (542/1000 genes mapped to at least one pathway)
-  - GMT file cached in `.cache/` after first download
+  - 50 Hallmark pathways (default fixed fallback when using `--pathway-init`).
+  - GMT file cached in `.cache/` after first download.
+- **Custom Pathway Definitions** (`--custom-gmt` flag): Users can override the default Hallmarks by providing a URL or local path to a `.gmt` file, enabling custom database integrations (e.g., KEGG, Reactome, or highly specific tissue masks).
 
 - **Sparsity Regularization** (`--sparsity-lambda` flag): L1 penalty on `gene_reconstructor` weights to encourage pathway-like groupings when using data-driven (random) initialization.
 
@@ -79,8 +85,9 @@ python -m spatial_transcript_former.train \
     --model interaction --num-pathways 50 --sparsity-lambda 0.01 ...
 ```
 
+- **Spatial Pathway Maps**: Visualize pathway activations as spatial heatmaps overlaid on histology using `stf-predict`. See the [README](../README.md) for inference instructions.
+
 ### Future Work
 
-- **KEGG/Reactome**: More granular pathway databases for finer-grained analysis.
-- **Post-Hoc Enrichment**: `gseapy` integration for pathway activation maps from model outputs.
-- **Spatial Pathway Maps**: Visualize pathway activations as spatial heatmaps overlaid on histology.
+- **Post-Hoc Enrichment**: `gseapy` integration for pathway activation maps from model outputs without architectural bottlenecks.
+- **End-to-End Risk Assessment Module**: Developing a downstream prediction system that takes the spatially-resolved pathway activations and gene expressions derived from the model and maps them directly to clinical risk and survival outcomes.