diff --git a/CITATION.cff b/CITATION.cff new file mode 100644 index 0000000..d1729e3 --- /dev/null +++ b/CITATION.cff @@ -0,0 +1,24 @@ +cff-version: 1.2.0 +message: "If you use this software in academic work, please cite it as below." +title: SpatialTranscriptFormer +abstract: >- + A transformer framework that predicts spatially-resolved biological pathway + activity directly from H&E histology. Learnable pathway tokens cross-attend to + histology patch features (a configurable quad-flow interaction), and a + dot-product + Softplus head produces per-spot pathway-activity scores supervised + against offline-computed targets. +type: software +authors: + - family-names: Wilson + given-names: "Benjamin Isaac" +version: "0.1.0" +license: Apache-2.0 +repository-code: "https://github.com/BenjaminIsaac0111/SpatialTranscriptFormer" +url: "https://github.com/BenjaminIsaac0111/SpatialTranscriptFormer" +keywords: + - spatial transcriptomics + - computational pathology + - deep learning + - pathway activity + - histology + - transformer diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d3f5ccd..d0c682d 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -7,13 +7,11 @@ Thank you for your interest in contributing! As a project at the intersection of > [!IMPORTANT] > This project is a **Work in Progress**. We are actively refining the core interaction logic and scaling behaviors. Expect breaking changes in the CLI and data schemas. -## Intellectual Property & Licensing +## Licensing of Contributions -SpatialTranscriptFormer is protected under a **Proprietary Source Code License**. +SpatialTranscriptFormer is released under the **[Apache License 2.0](LICENSE)**. -- **Academic/Non-Profit**: We encourage contributions from the research community. Contributions made under an academic affiliation are generally welcome. -- **Commercial/For-Profit**: Contributions from commercial entities or individuals intended for profit-seeking use require a separate agreement. -- **Assignment**: By submitting a Pull Request, you agree that your contributions will be licensed under the project's existing license, granting the author the right to include them in both the open-access and proprietary versions of the software. +By submitting a contribution (e.g. a Pull Request), you agree that your contribution is licensed to the project and its users under the same Apache-2.0 license ("inbound = outbound"), and you confirm you have the right to submit it under that license. Please don't include third-party code or data whose license is incompatible with Apache-2.0. ## Development Workflow diff --git a/LICENSE b/LICENSE index e7339c2..d645695 100644 --- a/LICENSE +++ b/LICENSE @@ -1,36 +1,202 @@ -# PROPRIETARY SOURCE CODE LICENSE (NON-COMMERCIAL + NEGOTIATED COMMERCIAL) -Copyright (c) 2026 Benjamin Isaac Wilson. All rights reserved. + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ -1. PERMITTED USE -This software and its source code are provided for the following purposes: - a. Evaluation: For prospective employers and collaborators to assess the Author's technical expertise. - b. Academic Research: For students, researchers, and faculty at academic institutions for the purpose of scientific investigation. - c. Non-Profit: For use by registered non-profit organizations in pursuit of their non-commercial missions. + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION -2. COMMERCIAL USE (NEGOTIATED AGREEMENT REQUIRED) -Use of this software, its source code, or derivative works by for-profit entities, or for any commercial or revenue-generating purposes, is strictly prohibited without a separate, written license agreement negotiated with and signed by the Author. + 1. Definitions. -Foundation Model Disclaimer: This software is designed to be backbone-agnostic and relies on third-party foundation models for feature extraction. A commercial agreement with the Author covers ONLY the SpatialTranscriptFormer source code, architecture, and intellectual property. It DOES NOT grant commercial use rights for restricted third-party models (e.g., CTransPath or Phikon). To use this software commercially, you must pair it with a backbone model that has a compatible commercial-friendly license (e.g., Apache 2.0 models like Hibou, Virchow, or H-Optimus-0). + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. -Dataset Disclaimer: This project utilizes third-party datasets, primarily HEST-1k. These datasets are subject to their own respective licenses (e.g., CC BY-NC-SA 4.0 for HEST-1k). A commercial agreement with the Author DOES NOT grant any rights to use, redistribute, or commercially exploit restricted third-party data. Commercial users must ensure they have independent legal rights to any training or evaluation data used within this framework. + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. -MSigDB Disclaimer: The Pathway Bottleneck and Reconstruction components utilize MSigDB Hallmark gene sets (Broad Institute). Content from MSigDB (v6.0-v7.5.1, v2022.1+) is protected by copyright (c) 2004-2025 Broad Institute, Inc., MIT, and Regents of the University of California, and is provided under CC BY 4.0. Users are responsible for complying with attribution requirements and any third-party restrictions (e.g., KEGG, BioCarta) that may apply to specific subsets of the data. + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. -Interested parties should contact the Author directly to discuss commercial licensing terms. + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. -1. ATTRIBUTION -Any publication, presentation, or public report that uses this software or its results must provide clear attribution to the Author (Benjamin Isaac Wilson) and reference this repository. + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. -2. RESTRICTIONS -Unless explicitly authorized in writing by the Author, the following actions are strictly prohibited: - a. Redistribution: You may not redistribute, sell, or publish the source code or any derivative works to third parties. - b. Sublicensing: You may not grant sublicenses to others. + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. -3. OWNERSHIP -All intellectual property rights, including but not limited to copyright, patents, and trade secrets, remain solely with the Author. Core architectural components, original methodologies, and academic attributions are documented in the [IP Statement](docs/IP_STATEMENT.md). + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). -4. NO WARRANTY -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. -For inquiries regarding commercial licensing or further collaboration, please contact the Author directly. + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/NOTICE b/NOTICE new file mode 100644 index 0000000..b09a3a0 --- /dev/null +++ b/NOTICE @@ -0,0 +1,38 @@ +SpatialTranscriptFormer +Copyright 2026 Benjamin Isaac Wilson + +Licensed under the Apache License, Version 2.0 (the "License"); you may not use +this file except in compliance with the License. A copy of the License is +provided in the LICENSE file and at http://www.apache.org/licenses/LICENSE-2.0 + +================================================================================ +Third-party components and data +================================================================================ + +This framework integrates with, and/or interoperates with, third-party +software, model weights, and datasets that are governed by their own licenses. +These are NOT redistributed in this repository: they are downloaded or supplied +by the user at runtime and remain subject to their original terms. The Apache +2.0 license over this source code does not grant any rights to them. + +- HEST-1k dataset (Mahmood Lab) — CC BY-NC-SA 4.0 (non-commercial, share-alike). + Used here for research/benchmarking only. Users must obtain independent rights + for any commercial or non-research use. + +- MSigDB Hallmark gene sets (Broad Institute) — copyright (c) 2004-2025 Broad + Institute, Inc., MIT, and the Regents of the University of California; provided + under CC BY 4.0 (attribution required). Individual gene sets may carry + additional third-party terms (e.g. KEGG, BioCarta). + +- Pathology foundation-model backbones (e.g. CTransPath, Phikon, GigaPath, + Hibou) — each subject to its own license. Some are gated or restricted and may + prohibit commercial use. Select a backbone whose license fits your use case. + +================================================================================ +Conceptual attribution +================================================================================ + +The multimodal pathway-histology interaction framing was conceptually inspired +by SURVPATH (Jaume et al., 2024; https://arxiv.org/abs/2303.15545). No SURVPATH +source code is used, copied, or adapted in this project; the implementation here +is original. diff --git a/README.md b/README.md index 9d9dc7c..3d26b09 100644 --- a/README.md +++ b/README.md @@ -28,10 +28,10 @@ predictor = Predictor(model, device="cuda") # coords: (N, 2) tensor of spatial coordinates (from your WSI tiling) features = extractor.extract_batch(image_patches, batch_size=64) # → (N, 768) -# 3. Predict pathway activity scores from extracted features -predictions = predictor.predict_wsi(features, coords) # → (1, P) +# 3. Predict per-spot pathway activity from extracted features +predictions = predictor.predict_wsi(features, coords, return_dense=True) # → (1, N, P) -# 4. Integrate with Scanpy +# 4. Integrate with Scanpy (one pathway-activity vector per spot) inject_predictions(adata, coords, predictions[0], pathway_names=model.pathway_names) ``` @@ -50,18 +50,18 @@ For more details, see the **[Python API Reference](docs/API.md)**. ## License -This project is protected by a **Proprietary Source Code License**. See the [LICENSE](LICENSE) file for full details. +SpatialTranscriptFormer is released under the **[Apache License 2.0](LICENSE)** — you are free to use, modify, and redistribute it (including commercially), provided you retain the copyright/license notices and state significant changes. See [LICENSE](LICENSE) and [NOTICE](NOTICE). -- ✅ **Permitted**: Evaluation for employment, Academic Research, and Non-Profit use. -- 🤝 **For-Profit Use**: Permitted only under a **negotiated agreement** with the author. - - **Note on Foundation Models**: This architecture is backbone-agnostic. Any negotiated commercial agreement covers *only* the SpatialTranscriptFormer source code and IP. It does **not** grant commercial rights to use restricted third-party weights (e.g., CTransPath, Phikon). To use this system commercially, you must select a foundation model with a compatible open or commercial license (e.g., Hibou, Virchow, or H-Optimus-0). - - **Note on HEST-1k Dataset**: The benchmark data used in this project is sourced from the **HEST-1k** dataset (Mahmood Lab), which is licensed under **CC BY-NC-SA 4.0**. This data is strictly for non-commercial research and cannot be used for commercial training or clinical deployment without explicit permission from the original authors. - - **Note on MSigDB**: This project uses data from the **Molecular Signatures Database (MSigDB)** (versions v6.0–v7.5.1, and v2022.1+). The contents are protected by copyright © 2004–2025 Broad Institute, Inc., MIT, and Regents of the University of California, and are distributed under the **CC BY 4.0** license. While this allows for commercial use, users must provide appropriate attribution. Note that individual gene sets within MSigDB may be subject to additional terms from third-party sources (e.g., KEGG). -- ❌ **Prohibited**: Redistribution and unauthorized commercial exploitation. +> [!IMPORTANT] +> Apache-2.0 covers **this repository's source code only**. It does **not** grant rights to the third-party components this framework relies on at runtime, which keep their own licenses and are not redistributed here: +> +> - **Foundation-model backbones** (e.g. CTransPath, Phikon, GigaPath, Hibou) — each has its own license; some are gated or prohibit commercial use. Choose one whose license fits your use case (e.g. Apache-2.0 models such as Hibou, Virchow, or H-Optimus-0 for commercial use). +> - **HEST-1k dataset** (Mahmood Lab) — **CC BY-NC-SA 4.0** (non-commercial, share-alike): research/benchmarking only; obtain independent rights for commercial or clinical use. +> - **MSigDB Hallmark gene sets** (Broad Institute, v6.0–v7.5.1 / v2022.1+) — **CC BY 4.0**, attribution required; some subsets carry extra terms (e.g. KEGG). © 2004–2025 Broad Institute, Inc., MIT, and the Regents of the University of California. -## Intellectual Property +## Attribution & Provenance -The core architectural innovations, including the **SpatialTranscriptFormer** interaction logic and spatial masking strategies, are the unique Intellectual Property of the author. For a detailed breakdown, see the [IP Statement](docs/IP_STATEMENT.md). +This is original work by Benjamin Isaac Wilson. The pathway↔histology interaction framing was **conceptually inspired by** SURVPATH (Jaume et al., 2024) — no SURVPATH source code is used or adapted. For the design contributions and third-party attributions, see the [Attribution & Design Notes](docs/IP_STATEMENT.md) and [NOTICE](NOTICE). If you use this work in academic research, please cite this repository and its author. --- @@ -105,7 +105,7 @@ python scripts/run_preset.py --preset stf_small ### 4. Inference & Visualization ```bash -stf-predict --data-dir A:\hest_data --sample-id MEND29 --model-path checkpoints/best_model.pth --model-type interaction +stf-predict --run-dir checkpoints --sample-id MEND29 --output-dir results ``` Visualization plots and spatial pathway activation maps will be saved to the `./results` directory. For the full guide, see the **[HEST Recipe Docs](src/spatial_transcript_former/recipes/hest/README.md)**. @@ -168,4 +168,4 @@ Active research and development is tracked in the **[Research & Improvement Road ## 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. +We welcome contributions! Please see [CONTRIBUTING.md](CONTRIBUTING.md) for details on our coding standards and the process for submitting pull requests. Contributions are accepted under the project's Apache-2.0 license (inbound = outbound). diff --git a/docs/API.md b/docs/API.md index 8f5048d..fb43355 100644 --- a/docs/API.md +++ b/docs/API.md @@ -1,6 +1,6 @@ # Python API Reference -The SpatialTranscriptFormer package exposes a clean API for training, inference, and integration with the Scanpy/AnnData ecosystem. +The SpatialTranscriptFormer package exposes a clean API for training, inference, and integration with the Scanpy/AnnData ecosystem. The model predicts spatially-resolved **biological pathway activity scores** (not gene expression) from histology features. ```python from spatial_transcript_former import ( @@ -8,8 +8,8 @@ from spatial_transcript_former import ( Trainer, # High-level training orchestrator Predictor, # Inference wrapper FeatureExtractor, # Backbone feature extraction - save_pretrained, # Save checkpoint directory - load_pretrained, # Load checkpoint directory + save_pretrained, # Save a checkpoint directory + load_pretrained, # Load a checkpoint directory inject_predictions, # AnnData integration ) ``` @@ -21,26 +21,27 @@ from spatial_transcript_former import ( ### End-to-End Inference (New Data) ```python +import pandas as pd +import scanpy as sc from spatial_transcript_former import SpatialTranscriptFormer, Predictor, FeatureExtractor from spatial_transcript_former.predict import inject_predictions -import scanpy as sc -# 1. Load model from checkpoint directory +# 1. Load a model from a checkpoint directory (config.json + model.pth) model = SpatialTranscriptFormer.from_pretrained("./checkpoints/my_run/") -print(model.gene_names[:3]) # ['TP53', 'EGFR', 'MYC'] +print(model.pathway_names[:3]) # e.g. ['HALLMARK_HYPOXIA', ...] or None # 2. Extract features from raw patches extractor = FeatureExtractor(backbone="phikon", device="cuda") features = extractor.extract_batch(image_tensor, batch_size=64) # (N, 768) -# 3. Predict gene expression +# 3. Predict per-spot pathway activity from the features predictor = Predictor(model, device="cuda") -predictions = predictor.predict_wsi(features, coords) # (1, G) +preds = predictor.predict_wsi(features, coords, return_dense=True) # (1, N, P) -# 4. Inject into AnnData for Scanpy analysis -adata = sc.AnnData(obs=pd.DataFrame(index=[f"spot_{i}" for i in range(N)])) -inject_predictions(adata, coords, predictions[0], gene_names=model.gene_names) -sc.pl.spatial(adata, color="TP53") +# 4. Inject into AnnData for Scanpy analysis (one activity vector per spot) +adata = sc.AnnData(obs=pd.DataFrame(index=[f"spot_{i}" for i in range(len(coords))])) +inject_predictions(adata, coords, preds[0], pathway_names=model.pathway_names) +sc.pl.spatial(adata, color="HALLMARK_HYPOXIA") ``` ### Saving a Trained Model @@ -49,16 +50,16 @@ sc.pl.spatial(adata, color="TP53") from spatial_transcript_former import save_pretrained # After training, export a self-contained checkpoint -save_pretrained(model, "./release/v1/", gene_names=gene_list) +save_pretrained(model, "./release/v1/", pathway_names=pathway_list) ``` This creates: ``` release/v1/ -├── config.json # Architecture parameters -├── model.pth # Model weights (state_dict) -└── gene_names.json # Ordered gene symbols +├── config.json # Architecture parameters (+ pathway_format_version) +├── model.pth # Model weights (state_dict) +└── pathway_names.json # Ordered pathway names (optional) ``` --- @@ -67,24 +68,35 @@ release/v1/ ### `SpatialTranscriptFormer` -The core transformer model. Predicts gene expression from histology patch features and spatial coordinates. +The core transformer. Predicts pathway activity scores from histology patch features and spatial coordinates. -#### `SpatialTranscriptFormer.__init__(...)` +#### `SpatialTranscriptFormer(...)` | Parameter | Type | Default | Description | |---|---|---|---| -| `num_genes` | `int` | *required* | Number of output genes | -| `num_pathways` | `int` | `50` | Number of pathway bottleneck tokens | -| `backbone_name` | `str` | `"resnet50"` | Backbone identifier (`resnet50`, `phikon`, `ctranspath`, etc.) | +| `num_pathways` | `int` | `50` | Number of pathway tokens **and** the output dimension | +| `backbone_name` | `str` | `"resnet50"` | Backbone identifier (`resnet50`, `phikon`, `ctranspath`, `gigapath`, `hibou-b`, …) | | `pretrained` | `bool` | `True` | Load pretrained backbone weights | -| `token_dim` | `int` | `256` | Common embedding dimension | +| `token_dim` | `int` | `256` | Transformer embedding dimension | | `n_heads` | `int` | `4` | Number of attention heads | -| `n_layers` | `int` | `2` | Number of transformer layers | +| `n_layers` | `int` | `2` | Transformer layers (minimum 2 when `h2h` is disabled) | | `dropout` | `float` | `0.1` | Dropout probability | -| `pathway_init` | `Tensor` | `None` | `(P, G)` biological pathway membership matrix | -| `use_spatial_pe` | `bool` | `True` | Enable learned spatial positional encodings | -| `output_mode` | `str` | `"counts"` | Output head: `"counts"` (Softplus) or `"zinb"` (Zero-Inflated NB) | -| `interactions` | `list[str]` | all | Attention interactions: `p2p`, `p2h`, `h2p`, `h2h` | +| `use_spatial_pe` | `bool` | `True` | Enable learned 2-D spatial positional encodings (requires `rel_coords` at call time) | +| `interactions` | `list[str]` | all four | Attention quadrants to enable: `p2p`, `p2h`, `h2p`, `h2h` | + +> The model predicts pathway activity **directly** via a scaled dot-product + Softplus head — there is no `num_genes`, gene-reconstruction head, or `output_mode`. See [MODELS.md](MODELS.md). + +#### `forward(x, rel_coords=None, mask=None, return_dense=False, return_attention=False)` + +| Arg | Description | +|---|---| +| `x` | `(B, 3, H, W)` image patch, **or** `(B, S, D)` pre-computed features | +| `rel_coords` | `(B, S, 2)` slide-stationary coordinates (required when `use_spatial_pe=True`) | +| `mask` | `(B, S)` bool padding mask (`True` = padded/ignore) | +| `return_dense` | If `True`, return per-patch predictions instead of a pooled slide vector | +| `return_attention` | If `True`, also return a list of per-layer attention maps | + +**Returns:** `(B, num_pathways)`, or `(B, S, num_pathways)` if `return_dense=True`. If `return_attention=True`, returns `(scores, attentions)`. #### `SpatialTranscriptFormer.from_pretrained(checkpoint_dir, device="cpu", **kwargs)` @@ -92,16 +104,16 @@ Load a model from a checkpoint directory created by `save_pretrained`. ```python model = SpatialTranscriptFormer.from_pretrained("./checkpoint/", device="cuda") -model.gene_names # List[str] or None +model.pathway_names # list[str] or None ``` | Parameter | Type | Description | |---|---|---| -| `checkpoint_dir` | `str` | Path to directory with `config.json` + `model.pth` | +| `checkpoint_dir` | `str` | Directory with `config.json` + `model.pth` | | `device` | `str` | Torch device (`"cpu"`, `"cuda"`) | | `**kwargs` | | Override any `config.json` value (e.g. `dropout=0.0`) | -**Returns:** `SpatialTranscriptFormer` in eval mode with `.gene_names` attribute. +**Returns:** `SpatialTranscriptFormer` in eval mode with a `.pathway_names` attribute. --- @@ -109,50 +121,44 @@ model.gene_names # List[str] or None Stateful inference wrapper. Manages device placement, eval mode, and optional AMP. -#### `Predictor.__init__(model, device="cpu", use_amp=False)` +#### `Predictor(model, device="cpu", use_amp=False)` ```python predictor = Predictor(model, device="cuda", use_amp=True) +predictor.pathway_names # forwarded from the model, or None ``` -#### `Predictor.predict_patch(image, return_pathways=False)` +#### `Predictor.predict_patch(image)` Single-patch inference from a raw image tensor. ```python result = predictor.predict_patch(image) # image: (1, 3, 224, 224) or (3, 224, 224) -# result: (1, num_genes) +# result: (1, num_pathways) ``` -> **Note:** When the model uses spatial PE, a zero-coordinate is automatically injected — no need to provide coordinates for single patches. +> When the model uses spatial PE, a zero coordinate is injected automatically — no coordinates are needed for single patches. -#### `Predictor.predict_wsi(features, coords, return_pathways=False, return_dense=False)` +#### `Predictor.predict_wsi(features, coords, return_dense=False)` -Whole-slide inference from pre-extracted feature embeddings. +Whole-slide inference from pre-extracted feature embeddings. Coordinates are re-centred/standardised internally to match the training-time slide-stationary scaling. ```python -# Global prediction (one vector per slide) -result = predictor.predict_wsi(features, coords) # (1, G) - -# Dense prediction (one vector per patch) -result = predictor.predict_wsi(features, coords, return_dense=True) # (1, N, G) +result = predictor.predict_wsi(features, coords) # (1, P) — pooled slide vector +result = predictor.predict_wsi(features, coords, return_dense=True) # (1, N, P) — per-spot ``` | Parameter | Type | Description | |---|---|---| | `features` | `Tensor` | `(N, D)` or `(1, N, D)` embeddings | | `coords` | `Tensor` | `(N, 2)` or `(1, N, 2)` spatial coordinates | -| `return_pathways` | `bool` | Also return pathway scores | -| `return_dense` | `bool` | Per-patch predictions instead of global | +| `return_dense` | `bool` | Per-patch predictions instead of a pooled slide vector | -> **Validation:** Raises `ValueError` with a clear message if the feature dimension doesn't match the model's expected backbone dimension. +> Raises `ValueError` with a clear message if the feature dimension doesn't match the model's expected backbone dimension. #### `Predictor.predict(features, coords=None, **kwargs)` -Unified entry point — auto-dispatches: - -- 4D tensor `(B, 3, H, W)` → `predict_patch` -- 2D tensor `(N, D)` → `predict_wsi` (requires `coords`) +Unified entry point — dispatches a `(B, 3, H, W)` image to `predict_patch`, otherwise to `predict_wsi` (requires `coords`). --- @@ -160,50 +166,56 @@ Unified entry point — auto-dispatches: Wraps a backbone model and its normalization transform for one-line feature extraction. -#### `FeatureExtractor.__init__(backbone="resnet50", device="cpu", pretrained=True, transform=None)` +#### `FeatureExtractor(backbone="resnet50", device="cpu", pretrained=True, transform=None)` ```python extractor = FeatureExtractor(backbone="phikon", device="cuda") -extractor.feature_dim # 768 -extractor.backbone_name # "phikon" +extractor.feature_dim # 768 +extractor.backbone_name # "phikon" ``` | Backbone | `feature_dim` | Source | |---|---|---| | `resnet50` | 2048 | torchvision | -| `ctranspath` | 768 | HuggingFace (CTransPath) | -| `phikon` | 768 | Owkin Phikon (HuggingFace) | +| `ctranspath` | 768 | CTransPath | +| `phikon` | 768 | Owkin Phikon | | `vit_b_16` | 768 | torchvision | -| `gigapath` | 1536 | ProvGigaPath *(gated)* | +| `gigapath` | 1536 | Prov-GigaPath *(gated)* | | `hibou-b` | 768 | Hibou-B *(gated)* | | `hibou-l` | 1024 | Hibou-L *(gated)* | #### `extractor(images)` / `extractor.extract_batch(images, batch_size=64)` ```python -features = extractor(images) # (N, D) — all at once -features = extractor.extract_batch(images, batch_size=64) # batched, returns on CPU +features = extractor(images) # (N, D) — all at once, on device +features = extractor.extract_batch(images, batch_size=64) # (N, D) — batched, returned on CPU ``` -Images should be float tensors in `[0, 1]` range, shape `(N, 3, H, W)`. +Images should be float tensors in `[0, 1]`, shape `(N, 3, H, W)`. --- -### `save_pretrained(model, save_dir, gene_names=None)` +### `save_pretrained(model, save_dir, pathway_names=None)` Save a self-contained checkpoint directory. ```python -save_pretrained(model, "./release/v1/", gene_names=["TP53", "EGFR", ...]) +save_pretrained(model, "./release/v1/", pathway_names=["HALLMARK_HYPOXIA", ...]) ``` | Parameter | Type | Description | |---|---|---| | `model` | `SpatialTranscriptFormer` | Trained model instance | | `save_dir` | `str` | Output directory (created if needed) | -| `gene_names` | `list[str]` | Optional ordered gene symbols (must match `num_genes`) | +| `pathway_names` | `list[str]` | Optional ordered pathway names (must match `num_pathways`) | + +**Raises:** `ValueError` if `pathway_names` length doesn't match `num_pathways`. -**Raises:** `ValueError` if `gene_names` length doesn't match `num_genes`. +### `load_pretrained(checkpoint_dir, device="cpu", **override_kwargs)` + +Reconstruct a model from `config.json` + `model.pth` (and optional `pathway_names.json`). `**override_kwargs` override config values. **Raises** `ValueError` if the checkpoint's `pathway_format_version` doesn't match the current preprocessing format (re-train against current targets). + +> `SpatialTranscriptFormer.from_pretrained(...)` is a thin wrapper around `load_pretrained`. ### AnnData & Scanpy — A Primer @@ -214,78 +226,61 @@ If you're coming from a pure deep-learning background, AnnData and Scanpy may be An `AnnData` object is a structured container for observations × variables matrices, designed for genomics. Think of it as a spreadsheet with labelled sidecars: ``` - var (genes) - ┌──────────────────┐ - │ TP53 EGFR MYC │ - ┌────┼──────────────────┤ - obs │ s0 │ 0.3 1.2 0.8 │ ← adata.X (the main data matrix) - (spots/ │ s1 │ 0.1 0.5 1.1 │ - cells) │ s2 │ 0.9 0.2 0.4 │ - └────┴──────────────────┘ + var (pathways) + ┌──────────────────────┐ + │ HYPOXIA EMT MYC │ + ┌────┼──────────────────────┤ + obs │ s0 │ 0.3 1.2 0.8 │ ← adata.X (the main data matrix) + (spots)│ s1 │ 0.1 0.5 1.1 │ + │ s2 │ 0.9 0.2 0.4 │ + └────┴──────────────────────┘ ``` | Slot | What it stores | Our usage | |---|---|---| -| `adata.X` | Main matrix `(N, G)` | Predicted gene expression | +| `adata.X` | Main matrix `(N, P)` | Predicted pathway activity | | `adata.obs` | Per-observation metadata | Spot/cell barcodes, cluster labels | -| `adata.var` | Per-variable metadata | Gene symbols as the index | +| `adata.var` | Per-variable metadata | Pathway names as the index | | `adata.obsm["spatial"]` | Observation-level embeddings | `(N, 2)` spatial coordinates | -| `adata.obsm["spatial_pathways"]` | Additional embeddings | `(N, P)` pathway scores | -| `adata.uns` | Unstructured metadata | Pathway names, model config | +| `adata.uns` | Unstructured metadata | Model config, run info | #### What is Scanpy? -[Scanpy](https://scanpy.readthedocs.io/) (`sc`) is the analysis library that operates on AnnData objects. Once predictions are inside an `adata`, you instantly get access to: +[Scanpy](https://scanpy.readthedocs.io/) (`sc`) is the analysis library that operates on AnnData objects. Once predictions are inside an `adata`, you get clustering, differential testing, spatial plotting, and trajectory analysis for free: ```python import scanpy as sc -# Spatial plotting — visualise gene expression on tissue coordinates -sc.pl.spatial(adata, color="TP53") - -# Clustering — find groups of spots with similar expression -sc.tl.leiden(adata) - -# Differential expression — find marker genes per cluster -sc.tl.rank_genes_groups(adata, groupby="leiden") - -# Dimensionality reduction -sc.tl.pca(adata) -sc.tl.umap(adata) -sc.pl.umap(adata, color="leiden") +sc.pl.spatial(adata, color="HALLMARK_HYPOXIA") # spatial pathway-activity map +sc.tl.leiden(adata) # cluster spots by activity profile +sc.tl.rank_genes_groups(adata, groupby="leiden") # marker pathways per cluster ``` -#### Why does this matter? - -By injecting predictions into AnnData, our model's output becomes **instantly compatible** with the entire Scanpy ecosystem — clustering, differential testing, spatial plotting, trajectory analysis — without any custom code. Biologists can take our predictions and run their standard workflows immediately. +By injecting predictions into AnnData, the model's output becomes instantly compatible with the entire Scanpy ecosystem. --- -### `inject_predictions(adata, coords, predictions, ...)` +### `inject_predictions(adata, coords, pathway_scores, pathway_names=None)` -Inject predictions into an AnnData object for Scanpy integration. +Inject pathway-activity predictions into an AnnData object for Scanpy integration. ```python inject_predictions( adata, - coords, # → adata.obsm["spatial"] - predictions, # → adata.X - gene_names=["TP53", "EGFR", ...], # → adata.var_names - pathway_scores=pathway_activations, # → adata.obsm["spatial_pathways"] - pathway_names=["APOPTOSIS", ...], # → adata.uns["pathway_names"] + coords, # → adata.obsm["spatial"] + pathway_scores, # → adata.X (N, P) + pathway_names=["HALLMARK_HYPOXIA", ...], # → adata.var_names ) ``` | Parameter | Type | Description | |---|---|---| -| `adata` | `AnnData` | Target AnnData object | +| `adata` | `AnnData` | Target object (must have `N` observations matching `coords`) | | `coords` | `ndarray` | `(N, 2)` spatial coordinates | -| `predictions` | `ndarray` | `(N, G)` gene predictions | -| `gene_names` | `list[str]` | Optional gene symbols | -| `pathway_scores` | `ndarray` | Optional `(N, P)` pathway scores | -| `pathway_names` | `list[str]` | Optional pathway names | +| `pathway_scores` | `ndarray` | `(N, P)` predicted pathway activity | +| `pathway_names` | `list[str]` | Optional `P` pathway names (set as `adata.var_names`) | -> **Lazy loading:** `anndata` is only imported when this function is called, so it's not required for basic inference. +> **Lazy loading:** `anndata` is imported only when this function is called, so it isn't required for basic inference. --- @@ -293,16 +288,15 @@ inject_predictions( ``` checkpoint/ -├── config.json # Architecture (JSON) -├── model.pth # Weights (PyTorch state_dict) -└── gene_names.json # Gene symbols (JSON array, optional) +├── config.json # Architecture (JSON) +├── model.pth # Weights (PyTorch state_dict) +└── pathway_names.json # Pathway names (JSON array, optional) ``` **`config.json` example:** ```json { - "num_genes": 460, "num_pathways": 50, "backbone_name": "phikon", "token_dim": 256, @@ -310,29 +304,24 @@ checkpoint/ "n_layers": 2, "dropout": 0.1, "use_spatial_pe": true, - "output_mode": "counts", - "interactions": ["h2h", "h2p", "p2h", "p2p"] + "interactions": ["h2h", "h2p", "p2h", "p2p"], + "pathway_format_version": 2 } ``` -**`gene_names.json` example:** - -```json -["TP53", "EGFR", "MYC", "BRCA1", ...] -``` +`pathway_format_version` records the preprocessing/target schema the model was trained against; `load_pretrained` refuses checkpoints whose version doesn't match the current pipeline. --- ## Training API -The training pipeline lives in the `spatial_transcript_former.training` subpackage. You can use it via the **CLI** or **programmatically** in your own scripts. +The training pipeline lives in the `spatial_transcript_former.training` subpackage. Use it via the **CLI** or **programmatically**. ### CLI Quick Start Training is launched via the `stf-train` entry point (or `python -m spatial_transcript_former.train`): ```bash -# Minimal: train on precomputed features with whole-slide mode stf-train \ --model interaction \ --backbone phikon \ @@ -340,116 +329,67 @@ stf-train \ --precomputed \ --whole-slide \ --use-spatial-pe \ - --pathway-init \ - --loss mse_pcc \ + --loss mse_ccc \ --epochs 100 \ --lr 1e-4 \ --warmup-epochs 10 - -# Resume from checkpoint -stf-train --model interaction --resume --output-dir ./checkpoints ``` -### CLI Arguments - -#### Data - -| Flag | Default | Description | -| --- | --- | --- | -| `--data-dir` | from config | Root HEST data directory | -| `--feature-dir` | auto | Explicit pre-extracted feature directory | -| `--num-genes` | 1000 | Number of output genes | -| `--precomputed` | off | Use pre-extracted backbone features | -| `--whole-slide` | off | Dense whole-slide prediction mode | -| `--organ` | all | Filter samples by organ type | -| `--max-samples` | all | Limit samples (for debugging) | - -#### Model - -| Flag | Default | Description | -| --- | --- | --- | -| `--model` | `he2rna` | Architecture: `interaction`, `he2rna`, `vit_st`, `attention_mil`, `transmil` | -| `--backbone` | `resnet50` | Backbone: `resnet50`, `phikon`, `ctranspath`, `vit_b_16`, etc. | -| `--num-pathways` | 50 | Pathway bottleneck tokens | -| `--token-dim` | 256 | Embedding dimension | -| `--n-heads` | 4 | Attention heads | -| `--n-layers` | 2 | Transformer layers | -| `--use-spatial-pe` | off | Learned spatial positional encoding | -| `--interactions` | all | Attention mask: `p2p p2h h2p h2h` | -| `--pathway-init` | off | Initialize gene head from MSigDB Hallmarks | - -#### Training - -| Flag | Default | Description | -| --- | --- | --- | -| `--epochs` | 10 | Total training epochs | -| `--batch-size` | 32 | Batch size | -| `--lr` | 1e-4 | Learning rate | -| `--warmup-epochs` | 10 | Linear warmup before cosine annealing | -| `--weight-decay` | 0.0 | AdamW weight decay | -| `--grad-accum-steps` | 1 | Gradient accumulation steps | -| `--use-amp` | off | Mixed precision (FP16) | -| `--compile` | off | `torch.compile` the model | -| `--resume` | off | Resume from latest checkpoint | - -#### Loss - -| Flag | Default | Description | -| --- | --- | --- | -| `--loss` | `mse_pcc` | Loss function: `mse`, `pcc`, `mse_pcc`, `zinb`, `poisson`, `logcosh` | -| `--pcc-weight` | 1.0 | PCC term weight in `mse_pcc` | -| `--pathway-loss-weight` | 0.0 | Auxiliary pathway PCC loss weight (0 = disabled) | - ---- +See the **[Training Guide](TRAINING_GUIDE.md)** for the full, authoritative flag reference and HEST recipes. Pathway targets must be pre-computed first with `stf-compute-pathways` (see [PATHWAY_MAPPING.md](PATHWAY_MAPPING.md)). ### Trainer (High-Level) -The `Trainer` class handles LR scheduling, AMP, checkpointing, logging, and early stopping: +`Trainer` handles LR scheduling (linear warmup → cosine), AMP, checkpointing, SQLite logging, and early stopping. All arguments after `criterion` are **keyword-only**. ```python from spatial_transcript_former import SpatialTranscriptFormer, Trainer from spatial_transcript_former.training import CompositeLoss, EarlyStoppingCallback -model = SpatialTranscriptFormer(num_genes=460, backbone_name="phikon", ...) +model = SpatialTranscriptFormer(num_pathways=50, backbone_name="phikon") trainer = Trainer( model=model, train_loader=train_dl, val_loader=val_dl, - criterion=CompositeLoss(alpha=1.0), + criterion=CompositeLoss(alpha=1.0, pcc_type="ccc"), epochs=100, warmup_epochs=10, device="cuda", output_dir="./checkpoints", + model_name="interaction", use_amp=True, + whole_slide=True, callbacks=[EarlyStoppingCallback(patience=15)], ) -results = trainer.fit() # returns {"best_val_loss", "history", ...} -trainer.save_pretrained("./release/v1/") # inference-ready export +results = trainer.fit() # {"best_val_loss", "epochs_completed", "history"} +trainer.save_pretrained("./release/v1/", pathway_names=pathway_list) ``` #### Trainer Parameters | Parameter | Default | Description | -| --- | --- | --- | +|---|---|---| | `model` | *required* | Any `nn.Module` | -| `train_loader` | *required* | Training `DataLoader` | -| `val_loader` | *required* | Validation `DataLoader` | +| `train_loader` / `val_loader` | *required* | Training / validation `DataLoader` | | `criterion` | *required* | Loss function | -| `optimizer` | `None` | Custom optimizer (default: `AdamW`) | -| `lr` | `1e-4` | Learning rate (if no custom optimizer) | -| `epochs` | `100` | Total training epochs | +| `optimizer` | `None` | Custom optimizer (default: `AdamW` from `lr` / `weight_decay`) | +| `lr` | `1e-4` | Learning rate (when no custom optimizer) | +| `weight_decay` | `0.0` | Weight decay (when no custom optimizer) | +| `epochs` | `100` | Total epochs | | `warmup_epochs` | `10` | Linear warmup before cosine annealing | +| `device` | `"cuda"` | Torch device | +| `output_dir` | `"./checkpoints"` | Directory for checkpoints/logs | +| `model_name` | `"model"` | Used in checkpoint filenames (`best_model_.pth`) | | `use_amp` | `False` | Mixed precision (FP16) | | `grad_accum_steps` | `1` | Gradient accumulation | -| `whole_slide` | `False` | Dense whole-slide mode | -| `output_dir` | `./checkpoints` | Directory for checkpoints/logs | -| `callbacks` | `[]` | List of `TrainerCallback` instances | -| `resume` | `False` | Resume from checkpoint | +| `whole_slide` | `False` | Dense whole-slide training mode | +| `val_whole_slide` | `None` | Whole-slide mode for validation (defaults to `whole_slide`) | +| `callbacks` | `None` | List of `TrainerCallback` instances | +| `resume` | `False` | Resume from a checkpoint in `output_dir` | #### Callbacks -Subclass `TrainerCallback` to hook into the training loop: +Subclass `TrainerCallback` to hook into the loop: ```python from spatial_transcript_former.training import TrainerCallback @@ -457,32 +397,25 @@ from spatial_transcript_former.training import TrainerCallback class WandbCallback(TrainerCallback): def on_epoch_end(self, trainer, epoch, metrics): wandb.log(metrics, step=epoch) - - def should_stop(self, trainer, epoch, metrics): - return False # never stop early ``` | Hook | When | -| --- | --- | +|---|---| | `on_train_begin(trainer)` | Start of `fit()` | | `on_epoch_begin(trainer, epoch)` | Before each epoch | | `on_epoch_end(trainer, epoch, metrics)` | After validation | | `on_train_end(trainer, results)` | End of `fit()` | -| `should_stop(trainer, epoch, metrics)` | Return `True` to stop | +| `should_stop(trainer, epoch, metrics)` | Return `True` to stop early | -Built-in: `EarlyStoppingCallback(patience=15, min_delta=0.0)` - ---- +Built-in: `EarlyStoppingCallback(patience=15, min_delta=0.0)` (monitors `val_loss`). ### Programmatic Training (Low-Level) -For full control, use the engine functions directly: - ```python from spatial_transcript_former.training import train_one_epoch, validate, CompositeLoss -criterion = CompositeLoss(alpha=1.0) -optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) +criterion = CompositeLoss(alpha=1.0, pcc_type="ccc") +optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) for epoch in range(100): train_loss = train_one_epoch( @@ -496,95 +429,86 @@ for epoch in range(100): print(f"Epoch {epoch}: train={train_loss:.4f}, val={val_metrics['val_loss']:.4f}") ``` ---- +`validate(...)` returns a dict: `val_loss`, `val_mae`, `val_baseline_mae`, `val_pcc`, `val_ccc`, `val_pcc_per_pathway`, `val_ccc_per_pathway`, `pred_variance`, `spatial_coherence`, `attn_correlation` (entries are `None` when not applicable). ### Loss Functions (`training.losses`) -All losses accept `(B, G)` patch-level or `(B, N, G)` dense inputs, with optional `mask` for padded positions. +All accept `(B, P)` patch-level or `(B, N, P)` dense inputs, with optional `mask` (padded positions) and `pathway_mask` (invalid/zero-variance pathways). `CLIPAlignmentLoss` takes `mask` only. | Class | Formula / Description | -| --- | --- | -| `MaskedMSELoss` | Standard MSE with optional padding mask | -| `PCCLoss` | `1 - mean(PCC)` — gene-wise spatial Pearson correlation | -| `CompositeLoss` | `MSE + α · (1 - PCC)` — balances magnitude and spatial pattern | -| `ZINBLoss` | Zero-Inflated Negative Binomial NLL — for raw count data | -| `MaskedHuberLoss` | Huber (SmoothL1) — robust to outlier spots | -| `AuxiliaryPathwayLoss` | Wraps any base loss + PCC on pathway bottleneck scores | +|---|---| +| `MaskedMSELoss` | MSE over valid positions | +| `MaskedHuberLoss` | Huber / Smooth-L1 — robust to outlier spots (`delta=1.0`) | +| `PCCLoss` | `1 − mean(PCC)` — per-pathway spatial Pearson correlation (scale-invariant) | +| `CCCLoss` | `1 − mean(CCC)` — concordance correlation; penalises mean/variance offsets too | +| `CLIPAlignmentLoss` | Batch-discriminative anti-collapse regulariser in pathway space (opt-in) | +| `CompositeLoss` | `MSE/Huber + α·(1 − PCC/CCC) [+ clip_weight·L_CLIP]` | -### Training Engine (`training.engine`) - -| Function | Description | -| --- | --- | -| `train_one_epoch(model, loader, criterion, optimizer, device, ...)` | One epoch of training. Handles gradient accumulation, AMP, and both patch/WSI modes. Returns average loss. | -| `validate(model, loader, criterion, device, ...)` | Validation pass. Returns `dict` with `val_loss`, `val_mae`, `val_pcc`, `pred_variance`, and optional `attn_correlation`. | +`CompositeLoss(alpha=1.0, eps=1e-8, mse_type="mse"|"huber", pcc_type="pcc"|"ccc", clip_weight=0.0, clip_temperature=0.07)`. The CLI `--loss` values map onto these: `mse`, `pcc`, `ccc`, `mse_pcc`, `mse_ccc`, `mse_ccc_clip`, `mse_huber`. ### Experiment Logger (`training.experiment_logger`) -Offline-friendly logger (no W&B dependency). Writes metrics to SQLite and a JSON summary. +Offline-friendly logger (no W&B dependency). Writes per-epoch metrics to SQLite and a JSON summary. ```python +from spatial_transcript_former.training import ExperimentLogger + logger = ExperimentLogger(output_dir, config_dict) -logger.log_epoch(epoch, {"train_loss": 0.1, "val_loss": 0.2, "val_pcc": 0.65}) +logger.log_epoch(epoch, {"train_loss": 0.1, "val_loss": 0.2, "val_ccc": 0.65}) logger.finalize(best_val_loss=0.15) ``` | Output File | Contents | -| --- | --- | -| `training_logs.sqlite` | Per-epoch metrics table | -| `results_summary.json` | Config + final metrics + runtime | +|---|---| +| `training_logs.sqlite` | Per-epoch metrics (table `metrics`, columns added dynamically) | +| `results_summary.json` | `config` + final metrics + runtime | ### Checkpoint Lifecycle -During training, checkpoints are managed by `training.checkpoint` (the *internal* module — distinct from the public `save_pretrained`): +During training, full training state (for `--resume`) is managed by `training.checkpoint` — the *internal* module, distinct from the public `save_pretrained`: | Function | Purpose | -| --- | --- | -| `save_checkpoint(model, optimizer, scaler, schedulers, ...)` | Saves full training state for `--resume` | -| `load_checkpoint(model, optimizer, scaler, schedulers, ...)` | Restores training state | - -After training is complete, use the public `save_pretrained` to export a clean, inference-ready checkpoint: +|---|---| +| `save_checkpoint(model, optimizer, scaler, schedulers, epoch, best_val_metric, output_dir, model_name)` | Writes `latest_model_.pth` (full training state) | +| `load_checkpoint(model, optimizer, scaler, schedulers, output_dir, model_name, device)` | Restores it → `(start_epoch, best_val_metric, loaded_schedulers)` | -```python -from spatial_transcript_former import save_pretrained - -# Export for inference (strips optimizer/scheduler state) -save_pretrained(model, "./release/v1/", gene_names=gene_list) -``` +After training, use the public `save_pretrained` to export a clean, inference-ready checkpoint directory. --- ## Bring Your Own Data -All datasets implement the `SpatialDataset` contract (in `data.base`). The contract requires `__getitem__` to return: +All datasets implement the `SpatialDataset` contract (in `data.base`). `__getitem__` must return a **5-tuple**: ```python -(features, gene_counts, rel_coords) -# features: (S, D) tensor — patch embeddings (S = 1 + neighbours) -# gene_counts: (G,) tensor — expression targets -# rel_coords: (S, 2) tensor — relative spatial coordinates +(features, gene_counts, pathway_targets, rel_coords, mask) +# features: (S, D) tensor — patch embeddings (S = 1 + neighbours), or an image tensor +# gene_counts: legacy slot — pass None (the model is pathway-exclusive) +# pathway_targets: (P,) tensor — the supervised target (or None for inference-only) +# rel_coords: (S, 2) tensor — coordinates relative to the centre patch +# mask: (S,) bool tensor — True = padded, or None ``` ### Minimal Implementation ```python -from spatial_transcript_former.data.base import SpatialDataset import torch +from spatial_transcript_former.data.base import SpatialDataset class MyVisiumDataset(SpatialDataset): - def __init__(self, features, gene_matrix, coords): + def __init__(self, features, pathways, coords): self._features = torch.as_tensor(features, dtype=torch.float32) - self._genes = torch.as_tensor(gene_matrix, dtype=torch.float32) + self._pathways = torch.as_tensor(pathways, dtype=torch.float32) self._coords = torch.as_tensor(coords, dtype=torch.float32) - self.num_genes = self._genes.shape[1] + self.num_pathways = self._pathways.shape[1] def __len__(self): return len(self._features) def __getitem__(self, idx): - feat = self._features[idx].unsqueeze(0) # (1, D) - genes = self._genes[idx] # (G,) - rel_coord = torch.zeros(1, 2) # centre = [0,0] - return feat, genes, rel_coord + feat = self._features[idx].unsqueeze(0) # (1, D) + rel = torch.zeros(1, 2) # centre = [0, 0] + return feat, None, self._pathways[idx], rel, None ``` ### Training Your Custom Dataset @@ -593,16 +517,17 @@ class MyVisiumDataset(SpatialDataset): from torch.utils.data import DataLoader, random_split from spatial_transcript_former import SpatialTranscriptFormer, Trainer from spatial_transcript_former.training import CompositeLoss, EarlyStoppingCallback +from spatial_transcript_former.recipes.hest.dataset import collate_fn_patch -dataset = MyVisiumDataset(features, gene_matrix, coords) +dataset = MyVisiumDataset(features, pathways, coords) train_ds, val_ds = random_split(dataset, [0.8, 0.2]) -model = SpatialTranscriptFormer(num_genes=dataset.num_genes, backbone_name="phikon") +model = SpatialTranscriptFormer(num_pathways=dataset.num_pathways, backbone_name="phikon") trainer = Trainer( model=model, - train_loader=DataLoader(train_ds, batch_size=32, shuffle=True), - val_loader=DataLoader(val_ds, batch_size=64), + train_loader=DataLoader(train_ds, batch_size=32, shuffle=True, collate_fn=collate_fn_patch), + val_loader=DataLoader(val_ds, batch_size=64, collate_fn=collate_fn_patch), criterion=CompositeLoss(), epochs=100, callbacks=[EarlyStoppingCallback(patience=15)], @@ -611,4 +536,4 @@ results = trainer.fit() trainer.save_pretrained("./my_model/") ``` -See `recipes/custom/README.md` for the full guide. +See [`recipes/custom/README.md`](../src/spatial_transcript_former/recipes/custom/README.md) for the full guide. diff --git a/docs/DATALOADER.md b/docs/DATALOADER.md index fee010b..4c5e3c6 100644 --- a/docs/DATALOADER.md +++ b/docs/DATALOADER.md @@ -15,7 +15,7 @@ This path is used when training or evaluating directly on pixel-space images. * **Neighbourhood Context**: Can retrieve a patch along with its $K$ nearest neighbours. * **Dihedral Augmentation**: Randomly rotates or flips patch pixels and coordinates in sync. * **`get_hest_dataloader`**: High-level orchestrator that creates a `DataLoader` over raw patches for a list of sample IDs, combining individual datasets using `ConcatDataset`. -* **Returned Tuples**: Yields `(patches, None, rel_coords)` where the second element (formerly gene expression counts) is `None`. +* **Returned Tuples**: Yields `(patches, None, pathway_acts, rel_coords, mask)`. The second slot is the legacy gene-counts slot (always `None` now); `pathway_acts` holds the pre-computed targets (or `None` when no `pathway_targets_dir` is given). ### 2. Pre-Computed Feature Loading Path @@ -27,7 +27,7 @@ This is the default path used by the SpatialTranscriptFormer training pipeline ( * **Patch Mode**: Returns a single spot feature vector, its local neighbourhood features (optionally with random dropout augmentation), pre-computed pathway targets, and relative coordinates. * **Whole-Slide Mode**: Returns all spots on the slide as a single sequence. * **`get_hest_feature_dataloader`**: Builds a `DataLoader` over the feature datasets. - * In **patch mode**, yields standard batched tensors `(feats, None, pathway_acts, coords)`. + * In **patch mode**, yields standard batched tensors `(feats, None, pathway_acts, coords, mask)`. * In **whole-slide mode**, pads variable-length slides to the longest slide in the batch and appends a boolean padding mask. Yields `(padded_feats, None, padded_pathways, padded_coords, mask)`. --- @@ -50,10 +50,11 @@ dataloader = get_hest_feature_dataloader( pathway_targets_dir="./hest_data/pathway_activities" ) -for feats, _, pathway_acts, rel_coords in dataloader: - # feats shape: (BS, 1 + n_neighbors, feature_dim) +for feats, _, pathway_acts, rel_coords, mask in dataloader: + # feats shape: (BS, 1 + n_neighbors, feature_dim) # pathway_acts shape: (BS, num_pathways) - # rel_coords shape: (BS, 1 + n_neighbors, 2) + # rel_coords shape: (BS, 1 + n_neighbors, 2) + # mask shape: (BS, 1 + n_neighbors) bool, True = padded ... ``` diff --git a/docs/IP_STATEMENT.md b/docs/IP_STATEMENT.md index 61c03e3..b11f246 100644 --- a/docs/IP_STATEMENT.md +++ b/docs/IP_STATEMENT.md @@ -1,45 +1,60 @@ -# Intellectual Property Statement - -This document identifies the core technological and methodological contributions of Benjamin Isaac Wilson within the **SpatialTranscriptFormer** project. The Author retains all intellectual property rights to these components. - -## Scientific Context & Attribution - -While this implementation introduces novel spatial constraints and transcriptomics-specific bottlenecks, the core multimodal interaction logic is inspired by and derived from the **SURVPATH** architecture ([Jaume et al., 2024](https://arxiv.org/abs/2303.15545)). This work focuses on the **adaptation and extension** of those concepts to the domain of high-dimensional spatial gene expression regression. - -## 1. Model Architecture (SpatialTranscriptFormer) - -The primary innovation is the **multimodal bottleneck transformer** designed for spatial transcriptomics regression. Key IP components include: - -- **Pathway-Histology Interaction Layer**: The specific implementation of Cross-Attention where learnable pathway tokens (queries) interact with high-resolution histology tokens (keys/values). -- **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. -- **Coordinate-Aware Augmentation**: Specific strategies for utilizing relative patch coordinates as inputs/constraints in the transformer pipeline. - -## 3. Data Integration & Alignment - -- **Gene-Lock Alignment Logic**: The specific implementation of automated gene set alignment across heterogeneous spatial transcriptomics technologies (Visium, Xenium, etc.). -- **Multi-Resolution Patch Loading**: Specialized dataloader logic for handling whole-slide histology integration with localized gene expression matrices. - -## 4. Visualization & Interpretability - -- **Histology-Gene Overlay Heatmaps**: The specific implementation of scaling and aligning predicted gene expression distributions onto full-resolution WSIs for clinical interpretation. - ---- - -For technical inquiries or discussion regarding these methodologies, please contact the Author. +# Attribution, Provenance & Design Notes + +> **Licensing note:** SpatialTranscriptFormer is released under the Apache +> License 2.0 (see [LICENSE](../LICENSE) and [NOTICE](../NOTICE)). This document +> is *descriptive* — it records design provenance and attributions for readers +> and reviewers. It does not itself grant or reserve any rights. + +## Provenance + +This implementation is original work. The multimodal interaction framing — +treating learnable biological-pathway tokens as a bottleneck that cross-attends +to histology patch tokens — was **conceptually inspired by** SURVPATH +([Jaume et al., 2024](https://arxiv.org/abs/2303.15545)). **No SURVPATH source +code is used, copied, or adapted here**; the model, losses, data pipeline, and +training code were written from scratch for the spatial-transcriptomics +pathway-activity task. + +## Design contributions + +The notable design choices of this project (described for readers/reviewers, not +as legal claims): + +- **Pathway-exclusive prediction.** The model predicts pre-computed spatial + pathway-activity targets directly, with no intermediate gene-reconstruction + step. Targets are computed offline (QC → CP10k normalisation → mean pathway + aggregation; see [PATHWAY_MAPPING.md](PATHWAY_MAPPING.md)), decoupling + biological-knowledge integration from training and removing the circular + auxiliary-loss dependency used in earlier prototypes. + +- **Dot-product + Softplus scoring head.** Per-spot pathway scores are + `softplus(scale · (patch · pathwayᵀ / √d) + bias)`, with learnable per-pathway + scale and bias. Softplus yields non-negative scores matching the + mean-log1p-CP10k targets; the √d scaling keeps affinities well-conditioned (as + in scaled dot-product attention). This is *not* cosine similarity. + +- **Configurable quad-flow interaction masking.** Attention between pathway (P) + and histology (H) tokens is selectable per quadrant (`p2p`, `p2h`, `h2p`, + `h2h`), enabling ablations and the "pathway bottleneck" variant (block `h2h` + so all inter-patch information must flow through named pathways). + +- **Slide-stationary spatial positional encoding.** Patch tokens receive a + learned 2-D encoding of centred/standardised coordinates, so the model learns + spatially-varying pathway maps rather than slide-level averages. + +- **Moran's I as diagnostic + collapse detector.** Spatial autocorrelation is + used to curate targets and to detect representation collapse during validation + (correlation of predicted vs. ground-truth per-pathway Moran's I), rather than + as a loss term. See [spatial_stats.py](../src/spatial_transcript_former/data/spatial_stats.py). + +## Third-party attributions + +See [NOTICE](../NOTICE) for the licenses and attribution requirements of the +datasets, gene sets, and foundation-model backbones this framework uses at +runtime (HEST-1k, MSigDB Hallmarks, CTransPath / Phikon / GigaPath / Hibou, …). + +## Citing + +If you use this work in academic research, please cite this repository and its +author, Benjamin Isaac Wilson, alongside the relevant third-party works listed +in [NOTICE](../NOTICE). diff --git a/docs/MODELS.md b/docs/MODELS.md index f9baa40..43afe3f 100644 --- a/docs/MODELS.md +++ b/docs/MODELS.md @@ -89,9 +89,9 @@ Together, these ensure the model learns *spatially-varying* pathway activation m │ └──────────┬───────────────────┘ │ │ │ │ │ ┌──────────▼──────────────────┐ │ -│ │ Cosine Similarity Scoring │ │ -│ │ │ │ -│ │ norm_patch @ norm_pathway.T │ │ +│ │ Dot-Product + Softplus Head │ │ +│ │ raw=(patch@path.T)/sqrt(d) │ │ +│ │ -> softplus(scale*raw+bias) │ │ │ │ → Pathway Scores (S, P) │ │ │ └──────────┬───────────────────┘ │ │ │ │ @@ -141,17 +141,15 @@ The transformer uses a custom attention mask that controls information flow betw > [!NOTE] > Disabling `h2h` creates the **Pathway Bottleneck** variant, where all inter-patch communication must flow through the pathway tokens. This requires **minimum 2 transformer layers**: Layer 1 lets pathways gather information from patches, and Layer 2 lets patches read the contextualised pathway tokens. -#### Cosine Similarity Scoring +#### Dot-Product + Softplus Scoring Head -Pathway scores are computed via L2-normalized cosine similarity with a learnable temperature parameter $\tau$ (following CLIP): -$$s_{ij} = \cos(\hat{h}_i, \hat{p}_j) \times \tau$$ -where $\hat{h}_i$ and $\hat{p}_j$ are the L2-normalized processed patch and pathway tokens respectively. This produces scores in $[-\tau, +\tau]$ with meaningful relative differences, avoiding the saturation that occurs with raw dot-products. +Pathway scores are computed from a **scaled dot-product** between the processed patch and pathway tokens, shifted by a learnable per-pathway scale and bias, then passed through a Softplus non-linearity: +$$\hat{s}_{i,k} = \text{softplus}\!\left(\gamma_k \cdot \frac{h_i \cdot p_k}{\sqrt{d}} + \beta_k\right)$$ +where $h_i$ and $p_k$ are the processed patch and pathway tokens, $d$ is the token dimension, and $\gamma_k, \beta_k$ are learnable per-pathway scale and bias parameters. The $\sqrt{d}$ scaling keeps the affinities well-conditioned (as in scaled dot-product attention), and Softplus maps them to **non-negative** activity scores in $[0, \infty)$ that match the targets (mean log1p CP10k pathway aggregates). Note this is **not** cosine similarity — the tokens are deliberately *not* L2-normalised, so the head can represent the unbounded, non-negative magnitudes the targets require. #### Direct Pathway Prediction -The model's output *is* the pathway activity score vector. There is no intermediate gene reconstruction layer. The cosine similarity scores between (processed) patch tokens and pathway tokens directly serve as the prediction: -$$\hat{s}_{i,k} = \cos(\hat{h}_i, \hat{p}_k)$$ -where $\hat{h}_i$ and $\hat{p}_k$ are the L2-normalised patch and pathway tokens, respectively. These are supervised against pre-computed pathway activities (see [PATHWAY_MAPPING.md](PATHWAY_MAPPING.md)). +The model's output *is* the pathway activity score vector. There is no intermediate gene reconstruction layer. The scaled dot-product + Softplus scores $\hat{s}_{i,k}$ defined above directly serve as the prediction, supervised against the pre-computed pathway activities (see [PATHWAY_MAPPING.md](PATHWAY_MAPPING.md)). ### 2.4 Training Modes diff --git a/docs/PATHWAY_MAPPING.md b/docs/PATHWAY_MAPPING.md index 11f64e8..6600b28 100644 --- a/docs/PATHWAY_MAPPING.md +++ b/docs/PATHWAY_MAPPING.md @@ -140,9 +140,7 @@ This makes the pipeline applicable to any disease of interest without changing t After training, the `stf-predict` CLI and the `Predictor` Python API produce spatial heatmaps of pathway activation across the tissue. Each pathway gets its own overlay plot, allowing direct visual comparison between, e.g., hypoxia activation and VEGF signalling. ```bash -stf-predict --data-dir hest_data --sample-id MEND29 \ - --model-path checkpoints/best_model.pth \ - --model-type interaction +stf-predict --run-dir checkpoints --sample-id MEND29 --output-dir results ``` Plots are saved to `./results/` by default. diff --git a/docs/SC_BEST_PRACTICES.md b/docs/SC_BEST_PRACTICES.md index be1dc1f..2e6b5da 100644 --- a/docs/SC_BEST_PRACTICES.md +++ b/docs/SC_BEST_PRACTICES.md @@ -52,8 +52,8 @@ evaluation. | **GHIST** | Nat. Methods 2025 | Single-cell resolution via Xenium | STF's differentiators in this landscape are **pathway-level interpretability** and -**auxiliary biological supervision** — both under-represented in the 2025 benchmark's -*translational potential* category, where competing methods score poorly. +**direct biological (pathway) supervision** — both under-represented in the 2025 +benchmark's *translational potential* category, where competing methods score poorly. ### Prior Knowledge: Saezlab Ecosystem (EMBL Heidelberg) @@ -69,8 +69,8 @@ pathways (EGFR, MAPK, PI3K, TGFb, TNFa, Hypoxia, VEGF, WNT, p53, NFkB, JAK-STAT, Androgen, Estrogen, Trail) derived empirically from a large compendium of perturbation experiments. Each gene carries a **continuous signed weight** (direction and magnitude of pathway response). This is fundamentally different from MSigDB Hallmarks, which -use binary unweighted membership. PROGENy weights can be used directly as the -initialisation matrix for the linear gene reconstructor (`--pathway-init`). Holland +use binary unweighted membership. PROGENy weights are a natural warm-start for the +learnable pathway tokens (see item 9). Holland et al. (Genome Biology, 2020) benchmarked PROGENy on scRNA-seq and Visium data, confirming it performs well in the spot-level transcriptome regime. @@ -79,9 +79,8 @@ package that applies a family of statistical estimators (MLM, ULM, GSEA, AUCell, consensus) to any (resource, method) combination. `get_progeny()` and `run_mlm()` provide PROGENy activity scores out of the box. The official documentation includes a Visium spatial transcriptomics notebook. This is the practical bridge for computing -weighted pathway activity scores from HEST-1k expression matrices — directly -applicable as `AuxiliaryPathwayLoss` targets instead of the current binary MSigDB -membership scores. +weighted pathway activity scores from HEST-1k expression matrices — a richer +alternative to the current Hallmark mean-aggregation targets (see item 7). **CollecTRI** (Müller-Dott et al., Nucleic Acids Research, 2023) — The recommended successor to DoRothEA for transcription factor regulons. Integrates 12 resources for @@ -168,28 +167,34 @@ Explicit QC thresholds (minimum UMI count, minimum detected genes, maximum mitoc ### Training & Supervision -**Architectural direction: pathway activity as the primary task.** The current model is -trained to predict 1000 individual gene expression values, with pathway scores as an -auxiliary, interpretability-oriented output. This framing has two compounding problems: +**Architectural direction: pathway activity as the primary task.** ✅ **Implemented.** +This was the central reframing of the project and is now the shipped architecture: the +model predicts pathway activity directly via a dot-product + Softplus head, with no +gene-reconstruction step. The rationale is kept here as a record of *why* the change +was made. + +Earlier prototypes predicted ~1000 individual gene expression values with pathway +scores as an auxiliary output. That framing had two compounding problems: 1. **Noisy learning signal.** Even after log-normalisation and SVG-based vocabulary selection, MSE over 1000 genes is dominated by high-expression genes — not necessarily the spatially informative ones. Moran's I weighting fixes gene *selection* but not the *gradient* during training. -2. **The `AuxiliaryPathwayLoss` is circular.** Its ground-truth targets are computed - from the same gene expression being predicted (binary MSigDB membership sums). It - constrains the intermediate representation to match a re-parameterisation of the +2. **The (former) `AuxiliaryPathwayLoss` was circular.** Its targets were computed + from the same gene expression being predicted (MSigDB membership sums), so it + constrained the intermediate representation to match a re-parameterisation of the output, adding no independent biological information. -**The cleaner framing** is to invert the task hierarchy: +**The fix** was to invert the task hierarchy: ``` PRIMARY: H&E → pathway activity maps (spatially coherent, interpretable) -SECONDARY: pathway scores → gene expression (optional imputation head) +SECONDARY: pathway scores → gene expression (optional; not currently implemented) ``` -Pathway activity maps pre-computed offline via decoupleR + PROGENy are: +Pathway activity targets (currently the mean of log1p CP10k expression over each MSigDB +Hallmark's member genes, pre-computed offline) are: - **Spatially cleaner** than individual genes — Moran's I on pathway aggregates is typically 3–5× higher than on individual gene expression - **Adaptable** — the offline preprocessing step can be swapped for any biologically @@ -198,18 +203,19 @@ Pathway activity maps pre-computed offline via decoupleR + PROGENy are: - **Genuinely supervised** — the pathway scores are first-class outputs, not constraints on an intermediate representation -The `AuxiliaryPathwayLoss` in its current form is superseded by this framing. The -gene reconstructor becomes an optional secondary head rather than the primary loss -target. +The remaining open leads below concern *richer target backends* (decoupleR/PROGENy) +and an *optional* secondary gene head — not the core reframing, which is done. --- -**7. Pre-compute pathway activity targets (decoupleR + PROGENy)** 💡 — High priority, medium effort +**7. Richer pathway-activity targets (decoupleR + PROGENy)** 💡 — High priority, medium effort -For each HEST sample, run `decoupler.run_mlm(expression, net=get_progeny())` to -produce a `(spots × 14)` signed pathway activity matrix. Cache alongside the existing -feature `.h5` files. These become the primary prediction targets replacing the -`AuxiliaryPathwayLoss`. +The offline target pipeline (`stf-compute-pathways`) currently aggregates MSigDB +Hallmark membership (mean log1p CP10k). A higher-quality alternative: for each HEST +sample, run `decoupler.run_mlm(expression, net=get_progeny())` to produce a +`(spots × 14)` *signed, continuous* pathway activity matrix and cache it as an +alternative target set. Because targets are decoupled from the model, this is a +preprocessing-only change — no model edits required. The decoupleR MLM estimator fits: `expression[:, g] = a_p · W[g, p] + b_p + ε` and solves for `a_p` — the pathway activity at each spot — via OLS. The result is @@ -217,32 +223,30 @@ continuous, signed, and validated against known pathway perturbation experiments prior knowledge source (PROGENy, Hallmarks, LIANA+ LR pairs, CollecTRI TF regulons) can be swapped at the preprocessing step without touching the model. -**8. Moran's I weighted gene loss** 🔧 — High priority, low effort +**8. Moran's I weighted gene loss** 🚫 — Superseded (pathway-exclusive) -If gene prediction is retained as a secondary head, weight each gene's contribution -to the loss by its Moran's I score so that spatially coherent genes drive the -gradient: +Only relevant if an optional secondary gene-reconstruction head is ever added. In that +case each gene's loss contribution could be weighted by its Moran's I so spatially +coherent genes drive the gradient: ```python -L_gene = sum(w_g * MSE(pred_g, target_g) for g in genes) / sum(w_g) -# where w_g = MoranI_g (pre-computed during stf-build-vocab --svg-weight > 0) +L_gene = sum(w_g * MSE(pred_g, target_g) for g in genes) / sum(w_g) # w_g = MoranI_g ``` -This is the training-time counterpart to SVG-based vocabulary selection — both -ensure spatially informative genes dominate, one at selection time and one at -training time. +For the current pathway-exclusive model this does not apply — there is no gene head. -**9. PROGENy pathway token initialisation** 💡 — High priority, medium effort +**9. PROGENy pathway-token initialisation** 💡 — Medium priority, medium effort -Replace or supplement the binary MSigDB Hallmark membership matrix used in -`--pathway-init` with PROGENy's continuous signed footprint weights. Each gene -carries a directional weight reflecting its perturbation-validated response -magnitude, rather than binary membership. +The learnable pathway tokens are currently randomly initialised. A lead: warm-start +them from PROGENy's continuous signed footprint weights (each gene carries a +directional, perturbation-validated weight) instead of from scratch. -Hybrid strategy: initialise the first 14 pathway tokens from PROGENy weights -(signalling), the remaining 36 from Hallmarks membership (metabolic, immune, -proliferation). Requires additions to `data/pathways.py` and a new -`--pathway-init progeny` option in `training/builder.py`. +Hybrid strategy: initialise the first 14 tokens from PROGENy weights (signalling) and +the rest from Hallmark membership (metabolic, immune, proliferation). `get_pathway_init` +in `data/pathways.py` already builds a membership matrix; this would re-introduce a +token-initialisation pathway into the model and a corresponding option in +`training/builder.py` (the earlier `--pathway-init` flag was removed in the +pathway-exclusive cleanup and would need re-adding). **10. Cell–cell interaction pathway tokens (LIANA+)** 💡 — Low priority, high effort @@ -277,13 +281,15 @@ The architecture is platform-agnostic. Scaling to sub-cellular resolution platfo preprocessing pipelines, not the model itself. GHIST demonstrates that Xenium-derived training labels substantially improve resolution. -**14. Preprocessing data contract** 🔧 — Low priority, low effort +**14. Preprocessing data contract** ✅ — Largely implemented -Explicitly documenting which normalisation is applied at which stage (per-spot QC → -library-size normalisation → log1p → gene selection) and how the vocabulary was built -(which `--svg-weight`, which `--pathways`, run date) prevents silent mismatches between -training and inference. A short header in each output folder (`vocab_config.json`) -would suffice. +Each `stf-compute-pathways` output `.h5` records its provenance as attributes +(`format_version`, QC thresholds, spot counts, scored-pathway count; see +[PATHWAY_MAPPING.md](PATHWAY_MAPPING.md)), and `save_pretrained` stamps +`pathway_format_version` into `config.json` so `load_pretrained` / `--resume` refuse +checkpoints trained against incompatible targets. Remaining nicety: surface the same +provenance (QC params, `--pathways` subset, run date) in a per-run summary for quick +human inspection. --- diff --git a/docs/TRAINING_GUIDE.md b/docs/TRAINING_GUIDE.md index 2f47738..1ad4202 100644 --- a/docs/TRAINING_GUIDE.md +++ b/docs/TRAINING_GUIDE.md @@ -57,7 +57,7 @@ python -m spatial_transcript_former.train \ ## 2. Whole-Slide MIL (Multiple Instance Learning) -Aggregates all patches from a slide to predict the average pathway activities. Recommended to use **precomputed features** from the `stf-compute-features` CLI tool for speed. Foundation models like ctranspath can be used as backbones. +Aggregates all patches from a slide to predict the average pathway activities. Recommended to use **precomputed features** from the `stf-extract` CLI tool for speed. Foundation models like ctranspath can be used as backbones. ### Attention MIL (Weak Supervision) @@ -72,7 +72,6 @@ python -m spatial_transcript_former.train \ --precomputed \ --weak-supervision \ --use-amp \ - --log-transform \ --epochs 50 ``` @@ -89,7 +88,6 @@ python -m spatial_transcript_former.train \ --precomputed \ --weak-supervision \ --use-amp \ - --log-transform \ --epochs 50 ``` @@ -108,7 +106,6 @@ python -m spatial_transcript_former.train \ --data-dir A:\hest_data \ --model interaction \ --backbone ctranspath \ - --use-nystrom \ --num-pathways 50 \ --precomputed \ --whole-slide \ @@ -116,7 +113,7 @@ python -m spatial_transcript_former.train \ --epochs 100 ``` -### Using PROGENy-Style Priors (14 Pathways) - Currently Not Implemented +### Using PROGENy-Style Priors (14 Pathways) Use `--pathway-prior progeny` to set the number of pathway tokens to 14, matching the dimensionality of PROGENy-style signalling pathway databases. This is useful when your pre-computed targets were derived from a 14-pathway database rather than the 50 MSigDB Hallmarks. @@ -132,7 +129,7 @@ python -m spatial_transcript_former.train \ --epochs 100 ``` -> **Note**: `--pathway-prior progeny` sets `--num-pathways` to 14 automatically. Ensure your pre-computed pathway targets match this dimensionality. +> **Note**: `--pathway-prior progeny` only sets `--num-pathways` to 14 (the output dimensionality). Generating PROGENy *targets* is not yet built into `stf-compute-pathways` (which defaults to MSigDB Hallmarks) — supply matching 14-pathway targets via `--custom-gmt`. ### Recommended: Using Presets @@ -250,13 +247,12 @@ python -m spatial_transcript_former.train --resume --output-dir runs/my_experime | `--weak-supervision` | Bag-level training for MIL models. | Use with `attention_mil` or `transmil`. | | `--pathway-targets-dir` | Directory of pre-computed `.h5` pathway activity files. | Defaults to `/pathway_activities`. | | `--feature-dir` | Explicit path to precomputed features directory. | Overrides auto-detection. | -| `--loss` | Loss function: `mse`, `pcc`, `mse_pcc`. | `mse_pcc` recommended. | +| `--loss` | Loss function: `mse`, `pcc`, `ccc`, `mse_pcc`, `mse_ccc`, `mse_ccc_clip`, `mse_huber`. | `mse_pcc` / `mse_ccc` recommended. | | `--pcc-weight` | Weight ($\alpha$) for PCC term in composite loss. | Default 1.0. | | `--pathway-prior` | Pathway prior for token count (`hallmarks`=50). | Use `hallmarks` for MSigDB Hallmarks. | | `--interactions` | Enabled attention quadrants: `p2p`, `p2h`, `h2p`, `h2h`. | Default: `all` (Full Interaction). | | `--plot-pathways-list` | Names of pathways to visualize as heatmaps during validation. | e.g. `HYPOXIA ANGIOGENESIS` | | `--num-pathways` | Number of pathway bottleneck tokens (overridden by `--pathway-prior`). | Match your pre-computed targets. | -| `--mask-radius` | Euclidean distance for spatial attention gating. | Usually between 200 and 800. | | `--n-neighbors` | Number of context neighbors to load. | Set `> 0` for hybrid/GNN models. | | `--use-amp` | Mixed precision training. | Recommended on modern GPUs. | | `--grad-accum-steps` | Gradient accumulation steps. | Use when memory is limited. | @@ -268,7 +264,7 @@ python -m spatial_transcript_former.train --resume --output-dir runs/my_experime ## Tips for Success 1. **Pathway Pre-Computation**: Always run `stf-compute-pathways --data-dir hest_data` before training. The trainer will error if pathway targets are missing. -2. **Feature Extraction**: Run `hpc/prepare_data.slurm` or `scripts/extract_features.py` before training with `--precomputed`. +2. **Feature Extraction**: Run `stf-extract --data-dir hest_data --backbone ` before training with `--precomputed` (see [PRECOMPUTED_WORKFLOW.md](PRECOMPUTED_WORKFLOW.md)). 3. **Output**: Checkpoints, logs, and JSON summaries are saved to `--output-dir` (default: `./checkpoints`). 4. **Debug Mode**: Use `--max-samples 3 --epochs 1` to verify your setup before a full run. -5. **Results Aggregation**: Use `hpc/collect_results.py` to compare experiments across multiple runs. +5. **Results Aggregation**: Use `scripts/compare_models.py --collect-only` to compare experiments across multiple runs (see §4). diff --git a/pyproject.toml b/pyproject.toml index 19c1a38..d32ea96 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,6 +8,11 @@ version = "0.1.0" description = "A transformer-based model for spatial transcriptomics" readme = "README.md" requires-python = ">=3.9" +classifiers = [ + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + "Operating System :: OS Independent", +] dependencies = [ "pandas", "h5py", diff --git a/src/spatial_transcript_former/__init__.py b/src/spatial_transcript_former/__init__.py index ada9443..e7d85db 100644 --- a/src/spatial_transcript_former/__init__.py +++ b/src/spatial_transcript_former/__init__.py @@ -1,5 +1,5 @@ """ -SpatialTranscriptFormer — predict gene expression from histology. +SpatialTranscriptFormer — predict biological pathway activity from histology. Core public API:: diff --git a/src/spatial_transcript_former/data/base.py b/src/spatial_transcript_former/data/base.py index e777976..46a2f6e 100644 --- a/src/spatial_transcript_former/data/base.py +++ b/src/spatial_transcript_former/data/base.py @@ -20,50 +20,63 @@ class SpatialDataset(Dataset, ABC): Any dataset used with SpatialTranscriptFormer must subclass this and implement :meth:`__getitem__` and :meth:`__len__`. - ``__getitem__`` must return a 3-tuple:: + ``__getitem__`` must return a 5-tuple:: - (features, gene_counts, rel_coords) + (features, gene_counts, pathway_targets, rel_coords, mask) where: * **features** — ``(S, D)`` float tensor of patch embeddings (``S`` = 1 + K neighbours, ``D`` = backbone feature dim), or ``(3, H, W)`` / ``(S, 3, H, W)`` image tensor in raw-patch mode. - * **gene_counts** — ``(G,)`` float tensor of gene expression targets. + * **gene_counts** — legacy slot kept for tuple compatibility. Set to + ``None`` in the current pathway-exclusive pipeline (the model no longer + reconstructs gene expression). + * **pathway_targets** — ``(P,)`` float tensor of pre-computed pathway + activity targets — the supervised signal (see + :mod:`spatial_transcript_former.recipes.hest.compute_pathway_activities`). + May be ``None`` for inference-only datasets with no targets. * **rel_coords** — ``(S, 2)`` float tensor of spatial coordinates relative to the centre patch (centre is always ``[0, 0]``). + * **mask** — ``(S,)`` boolean padding mask (``True`` = padded/ignore), or + ``None`` when no padding is required. - Subclasses SHOULD also expose :attr:`num_genes` and (optionally) - :attr:`gene_names` as properties. + The matching collate function + (:func:`spatial_transcript_former.recipes.hest.dataset.collate_fn_patch`) + stacks each slot into a batch, passing ``None`` through for absent slots — + exactly what :func:`spatial_transcript_former.training.engine.train_one_epoch` + and :class:`~spatial_transcript_former.training.trainer.Trainer` consume. + + Subclasses SHOULD also expose :attr:`num_pathways` and (optionally) + :attr:`pathway_names`. Example:: class MyVisiumDataset(SpatialDataset): - def __init__(self, slide_path, genes, coords, features): + def __init__(self, features, pathways, coords): self._features = features # (N, D) - self._genes = genes # (N, G) + self._pathways = pathways # (N, P) self._coords = coords # (N, 2) + self.num_pathways = pathways.shape[1] def __len__(self): return len(self._features) def __getitem__(self, idx): - center = self._coords[idx] - rel = self._coords - center # simplest: no neighbour selection - return self._features[idx:idx+1], self._genes[idx], rel[idx:idx+1] - - @property - def num_genes(self): - return self._genes.shape[1] - - @property - def gene_names(self): - return None # or a list of gene symbols + feat = self._features[idx].unsqueeze(0) # (1, D) + rel = torch.zeros(1, 2) # centre = [0, 0] + return feat, None, self._pathways[idx], rel, None """ @abstractmethod - def __getitem__(self, idx) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Return ``(features, gene_counts, rel_coords)`` for index ``idx``.""" + def __getitem__(self, idx) -> Tuple[ + torch.Tensor, + Optional[torch.Tensor], + Optional[torch.Tensor], + torch.Tensor, + Optional[torch.Tensor], + ]: + """Return ``(features, gene_counts, pathway_targets, rel_coords, mask)``.""" ... @abstractmethod @@ -73,12 +86,12 @@ def __len__(self) -> int: # ------------------------------------------------------------------ # Optional attributes — subclasses can set these as instance - # attributes (self.num_genes = ...) or override as properties. + # attributes (self.num_pathways = ...) or override as properties. # ------------------------------------------------------------------ - #: Number of gene expression targets. Set in __init__ or override. - num_genes: int = 0 - #: Ordered list of gene symbols, or None if unavailable. - gene_names: Optional[List[str]] = None + #: Number of pathway activity targets. Set in __init__ or override. + num_pathways: int = 0 + #: Ordered list of pathway names, or None if unavailable. + pathway_names: Optional[List[str]] = None # --------------------------------------------------------------------------- diff --git a/src/spatial_transcript_former/models/interaction.py b/src/spatial_transcript_former/models/interaction.py index d4323cb..7a81c0f 100644 --- a/src/spatial_transcript_former/models/interaction.py +++ b/src/spatial_transcript_former/models/interaction.py @@ -46,11 +46,12 @@ def forward(self, rel_coords): class SpatialTranscriptFormer(nn.Module): - """Transformer for predicting gene expression from histology and spatial context. + """Transformer for predicting pathway activity from histology and spatial context. Integrates histology feature extraction with pathway-based bottleneck - attention to predict gene transcript counts. Follows a standard Vision - Transformer architecture where pathway tokens act as [CLS]-like bottlenecks. + attention to predict per-spot pathway activity scores. Follows a standard + Vision Transformer architecture where pathway tokens act as [CLS]-like + bottlenecks. Attributes: num_pathways (int): Number of pathway bottlenecks. @@ -120,7 +121,9 @@ def __init__( # 3. Learnable pathway tokens (one per pathway, shared across batch) self.pathway_tokens = nn.Parameter(torch.randn(1, num_pathways, token_dim)) - # Learnable scale and bias for dot-product Softplus head (Option 1) + # Learnable per-pathway scale and bias for the dot-product + Softplus + # scoring head (see forward()): they shift the raw patch·pathway + # affinity before Softplus maps it to a non-negative activity score. self.scale = nn.Parameter(torch.ones(1, 1, num_pathways)) self.bias = nn.Parameter(torch.zeros(1, 1, num_pathways)) @@ -183,7 +186,7 @@ def from_pretrained(cls, checkpoint_dir, device="cpu", **kwargs): The directory should contain ``config.json`` and ``model.pth`` (created by :func:`~spatial_transcript_former.checkpoint.save_pretrained`). - An optional ``gene_names.json`` will be loaded as ``model.gene_names``. + An optional ``pathway_names.json`` will be loaded as ``model.pathway_names``. Args: checkpoint_dir (str): Path to checkpoint directory. @@ -302,9 +305,11 @@ def forward( # Extract processed patch tokens processed_patch_tokens = out[:, p:, :] # (B, S, D) - # 5. Compute pathway scores via cosine similarity - # L2-normalize both sets of tokens to produce cosine similarities in [-1, 1] - # 5. Compute pathway scores via scaled dot product + softplus with learnable scale/bias + # 5. Compute pathway scores via a scaled dot-product + Softplus head. + # The raw affinity (patch · pathway) / sqrt(d) is shifted by a learnable + # per-pathway scale and bias, then Softplus maps it to a non-negative + # activity score (matching the mean-log1p-CP10k targets). This is *not* + # cosine similarity — the tokens are deliberately not L2-normalised. d_dim = processed_patch_tokens.shape[-1] if return_dense: diff --git a/src/spatial_transcript_former/models/mil.py b/src/spatial_transcript_former/models/mil.py index 2fcfaf6..ea65d21 100644 --- a/src/spatial_transcript_former/models/mil.py +++ b/src/spatial_transcript_former/models/mil.py @@ -1,8 +1,8 @@ """ -Multiple-Instance Learning (MIL) baselines for slide-level gene expression. +Multiple-Instance Learning (MIL) baselines for slide-level pathway activity. Models here aggregate an arbitrary number of patch features into a single -slide-level prediction, trained with bag-level (mean-expression) supervision. +slide-level prediction, trained with bag-level (mean-activity) supervision. * ``AttentionMIL`` — gated attention pooling (Ilse et al., 2018) * ``TransMIL`` — Nyström Transformer with PPEG positional encoding @@ -20,7 +20,7 @@ class AttentionMIL(nn.Module): """ Gated Attention Multiple Instance Learning (MIL) for regression. - Aggregates patch features to predict a global slide-level gene expression vector. + Aggregates patch features to predict a global slide-level pathway activity vector. Reference: Ilse et al. (2018). "Attention-based Deep Multiple Instance Learning." ICML. diff --git a/src/spatial_transcript_former/predict.py b/src/spatial_transcript_former/predict.py index 5189444..2138d81 100644 --- a/src/spatial_transcript_former/predict.py +++ b/src/spatial_transcript_former/predict.py @@ -6,7 +6,7 @@ * :class:`FeatureExtractor` — wraps a backbone (ResNet, Phikon, …) to turn raw image patches into feature embeddings. * :class:`Predictor` — wraps a trained :class:`SpatialTranscriptFormer` model - to predict gene expression from features + spatial coordinates. + to predict pathway activity from features + spatial coordinates. * :func:`inject_predictions` — injects predictions into an AnnData object for seamless Scanpy integration. @@ -14,6 +14,7 @@ training validation. """ +import json import os from typing import Dict, List, Optional, Union @@ -126,10 +127,10 @@ class Predictor: predictor = Predictor(model, device="cuda") # Single patch - genes = predictor.predict_patch(image_tensor) + pathways = predictor.predict_patch(image_tensor) # Whole slide (pre-extracted features) - genes = predictor.predict_wsi(features, coords) + pathways = predictor.predict_wsi(features, coords) """ def __init__( @@ -149,7 +150,7 @@ def __init__( self.model.eval() self.use_amp = use_amp - # Expose gene names if the model has them (set by from_pretrained) + # Expose pathway names if the model has them (set by from_pretrained) self.pathway_names: Optional[List[str]] = getattr(model, "pathway_names", None) @torch.no_grad() @@ -533,3 +534,119 @@ def plot_training_summary( plt.savefig(save_path, dpi=200, bbox_inches="tight", facecolor="#0f172a") plt.close(fig) print(f"Saved pathway summary to {save_path}") + + +# ═══════════════════════════════════════════════════════════════════════ +# CLI entry point (stf-predict) +# ═══════════════════════════════════════════════════════════════════════ + + +def _load_run_config(run_dir: str) -> dict: + """Read the training run's saved config from ``results_summary.json``.""" + summary_path = os.path.join(run_dir, "results_summary.json") + if not os.path.isfile(summary_path): + raise FileNotFoundError( + f"results_summary.json not found in {run_dir!r}. " + "Point --run-dir at a directory produced by stf-train." + ) + with open(summary_path, "r") as f: + summary = json.load(f) + # The trainer nests the run configuration under "config". + return summary.get("config", summary) + + +def _load_weights_into(model, run_dir: str, model_type: str, device: str) -> None: + """Load best (else latest) checkpoint weights into ``model`` in place.""" + candidates = [ + os.path.join(run_dir, f"best_model_{model_type}.pth"), + os.path.join(run_dir, f"latest_model_{model_type}.pth"), + ] + ckpt_path = next((p for p in candidates if os.path.isfile(p)), None) + if ckpt_path is None: + raise FileNotFoundError( + f"No checkpoint found in {run_dir!r} (looked for " + f"best_model_{model_type}.pth / latest_model_{model_type}.pth)." + ) + print(f"Loading checkpoint: {ckpt_path}") + ckpt = torch.load(ckpt_path, map_location=device, weights_only=True) + state_dict = ckpt.get("model_state_dict", ckpt) if isinstance(ckpt, dict) else ckpt + # Strip the torch.compile wrapper prefix if the checkpoint was compiled. + if any(k.startswith("_orig_mod.") for k in state_dict): + state_dict = {k.replace("_orig_mod.", "", 1): v for k, v in state_dict.items()} + model.load_state_dict(state_dict, strict=False) + + +def main(): + """CLI entry point for ``stf-predict``. + + Rebuilds a trained model from a run directory (``results_summary.json`` + + ``best_model_.pth``, as written by ``stf-train``), runs inference on + one HEST sample, and writes a spatial pathway-activity plot (prediction vs. + ground truth) to ``--output-dir``. + """ + import argparse + from types import SimpleNamespace + + from spatial_transcript_former.training.builder import setup_model + from spatial_transcript_former.visualization import run_inference_plot + + parser = argparse.ArgumentParser( + prog="stf-predict", + description="Run pathway-activity inference for a HEST sample and save " + "a spatial prediction-vs-truth plot.", + ) + parser.add_argument( + "--run-dir", + required=True, + help="Training run directory containing results_summary.json and " + "best_model_.pth (as written by stf-train).", + ) + parser.add_argument( + "--sample-id", required=True, help="HEST sample ID to run on (e.g. TENX156)." + ) + parser.add_argument( + "--output-dir", default="./results", help="Where to save the output plot." + ) + parser.add_argument( + "--data-dir", + default=None, + help="Override the data directory stored in the run config " + "(useful when predicting on a different machine).", + ) + parser.add_argument( + "--epoch", type=int, default=0, help="Epoch label for the output filename." + ) + parser.add_argument( + "--device", + default=None, + help="Torch device (default: cuda if available, else cpu).", + ) + cli = parser.parse_args() + + config = _load_run_config(cli.run_dir) + run_args = SimpleNamespace(**config) + + # Inference-time overrides: never compile or re-download backbone weights + # (the checkpoint already carries them), and route outputs as requested. + run_args.compile = False + run_args.pretrained = False + run_args.output_dir = cli.output_dir + if cli.data_dir is not None: + run_args.data_dir = cli.data_dir + if not getattr(run_args, "pathway_targets_dir", None): + run_args.pathway_targets_dir = os.path.join( + run_args.data_dir, "pathway_activities" + ) + + device = cli.device or ("cuda" if torch.cuda.is_available() else "cpu") + + model = setup_model(run_args, device) + _load_weights_into(model, cli.run_dir, run_args.model, device) + model.eval() + + print(f"Running inference for sample {cli.sample_id!r} on {device}...") + run_inference_plot(model, run_args, cli.sample_id, cli.epoch, device) + + +if __name__ == "__main__": + main() diff --git a/src/spatial_transcript_former/training/engine.py b/src/spatial_transcript_former/training/engine.py index 98a458a..551e55b 100644 --- a/src/spatial_transcript_former/training/engine.py +++ b/src/spatial_transcript_former/training/engine.py @@ -53,7 +53,7 @@ def _compute_masked_mse(preds, targets, mask): def _compute_bag_target(genes, mask): - """Average gene expression over valid spots (bag-level target for MIL).""" + """Average pathway targets over valid spots (bag-level target for MIL).""" mask_float = (~mask).float().unsqueeze(-1) # (B, N, 1) return (genes * mask_float).sum(dim=1) / mask_float.sum(dim=1) # (B, G) diff --git a/src/spatial_transcript_former/training/losses.py b/src/spatial_transcript_former/training/losses.py index 2c72f87..a019f35 100644 --- a/src/spatial_transcript_former/training/losses.py +++ b/src/spatial_transcript_former/training/losses.py @@ -54,13 +54,14 @@ class PCCLoss(nn.Module): """ Pearson Correlation Coefficient Loss. - Computes **gene-wise** spatial correlation, then averages across genes. - This measures whether the predicted spatial pattern (map) for each gene + Computes **per-pathway** spatial correlation, then averages across pathways. + This measures whether the predicted spatial pattern (map) for each pathway matches the pattern of the ground truth map, regardless of overall intensity. - PCC is scale-invariant: highly expressed genes and lowly expressed - genes contribute equally, making this objective robust to typical - spatial transcriptomics expression imbalances. + PCC is scale-invariant: high-activity and low-activity pathways contribute + equally, making this objective robust to the wide range of activity + magnitudes across pathways. (The ``G`` dimension in the shapes below is the + number of pathways; the loss is generic over any per-channel target.) """ def __init__(self, eps=1e-8):