diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index ead7736..3fbc459 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -12,7 +12,7 @@ jobs: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 with: - python-version: "3.10.20" + python-version: "3.12" - name: Install scalr requirements run: | pip install -r requirements.txt diff --git a/.github/workflows/publish-to-pypi.yml b/.github/workflows/publish-to-pypi.yml index 2a6a8ab..0be66c4 100644 --- a/.github/workflows/publish-to-pypi.yml +++ b/.github/workflows/publish-to-pypi.yml @@ -15,7 +15,7 @@ jobs: - name: Set up Python uses: actions/setup-python@v5 with: - python-version: "3.10.20" + python-version: "3.12" - name: Install pypa/build run: >- python3 -m diff --git a/.github/workflows/run_isort.yml b/.github/workflows/run_isort.yml index ee4332f..bc8c605 100644 --- a/.github/workflows/run_isort.yml +++ b/.github/workflows/run_isort.yml @@ -10,10 +10,10 @@ jobs: steps: - uses: actions/checkout@v2 - - name: Set up Python 3.10.20 + - name: Set up Python 3.12 uses: actions/setup-python@v2 with: - python-version: "3.10.20" + python-version: "3.12" - name: Install isort run: pip install isort - name: Run isort diff --git a/.github/workflows/run_pytest.yml b/.github/workflows/run_pytest.yml index be5af6f..6e41b09 100644 --- a/.github/workflows/run_pytest.yml +++ b/.github/workflows/run_pytest.yml @@ -11,7 +11,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.10.20"] + python-version: ["3.11", "3.12", "3.13", "3.14"] steps: - uses: actions/checkout@v3 @@ -22,6 +22,6 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - if [ -f requirements.txt ]; then pip install -r requirements.txt; fi + python -m pip install -e ".[dev]" - name: Running test cases run: pytest -v -s diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3b8cfeb --- /dev/null +++ b/.gitignore @@ -0,0 +1,24 @@ +# Python +__pycache__/ +*.py[cod] +*.egg-info/ +build/ +dist/ + +# Virtual environments +.venv/ +venv/ +env/ + +# Local run artifacts (quickstart/README commands write here by default) +nohup.out +example_data/ +models/ + +# IDE / OS +.vscode/ +.idea/ +.DS_Store + +# Docs build output +_build/ diff --git a/.venv/pyvenv.cfg b/.venv/pyvenv.cfg new file mode 100644 index 0000000..0c0e8ef --- /dev/null +++ b/.venv/pyvenv.cfg @@ -0,0 +1,5 @@ +home = /bin +include-system-site-packages = false +version = 3.12.3 +executable = /usr/bin/python3.12 +command = /bin/python3 -m venv /home/saurabh/infocusp_git/scaLR/.venv diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..890f5b3 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,46 @@ +# Changelog + +All notable changes to scaLR are documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). + +## [Unreleased] + +## [1.2.0] - 2026-10-13 + +### Added +- Hugging Face Hub integration: `scalr.models.push`/`pull` and `scalr models push`/`pull` to publish and download model artifacts (`scalr/models.py`). +- Dataset drift detection: `scalr.detect_drift` / `scalr drift` to check whether a query dataset's distribution has shifted from a model's training data (`scalr/drift.py`). +- Explainability API: `scalr.explain` / `scalr explain` for gene-attribution explanations of model predictions (`scalr/explain.py`). +- Streaming/chunked prediction (`streaming=True`, `chunk_size=...`) in `scalr.annotate`/`AnnotationModel.predict` for bounded-memory inference on large datasets. +- Support for custom/user-defined loss functions via the `name`/`params` builder pattern (`scalr/nn/loss/_loss.py`). +- `scalr` CLI: `train`, `annotate`, `explain`, `evaluate`, `drift`, and `models {list,info,push,pull}` subcommands, plus a `[dev]` extra and Development setup docs in the README. +- `.gitignore` for local run artifacts (`.venv/`, `__pycache__/`, `example_data/`, `models/`, build output). + +### Changed +- Feature scoring in `FeatureExtractionPipeline` parallelized across feature subsets. +- `DataIngestionPipeline`: `train_val_test` config checks consolidated into a single `_validate_train_val_test_config` method, run up front instead of being interleaved with split generation. + +### Fixed +- Gene alignment when a dataset has duplicate gene names (`scalr/genes.py`, `scalr/feature/feature_subsetting.py`). +- Multi-worker config propagation during chunked preprocessing. +- `anndata` bumped past a broken 0.11.x chunked-read regression; fixed a `pydeseq2` `np.matrix` incompatibility in pseudobulk DGE (`scalr/analysis/dge_pseudobulk.py`). +- Disabled `AnnCollection` dtype harmonization, which was corrupting chunked `.h5ad` reads. + +### Removed +- Dropped Python 3.10 support (`requires-python`, CI matrix, README). The `anndata>=0.12` floor needed for the chunked-read fix above requires Python >=3.11, so 3.10 installs could no longer resolve a compatible `anndata` version. + +### Performance +- Lazy-import scanpy/SHAP-backed submodules (`scalr/analysis`, `scalr/feature/scoring`) to reduce base import time. + +### Documentation +- `README.md` restructured into a concise entry point, with detailed content split out into linked guides: [docs/guide_simple_api.md](docs/guide_simple_api.md), [docs/guide_cli.md](docs/guide_cli.md), [docs/guide_pipeline.md](docs/guide_pipeline.md), and [docs/tutorials.md](docs/tutorials.md). +- Added [CONTRIBUTING.md](CONTRIBUTING.md) for development setup, tests, and formatting instructions. +- Added Sphinx autodoc coverage for the entire simple-API layer (`docs/scalr.simple_api.rst`: `api`, `artifact`, `benchmark`, `calibration`, `cli`, `doublet`, `drift`, `explain`, `feature_stability`, `genes`, `hierarchy`, `leakage`, `metrics`, `model_card`, `models`, `refinement`, `result`, `validation`), which previously had no published documentation at all. +- Fixed broken/orphaned Sphinx toctree entries (`docs/modules.rst`, `docs/scalr.rst`) and resolved a duplicate-object-description warning on `scalr.result.PredictionResult`. +- Added the previously-missing SHAP analysis tutorial to the tutorials list ([docs/tutorials.md](docs/tutorials.md)). +- Fixed stale `anndata==0.10.9` pins in `dge.ipynb`, `gene_recall_curve.ipynb`, `normalization.ipynb`, and `batch_correction.ipynb` (incompatible with the `anndata>=0.12` floor above), and a stale feature-branch image URL in `scalr_pipeline.ipynb`/`scalr_pipeline_local_run.ipynb`. + +## [1.1.0] + +Baseline prior to this changelog's introduction. See `git log` for history before this point. diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..0e3f225 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,90 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## What this is + +scaLR (Single-cell analysis using Low Resource) is a Python library/pipeline for scRNA-seq cell-type classification, feature selection, and biomarker analysis, designed to work on large datasets with limited compute. It has two layers: + +1. **Simple library API** (`scalr.train`, `scalr.annotate`, `scalr.load_model`, `scalr.validate`, `scalr.align_genes`) — an inference-first API added on top of the engine for quick, config-free use. See `scalr/api.py`, `scalr/artifact.py`, `scalr/result.py`. +2. **Configuration-driven pipeline** (`pipeline.py` + a YAML config) — the original, full-featured research pipeline: data ingestion/splitting/preprocessing → feature extraction → model training → evaluation & downstream analysis. This remains the primary way to run large, customized experiments. + +Both layers share the same underlying engine in `scalr/` (models, dataloaders, trainers, losses, splitters, preprocessors, feature scorers/selectors, analysis modules). + +## Commands + +Install (editable, for development): +```bash +pip install -r requirements.txt +pip install -e . +``` +CPU is the default `torch` install; for CUDA, install a CUDA-specific `torch` build per the [PyTorch instructions](https://pytorch.org/get-started/locally/) before or after. `requirements.txt` uses bounded version ranges so pip can select wheels for the active Python version while avoiding major-version compatibility breaks. + +`pip install -e .` also registers the `scalr` console script (`scalr validate|train|annotate|evaluate|models ...`, see `scalr/cli.py`), a thin wrapper around the simple API below. + +Run the full config-driven pipeline: +```bash +python pipeline.py --config /path/to/config.yaml -l -m +``` +(`-l`/`--log` enables experiment logging, `-m`/`--memoryprofiler` enables memory profiling.) + +Tests (pytest, no separate test directory — test files live next to the code as `test_*.py`): +```bash +pytest # run everything +pytest -v -s # verbose, as CI does +pytest scalr/test_api.py # a single test file +pytest scalr/test_api.py::test_train_returns_annotation_model_with_metrics # a single test +``` + +Formatting/import order (both enforced in CI, config in `.style.yapf` / `.isort.cfg`, both based on the `google` style, 80-col): +```bash +yapf -ir --style .style.yapf . +isort . --settings-path .isort.cfg +``` +`pre-commit` is configured (`.pre-commit-config.yaml`) to run both automatically. + +Docs (Sphinx, published to GitHub Pages on push to `main`): +```bash +sphinx-build docs _build +``` + +## Architecture + +### Config-driven pipeline (`scalr/*_pipeline.py`, driven by `pipeline.py`) + +`pipeline.py` reads one YAML config and runs whichever top-level sections are present (sections can be omitted/commented out to skip stages): + +1. **`data`** → `DataIngestionPipeline` (`data_ingestion_pipeline.py`): splits data into train/val/test (via `scalr/data/split/`), applies preprocessing (`scalr/data/preprocess/`), generates label mappings. +2. **`feature_selection`** → `FeatureExtractionPipeline` (`feature_extraction_pipeline.py`): trains on feature subsets batch-wise, scores features (`scalr/feature/scoring/`: linear-weight or SHAP-based), and selects top-k features (`scalr/feature/selector/`). +3. **`final_model`**/training → `ModelTrainingPipeline` (`model_training_pipeline.py`): builds model/optimizer/loss/callbacks and trains via a `Trainer` (`scalr/nn/trainer/`). +4. **`analysis`** → `EvalAndAnalysisPipeline` (`eval_and_analysis_pipeline.py`): evaluation metrics/classification report, gene/biomarker analysis, and pluggable downstream analyses (`scalr/analysis/`: DGE via pseudobulk or linear mixed models, ROC-AUC, gene recall curves, heatmaps). + +Each stage returns an updated config dict so the full resolved config (including defaults) gets written back out — this is how `config/README.md` params map to code. + +### The `name`/`params` builder pattern + +Almost every pluggable component (models, dataloaders, loss functions, splitters, preprocessors, feature scorers/selectors, callbacks, analysers) is instantiated from a config dict of the form `{name: , params: {...}}` via `scalr.utils.build_object` (`scalr/utils/misc_utils.py`), which looks up `` in the given module, merges `params` over the class's `get_default_params()`, and constructs it. Each such subpackage exposes a `build_*` convenience function next to its base class (e.g. `build_model` in `scalr/nn/model/_model.py`, `build_preprocessor` in `scalr/data/preprocess/_preprocess.py`, `build_splitter`, `build_analyser`, `build_dataloader`, `build_loss_fn`). When adding a new implementation of any of these, follow the existing sibling classes' pattern: subclass the module's base class, implement its abstract methods, and implement `get_default_params()`. + +### Data handling + +- Input is `.h5ad` (AnnData) or a directory of chunked `.h5ad` files, read via `scalr.utils.read_data`/`write_data` (`scalr/utils/file_utils.py`), which returns an `AnnCollection` for chunked directories to avoid loading everything into memory. +- `adata.X` is expected to already be normalized (`log1p`, range 0–10) for the pipeline convention; `adata.obs` holds cell metadata including the classification target column; `adata.var` holds gene names. +- Chunkwise writing/transformation (`write_chunkwise_data`) is the mechanism used throughout to keep peak memory bounded regardless of dataset size — this is central to the "low resource" design goal, so avoid introducing full-dataset `.toarray()`/dense materialization in new code. +- `scalr.data.split` splitters (`StratifiedSplitter`, `GroupSplitter`, `StratifiedGroupSplitter`) matter for scientific validity: `GroupSplitter`/`StratifiedGroupSplitter` keep a `stratify` column (e.g. `donor_id`) confined to a single split to prevent train/test leakage — prefer these over plain stratified splitting whenever a donor/patient/sample identifier is available. + +### Simple API layer (`scalr/api.py`, `scalr/artifact.py`, `scalr/result.py`, plus one module per capability) + +This is a newer, self-contained layer built on top of the same engine (`scalr.nn.model`, etc.) but bypassing the config/pipeline machinery for simple in-memory workflows. Each capability lives in its own top-level module and is wired into `AnnotationModel`/`api.py` as an optional parameter, not a separate code path: + +- `scalr.train(adata, labels_key, group_key=...)` does leakage-safe splitting (`scalr/leakage.py`), class-weighted training, temperature-scaling calibration (`scalr/calibration.py`), and macro/weighted-F1 + balanced-accuracy evaluation (`scalr/metrics.py`), returning an `AnnotationModel`. Pass `taxonomy={fine_label: broad_label, ...}` to enable hierarchical annotation. +- `AnnotationModel.save()`/`scalr.load_model()` (`scalr/artifact.py`) persist/restore a **self-contained model artifact directory** (`manifest.json`, `model.pt`, `model_config.json`, `label_mapping.json`, `features.json`, `preprocessing.json`, `calibration.json`, `metrics.json`, `taxonomy.json` if set, and an auto-generated `README.md` model card via `scalr/model_card.py`) — no external config needed to reload and run inference. `scalr.load_model` also resolves a bare name via the local model registry (`scalr/models.py`, `scalr.models.list/info/register`; `download()` intentionally raises — no hosted hub yet). +- `scalr.annotate(adata, model=...)` / `AnnotationModel.predict()` validate the input (`scalr/validation.py`), align its genes to the model's expected feature list (`scalr/genes.py`, zero-filling missing genes), run inference, and return a typed `PredictionResult` (`scalr/result.py`) with calibrated `confidence`, `entropy`, `margin`, and an `is_unknown` open-set abstention flag — always aligned to `adata.obs_names` (`PredictionResult.write_to_adata` refuses to write if names don't match, to prevent silently misaligned predictions). Optional `predict()`/`annotate()` params layer on: `level="broad"` (aggregate via taxonomy, `scalr/hierarchy.py`), `cluster_refinement="auto"|` (confidence-weighted per-cluster majority vote, `scalr/refinement.py` — always preserves the pre-refinement labels in `result.metadata["raw_labels"]`), and `flag_doublets=True` (sets `result.is_possible_doublet`, a heuristic screening signal from `scalr/doublet.py`, not a validated doublet call). +- `scalr.compute_feature_stability(adata, labels_key, ...)` (`scalr/feature_stability.py`) repeats feature selection over resampled subsets (via `sklearn.LogisticRegression`, independent of the DNN training loop) and reports per-gene selection frequency and Jaccard stability — for judging whether a selected biomarker set is a stable signal or a split-specific artifact. +- `scalr.run_benchmark(datasets, labels_key, ...)` (`scalr/benchmark.py`) runs train+predict per dataset and reports both quality (macro-F1, balanced accuracy) and cost (wall time, cells/sec, peak RSS via `resource.getrusage`) together, plus environment metadata — never scale one dataset's numbers without the other. +- `scalr/cli.py` exposes `validate`/`train`/`annotate`/`evaluate`/`models list`/`models info` as the `scalr` console script (registered in `pyproject.toml` under `[project.scripts]`); it's a thin argparse wrapper around this same API, not a separate implementation. + +When extending this layer, keep it usable without any YAML config — pass all such new functionality through function parameters. This layer works with in-memory `AnnData` directly rather than the chunked/config-driven flow used by the pipeline, so it's not currently intended for atlas-scale/out-of-core datasets. + +### Logging + +Two logger classes in `scalr/utils/logger.py`: `FlowLogger` prints high-level stage progress to stdout (use this for anything user-facing/visible by default). `EventLogger` writes detailed logs to a file and is **silent by default** (`NullHandler`) unless a `filepath` or `stdout=True` is given — don't use it where output should be visible without extra configuration. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..a9d9b01 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,38 @@ +# Contributing to scaLR + +## Development setup + +``` +python -m pip install -r requirements.txt +python -m pip install -e ".[dev]" +pre-commit install +``` + +## Tests + +pytest, test files live next to the code as `test_*.py`, no separate test directory: + +``` +pytest # run everything +pytest -v -s # verbose, as CI does +pytest scalr/test_api.py # a single test file +``` + +## Formatting / import order + +Both enforced in CI (config in `.style.yapf` / `.isort.cfg`), and run automatically by `pre-commit`: + +``` +yapf -ir --style .style.yapf . +isort . --settings-path .isort.cfg +``` + +## Docs + +Sphinx, published to GitHub Pages on push to `main`: + +``` +sphinx-build docs _build +``` + +See [CHANGELOG.md](CHANGELOG.md) for a history of notable changes. diff --git a/README.md b/README.md index 45d992a..97b92cc 100644 --- a/README.md +++ b/README.md @@ -1,10 +1,10 @@ - + # Single-cell analysis using Low Resource (scaLR) [![GitHub](https://img.shields.io/github/license/InFoCusp/scaLR)](https://github.com/infocusp/scaLR?tab=GPL-3.0-1-ov-file#) -[![Documentation](https://img.shields.io/badge/docs-v1.1.0-orange)](https://infocusp.github.io/scaLR/) +[![Documentation](https://img.shields.io/badge/docs-v1.2.0-orange)](https://infocusp.github.io/scaLR/) [![DOI](https://zenodo.org/badge/852710658.svg)](https://doi.org/10.5281/zenodo.13767942) @@ -22,356 +22,106 @@ **The below flowchart also explains the major steps of the scaLR platform.** -![image.jpg](img/Schematic-of-scPipeline.jpg) +![image.jpg](https://raw.githubusercontent.com/infocusp/scaLR/main/img/Schematic-of-scPipeline.jpg) -## Pre-requisites and installation scaLR +scaLR has two layers: a config-free **[simple API](https://github.com/infocusp/scaLR/blob/main/docs/guide_simple_api.md)** (`scalr.train`/`scalr.annotate`/`scalr.load_model`) for quick, in-memory use, and the original **[configuration-driven pipeline](https://github.com/infocusp/scaLR/blob/main/docs/guide_pipeline.md)** (`pipeline.py` + a YAML config) for large, customized research experiments. Both share the same underlying engine. +## Pre-requisites and installation -- ScaLR can be installed using git or pip. It is tested in Python 3.10.20 and it is recommended to use that environment. +- ScaLR supports Python 3.11-3.14. Dependencies use compatible version ranges instead of requiring one exact Python environment. New Python releases still need to be added to CI before they can be considered supported. ``` -conda create -n scaLR_env python=3.10.20 - +conda create -n scaLR_env python=3.12 conda activate scaLR_env ``` -- Using git +You can use `venv` instead of Conda: + +``` +python3.12 -m venv .venv +source .venv/bin/activate +python -m pip install --upgrade pip +``` + +Using git: ``` git clone https://github.com/infocusp/scaLR.git +cd scaLR -pip install -r requirements.txt +python -m pip install --upgrade pip +python -m pip install -r requirements.txt +python -m pip install -e . ``` -- Installation using pip + +Using pip: + ``` -pip install pyscaLR +python -m pip install pyscaLR +``` + +If pip reports that no compatible dependency version is available, check `python --version` and upgrade pip. The package supports Python 3.11-3.14, but individual dependency releases may lag behind a newly released Python version. + +**Note:** If you want to run the configuration-driven pipeline via the `pyscaLR` pip package, also clone/download `pipeline.py` and `config.yaml` from the git repository. + +The base install (`pip install pyscaLR` or `pip install -r requirements.txt`) pulls the CPU build of PyTorch and works out of the box on any machine. For CUDA acceleration, install a CUDA-enabled `torch` build for your CUDA version before or after installing scaLR, following the [official PyTorch instructions](https://pytorch.org/get-started/locally/) — scaLR auto-selects the device (`device="auto"`) at train/inference time. + +Contributing to scaLR itself (running tests, formatting, pre-commit hooks)? See [CONTRIBUTING.md](https://github.com/infocusp/scaLR/blob/main/CONTRIBUTING.md). + +## Quickstart: annotate cells in 3 lines + +```python +import scanpy as sc +import scalr + +adata = sc.read_h5ad("pbmc.h5ad") + +result = scalr.annotate(adata, model="models/pbmc_v1", device="auto") + +adata.obs["cell_type"] = result.labels +adata.obs["cell_type_confidence"] = result.confidence +adata.obs["cell_type_unknown"] = result.is_unknown ``` -**Note:** If the user wants to run the entire pipeline via installing pip pyscalr, they should clone/download these files(`pipeline.py` and `config.yaml`) from the git repository. - -## Input data format -- Currently the pipeline expects all datasets in [anndata](https://anndata.readthedocs.io/en/latest/tutorials/notebooks/getting-started.html) formats (`.h5ad` files only). -- The anndata object should contain cell samples as `obs` and genes as `var. ' -- `adata.X`: contains normalized gene counts/expression values (`log1p` normalization with range `0-10` expected). -- `adata.obs`: contains any metadata regarding cells, including a column for `target` which will be used for classification. The index of `adata.obs` is cell_barcodes. -- `adata.var`: contains all gene_names as an Index. - - -## How to run - -1. It is necessary that the user modify the configuration file, and each stage of the pipeline is available inside the config folder [config.yml] as per your requirements. Simply omit/comment out stages of the pipeline you do not wish to run. -2. Refer **config.yml** & **it's detailed config** [README](https://github.com/infocusp/scaLR/blob/main/config/README.md) file on how to use different parameters and files. -3. Then use the `pipeline.py` file to run the entire pipeline according to your configurations. This file takes as argument the path to config (`-c | --config`), along with optional flags to log all parts of the pipelines (`-l | --log`) and to analyze memory usage (`-m | --memoryprofiler`). -5. `python pipeline.py --config /path/to/config.yaml -l -m` to run the scaLR. - -## Example configs - -### Config for cell type classification and biomarker identification - -NOTE: Below are just suggestions for the model parameters. Feel free to play around with them for tuning the model & improving the results. - -An example configuration file for the current dataset, incorporating the edits below, can be found at '`scaLR/tutorials/pipeline/config_celltype.yaml`. Update the device as cuda or cpu as per the requirement. - -- **Device setup*** - - Update device: 'cuda' for GPU enabled runtype, else device: 'cpu' for CPU enabled runtype. -- **Experiment Config** - - The default exp_run number is 0.If not changed, the celltype classification experiment would be exp_run_0 with all the pipeline results. -- **Data Config** - - Update the full_datapath to `data/modified_adata.h5ad` (as we will include GeneRecallCurve in the downstream). - - Specify the num_workers value for effective parallelization. - - Set target to cell_type. -- **Feature Selection** - - Specify the num_workers value for effective parallelization. - - Update the model layers to [5000, 10], as there are only 10 cell types in the dataset. - - Change epoch to 10. -- **Final Model Training** - - Update the model layers to the same as for feature selection: [5000, 10]. - - Change epoch to 100. -- **Analysis** - - Downstream Analysis - - Uncomment the test_samples_downstream_analysis section. - - Update the reference_genes_path to `scaLR/tutorials/pipeline/grc_reference_gene.csv`. - - Refer to the section below: - ``` - # Config file for pipeline run for cell type classification. - - # DEVICE SETUP. - device: 'cuda' - - # EXPERIMENT. - experiment: - dirpath: 'scalr_experiments' - exp_name: 'exp_name' - exp_run: 0 - - # DATA CONFIG. - data: - sample_chunksize: 20000 - - train_val_test: - full_datapath: 'data/modified_adata.h5ad' - num_workers: 2 - - splitter_config: - name: GroupSplitter - params: - split_ratio: [7, 1, 2.5] - stratify: 'donor_id' - - # split_datapaths: '' - - # preprocess: - # - name: SampleNorm - # params: - # **args - - # - name: StandardScaler - # params: - # **args - - target: cell_type - - # FEATURE SELECTION. - feature_selection: - - # score_matrix: '/path/to/matrix' - feature_subsetsize: 5000 - num_workers: 2 - - model: - name: SequentialModel - params: - layers: [5000, 10] - weights_init_zero: True - - model_train_config: - trainer: SimpleModelTrainer - - dataloader: - name: SimpleDataLoader - params: - batch_size: 25000 - padding: 5000 - - optimizer: - name: SGD - params: - lr: 1.0e-3 - weight_decay: 0.1 - - loss: - name: CrossEntropyLoss - - epochs: 10 - - scoring_config: - name: LinearScorer - - features_selector: - name: AbsMean - params: - k: 5000 - - # FINAL MODEL TRAINING. - final_training: - - model: - name: SequentialModel - params: - layers: [5000, 10] - dropout: 0 - weights_init_zero: False - - model_train_config: - resume_from_checkpoint: null - - trainer: SimpleModelTrainer - - dataloader: - name: SimpleDataLoader - params: - batch_size: 15000 - - optimizer: - name: Adam - params: - lr: 1.0e-3 - weight_decay: 0 - - loss: - name: CrossEntropyLoss - - epochs: 100 - - callbacks: - - name: TensorboardLogger - - name: EarlyStopping - params: - patience: 3 - min_delta: 1.0e-4 - - name: ModelCheckpoint - params: - interval: 5 - analysis: - -     model_checkpoint: '' - -     dataloader: -         name: SimpleDataLoader -         params: -             batch_size: 15000 - -     gene_analysis: -         scoring_config: -             name: LinearScorer - -         features_selector: -             name: ClasswisePromoters -             params: -                 k: 100 -     test_samples_downstream_analysis: -         - name: GeneRecallCurve -           params: -             reference_genes_path: 'scaLR/tutorials/pipeline/grc_reference_gene.csv' -             top_K: 300 -             plots_per_row: 3 -             features_selector: -                 name: ClasswiseAbs -                 params: {} -         - name: Heatmap -           params: {} -         - name: RocAucCurve -           params: {} - ``` -### Config for clinical condition-specific biomarker identification and DGE analysis - -An example configuration file (`scaLR/tutorials/pipeline/config_clinical.yaml`). Update the device as CUDA or CPU as per the requirement. - -- Experiment Config - - Make sure to change the exp_run number if you have an experiment with the same number earlier related to cell classification. As we have done one experiment earlier, we'll change the number now to '1'. -- Data Config - - The full_datapath remains the same as above. - - Change the target to disease (this column contains data for clinical conditions, COVID-19/normal). -- Feature Selection - - Update the model layers to [5000, 2], as there are only two types of clinical conditions. - - epoch as 10. -- Final Model Training - -Update the model layers to the same as for feature selection: [5000, 2]. - - epoch as 100. -- Analysis - - Downstream Analysis - - Uncomment the full_samples_downstream_analysis section for example config file. - - We are not performing the 'gene_recall_curve' analysis in this case. It can be performed if the COVID-19/normal specific genes are available, but there are many possibilities of genes in the case of normal conditions. - - There are two options to perform differential gene expression (DGE) analysis: **DgePseudoBulk and DgeLMEM**. The parameters are updated as follows. Note that DgeLMEM may take a bit more time, as the multiprocessing is not very efficient with only 2 CPUs in the current Colab runtime. - - Refer to the section below: - - ``` - analysis: - -   model_checkpoint: '' - -   dataloader: -       name: SimpleDataLoader -       params: -           batch_size: 15000 - -   gene_analysis: -       scoring_config: -           name: LinearScorer - -       features_selector: -           name: ClasswisePromoters -           params: -               k: 100 -   full_samples_downstream_analysis: -       - name: Heatmap -         params: -           top_n_genes: 100 -       - name: RocAucCurve -         params: {} -       - name: DgePseudoBulk -         params: -             celltype_column: 'cell_type' -             design_factor: 'disease' -             factor_categories: ['COVID-19', 'normal'] -             sum_column: 'donor_id' -             cell_subsets: ['conventional dendritic cell', 'natural killer cell'] -       - name: DgeLMEM -         params: -           fixed_effect_column: 'disease' -           fixed_effect_factors: ['COVID-19', 'normal'] -           group: 'donor_id' -           celltype_column: 'cell_type' -           cell_subsets: ['conventional dendritic cell'] -           gene_batch_size: 1000 -           coef_threshold: 0.1 - ``` + +To train your own model: + +```python +model = scalr.train(adata, labels_key="cell_type", group_key="donor_id") +model.save("models/pbmc_v1") +``` + +`scalr.train` performs leakage-safe (donor-grouped) splitting, class-imbalance-aware training, confidence calibration, and reports macro-F1/balanced accuracy on held-out data. `scalr.annotate`/`model.predict` validate the input, align its genes to the model's expected feature set, and return a `PredictionResult` with calibrated `confidence`, `entropy`, `margin` and an `is_unknown` abstention flag — see [scalr/api.py](https://github.com/infocusp/scaLR/blob/main/scalr/api.py) and [scalr/result.py](https://github.com/infocusp/scaLR/blob/main/scalr/result.py). + +**The [Simple API guide](https://github.com/infocusp/scaLR/blob/main/docs/guide_simple_api.md) covers the rest of this layer in depth:** data validation, gene alignment, streaming/bounded-memory prediction, self-contained model artifacts, leakage-safe evaluation, hierarchical (coarse-to-fine) annotation, cluster-aware refinement & doublet screening, explainability, dataset drift detection, feature-selection stability, model cards/registry/Hugging Face Hub, and the benchmark suite. + +## Command-line interface + +The simple API is also available from the shell as the `scalr` command (`validate`, `train`, `annotate`, `explain`, `evaluate`, `drift`, `models`, and `analyze` for the config-driven pipeline): + +```bash +pip install -e . # registers the `scalr` command (or: pip install pyscaLR) + +scalr validate --input data.h5ad --labels-key cell_type --group-key donor_id +scalr train --input train.h5ad --labels-key cell_type --output models/pbmc_v1 +scalr annotate --input query.h5ad --model models/pbmc_v1 --output annotated.h5ad +``` + +**See the [CLI guide](https://github.com/infocusp/scaLR/blob/main/docs/guide_cli.md) for every subcommand, its flags, and what it writes.** + +## Configuration-driven pipeline + +For large, customized research experiments, run the full YAML-config-driven pipeline (data ingestion/splitting/preprocessing → feature extraction → model training → evaluation & downstream analysis): + +```bash +python pipeline.py --config /path/to/config.yaml -l -m +``` + +**See the [Pipeline guide](https://github.com/infocusp/scaLR/blob/main/docs/guide_pipeline.md)** for the input data format, how to run it, annotated example configs, and the experiment output structure, and [config/README.md](https://github.com/infocusp/scaLR/blob/main/config/README.md) for the full config parameter reference. ## Interactive tutorials -Detailed tutorials have been made on how to use some pipeline functionalities as a scaLR library. Find the links below. - -- **scaLR pipeline** [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infocusp/scaLR/blob/main/tutorials/pipeline/scalr_pipeline.ipynb) -- **Differential gene expression analysis** [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infocusp/scaLR/blob/main/tutorials/analysis/differential_gene_expression/dge.ipynb) -- **Gene recall curve** [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infocusp/scaLR/blob/main/tutorials/analysis/gene_recall_curve/gene_recall_curve.ipynb) -- **Normalization** [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infocusp/scaLR/blob/main/tutorials/preprocessing/normalization.ipynb) -- **Batch correction** [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infocusp/scaLR/blob/main/tutorials/preprocessing/batch_correction.ipynb) - -- **An example of jupyter notebook to [run scaLR in local machine](https://github.com/infocusp/scaLR/blob/main/tutorials/pipeline/scalr_pipeline_local_run.ipynb)**. - -## Experiment output structure -- **pipeline.py**: -The main script that performs an end-to-end run. - - `exp_dir`: root experiment directory for the storage of all step outputs of the platform specified in the config. - - `config.yml`: copy of config file to reproduce the user-defined experiment. - -- **data_ingestion**: -Reads the data and splits it into Train/Validation/Test sets for the pipeline. Then, it performs sample-wise normalization on the data. - - `exp_dir` - - `data` - - `train_val_test_split.json`: contains sample indices for train/validation/test splits. - - `label_mappings.json`: contains mappings of all metadata columns between labels and IDs. - - `train_val_test_split`: directory containing the train, validation, and test samples and data files. - -- **feature_extraction**: -Performs feature selection and extraction of new datasets containing a subset of features. - - `exp_dir` - - `feature_extraction` - - `chunked_models`: contains weights of each model trained on feature subset data (refer to feature subsetting algorithm). - - `feature_subset_data`: directory containing the new feature-subsetted train, val, and test samples anndatas. - - `score_matrix.csv`: combined scores of all individual models for each feature and class. shape: n_classes X n_features. - - `top_features.json`: a file containing a list of top features selected / to be subsetted from total features. - -- **final_model_training**: -Trains a final model based on `train_datapath` and `val_datapath` in config. - - `exp_dir` - - `model` - - `logs`: directory containing Tensorboard Logs for the training of the model. - - `checkpoints`: directory containing model weights checkpointed at every interval specified in config. - - `best_model`: the best model checkpoint contains information to use model for inference/resume training. - - `model_config.yaml`: config file containing model parameters. - - `mappings.json`: contains mapping of class_names to class_ids used by model during training. - - `model.pt`: contains model weights. - -- **eval_and_analysis**: -Performs evaluation of best model trained on user-defined metrics on the test set. Also performs various downstream tasks. - - `exp_dir` - - `analysis` - - `classification_report.csv`: contains classification report showing Precision, Recall, F1, and accuracy metrics for each class on the test set. - - `gene_analysis` - - `score_matrix.csv`: score of the final model, for each feature and class. shape: n_classes X n_features. - - `top_features.json`: a file containing a list of selected top features/biomarkers. - - `test_samples/full_samples` - - `heatmaps` - - `class_name.svg`: heatmap for top genes of a particular class w.r.t those genes association in other classes. E.g., B.svg, C.svg, etc. - - `roc_auc.svg`: contains ROC-AUC plot for all classes. - - `gene_recall_curve.svg`: contains gene recall curve plots. - - `gene_recall_curve_info.json`: contains reference genes list which are present in top_K ranked genes per class for each model. - - `pseudobulk_dge_result` - - `pbkDGE_celltype_factor_categories_0_vs_factor_categories_1.csv`: contains Pseudobulk DGE results between selected factor categories for a celltype. - - `pbkDGE_celltype_factor_categories_0_vs_factor_categories_1.svg`: volcano plot of Log2Foldchange vs -log10(p-value) of genes. - - `lmem_dge_result` - - `lmemDGE_celltype.csv`: contains LMEM DGE results between selected factor categories for a celltype. - - `lmemDGE_fixed_effect_factor_X.svg`: volcano plot of coefficient vs -log10(p-value) of genes. - + +Colab notebooks covering the pipeline, DGE/gene-recall/SHAP analysis, normalization, batch correction, and a local-run example — **see [Tutorials](https://github.com/infocusp/scaLR/blob/main/docs/tutorials.md)**. + ## Citation Jogani, S., Pol, A. S., Prajapati, M., Samal, A., Bhatia, K., Parmar, J., ... & Gupta, S. (2025). scaLR: a low-resource deep neural network-based platform for single cell analysis and biomarker discovery. Briefings in Bioinformatics, 26(3), bbaf243. diff --git a/config/README.md b/config/README.md index 92eff49..2b27267 100644 --- a/config/README.md +++ b/config/README.md @@ -93,6 +93,10 @@ If mentioned, the feature subsetting process will be skipped and this matrix wil default: `5000` - Chunks of features to subset data for training the model iteratively. +**num_workers** {int}: `int | null` +default: `1` +- Number of workers used to parallelize both feature-subset model training and, per chunked model, feature scoring afterwards. + **model** {dict}: Mention the name of the model class in `name` & its required parameters under `params` dict. Example below: @@ -132,7 +136,7 @@ Mention the model of training parameters like `trainer`, `dataloader`, `optimize - Current support/s in the platform for the following model configurations is/are: - trainer: `SimpleModelTrainer` - dataloader: `SimpleDataLoader|SimpleMetaDataLoader` - - loss: pytorch supported all inbuilt loss functions. The class name should match exactly with a torch.nn class names. + - loss: any built-in `torch.nn` loss class (name must match exactly), or a custom loss class subclassing `scalr.nn.loss.CustomLossBase` and registered under `scalr.nn.loss` (see [scalr/nn/loss/_loss.py](../scalr/nn/loss/_loss.py)). **scoring_config** {dict}: diff --git a/docs/conf.py b/docs/conf.py index 0ff726c..e4d97fd 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -2,7 +2,7 @@ project = 'scaLR' copyright = '2024, Infocusp Innovations' author = 'Infocusp Innovations' -release = 'v1.1.0' +release = 'v1.2.0' # -- General configuration --------------------------------------------------- import os diff --git a/docs/guide_cli.md b/docs/guide_cli.md new file mode 100644 index 0000000..1a41958 --- /dev/null +++ b/docs/guide_cli.md @@ -0,0 +1,109 @@ +# Command-line interface guide + +The [simple API](guide_simple_api.md) is also available from the shell as the `scalr` command. This is separate from, and does not replace, `python pipeline.py --config ...` (see the [Pipeline guide](guide_pipeline.md) for the configuration-driven pipeline); `scalr` wraps the in-memory simple API — see [scalr/cli.py](https://github.com/infocusp/scaLR/blob/main/scalr/cli.py). + +## 1. Install (registers the `scalr` command) + +```bash +pip install -r requirements.txt +pip install -e . # or: pip install pyscaLR +``` + +`pip install -e .` (or the packaged `pyscaLR`) registers the `scalr` console script via `[project.scripts]` in `pyproject.toml`. Verify it's on your `PATH`: + +```bash +scalr --help +``` + +## 2. Validate your data first + +```bash +scalr validate --input data.h5ad --labels-key cell_type --group-key donor_id +``` + +Prints the validation report (structure/numerical/label checks) and exits non-zero if the data is unusable — useful as a pre-flight check or CI gate. + +## 3. Train a model + +```bash +scalr train \ + --input train.h5ad \ + --labels-key cell_type \ + --group-key donor_id \ + --output models/pbmc_v1 \ + --epochs 15 \ + --device auto +``` + +Writes a self-contained model artifact directory to `models/pbmc_v1` (weights, manifest, calibration, a `README.md` model card, etc.) and prints macro-F1/balanced accuracy on the held-out split. + +## 4. Annotate new data with it + +```bash +scalr annotate \ + --input query.h5ad \ + --model models/pbmc_v1 \ + --output annotated.h5ad \ + --device auto \ + --flag-doublets +``` + +Writes predictions (`scalr_pred`, `scalr_confidence`, `scalr_unknown`, etc.) into `annotated.h5ad`'s `.obs`/`.obsm`/`.uns`. `--model` also accepts a name registered in the local model registry, not just a path. Add `--streaming --chunk-size 20000` for large inputs to gene-align and score in bounded-memory chunks instead of all at once. + +## 5. Explain predictions + +```bash +scalr explain --input query.h5ad --model models/pbmc_v1 --indices 10 20 30 --top-k 20 +scalr explain --model models/pbmc_v1 --class-name T_cell --top-k 50 +``` + +Prints, per requested cell, the genes whose expression most supported or contradicted the prediction; `--class-name` explains a class in general instead of specific cells. + +## 6. Evaluate against labeled test data + +```bash +scalr evaluate --input test.h5ad --model models/pbmc_v1 --labels-key cell_type +``` + +Prints macro-F1/weighted-F1/balanced-accuracy plus a per-class precision/recall/F1 table. + +## 7. Check for dataset drift + +```bash +scalr drift --input query.h5ad --model models/pbmc_v1 --fail-on-drift +``` + +Compares `query.h5ad` against the model's training-time gene-expression and library-size statistics and prints a drift report; `--fail-on-drift` exits non-zero if drift is detected, for use as a pre-flight/CI gate before trusting predictions on a new cohort. + +## 8. Manage the local model registry & Hugging Face Hub + +```bash +scalr models list +scalr models info human_pbmc +scalr models push models/pbmc_v1 your-org/human-pbmc # requires: pip install "pyscaLR[hub]" +scalr models pull your-org/human-pbmc # downloads and registers locally +``` + +Registering a model into the registry (`scalr.models.register(...)`) is currently Python-API-only, not yet a CLI subcommand. + +## 9. Run downstream analyses + +Configure the desired analyses under `analysis` in a YAML file. Supported analyses include `Heatmap`, `RocAucCurve`, `GeneRecallCurve`, `DgePseudoBulk`, and `DgeLMEM`. + +Run them through the CLI: + +```bash +scalr analyze \ + --config config/config.yaml \ + --log +``` + +Optional flags: + +```bash +--level INFO +--logpath scalr_experiments/analysis.log +--memoryprofiler +``` + +The command uses the existing configuration-driven pipeline and writes results under the `experiment.dirpath` configured in the YAML file. diff --git a/docs/guide_pipeline.md b/docs/guide_pipeline.md new file mode 100644 index 0000000..a0f7c20 --- /dev/null +++ b/docs/guide_pipeline.md @@ -0,0 +1,314 @@ +# Configuration-driven pipeline guide + +This is the original, full-featured research pipeline: data ingestion/splitting/preprocessing → feature extraction → model training → evaluation & downstream analysis, driven by a single YAML config and run via `pipeline.py`. For the newer config-free, in-memory API (`scalr.train`/`scalr.annotate`), see the [Simple API guide](guide_simple_api.md) instead. + +## Input data format + +- Currently the pipeline expects all datasets in [anndata](https://anndata.readthedocs.io/en/latest/tutorials/notebooks/getting-started.html) formats (`.h5ad` files only). +- The anndata object should contain cell samples as `obs` and genes as `var`. +- `adata.X`: contains normalized gene counts/expression values (`log1p` normalization with range `0-10` expected). +- `adata.obs`: contains any metadata regarding cells, including a column for `target` which will be used for classification. The index of `adata.obs` is cell_barcodes. +- `adata.var`: contains all gene_names as an Index. + +## How to run + +1. It is necessary that the user modify the configuration file, and each stage of the pipeline is available inside the config folder `config.yml` as per your requirements. Simply omit/comment out stages of the pipeline you do not wish to run. +2. Refer **config.yml** & its detailed config [README](https://github.com/infocusp/scaLR/blob/main/config/README.md) file on how to use different parameters and files. +3. Then use the `pipeline.py` file to run the entire pipeline according to your configurations. This file takes as argument the path to config (`-c | --config`), along with optional flags to log all parts of the pipelines (`-l | --log`) and to analyze memory usage (`-m | --memoryprofiler`). +4. `python pipeline.py --config /path/to/config.yaml -l -m` to run the scaLR. + +## Example configs + +### Config for cell type classification and biomarker identification + +NOTE: Below are just suggestions for the model parameters. Feel free to play around with them for tuning the model & improving the results. + +An example configuration file for the current dataset, incorporating the edits below, can be found at `scaLR/tutorials/pipeline/config_celltype.yaml`. Update the device as cuda or cpu as per the requirement. + +- **Device setup** + - Update device: 'cuda' for GPU enabled runtype, else device: 'cpu' for CPU enabled runtype. +- **Experiment Config** + - The default exp_run number is 0. If not changed, the celltype classification experiment would be exp_run_0 with all the pipeline results. +- **Data Config** + - Update the full_datapath to `data/modified_adata.h5ad` (as we will include GeneRecallCurve in the downstream). + - Specify the num_workers value for effective parallelization. + - Set target to cell_type. +- **Feature Selection** + - Specify the num_workers value for effective parallelization. + - Update the model layers to [5000, 10], as there are only 10 cell types in the dataset. + - Change epoch to 10. +- **Final Model Training** + - Update the model layers to the same as for feature selection: [5000, 10]. + - Change epoch to 100. +- **Analysis** + - Downstream Analysis + - Uncomment the test_samples_downstream_analysis section. + - Update the reference_genes_path to `scaLR/tutorials/pipeline/grc_reference_gene.csv`. + - Refer to the section below: + +```yaml +# Config file for pipeline run for cell type classification. + +# DEVICE SETUP. +device: 'cuda' + +# EXPERIMENT. +experiment: + dirpath: 'scalr_experiments' + exp_name: 'exp_name' + exp_run: 0 + +# DATA CONFIG. +data: + sample_chunksize: 20000 + + train_val_test: + full_datapath: 'data/modified_adata.h5ad' + num_workers: 2 + + splitter_config: + name: GroupSplitter + params: + split_ratio: [7, 1, 2.5] + stratify: 'donor_id' + + # split_datapaths: '' + + # preprocess: + # - name: SampleNorm + # params: + # **args + + # - name: StandardScaler + # params: + # **args + + target: cell_type + +# FEATURE SELECTION. +feature_selection: + + # score_matrix: '/path/to/matrix' + feature_subsetsize: 5000 + num_workers: 2 + + model: + name: SequentialModel + params: + layers: [5000, 10] + weights_init_zero: True + + model_train_config: + trainer: SimpleModelTrainer + + dataloader: + name: SimpleDataLoader + params: + batch_size: 25000 + padding: 5000 + + optimizer: + name: SGD + params: + lr: 1.0e-3 + weight_decay: 0.1 + + loss: + name: CrossEntropyLoss + + epochs: 10 + + scoring_config: + name: LinearScorer + + features_selector: + name: AbsMean + params: + k: 5000 + +# FINAL MODEL TRAINING. +final_training: + + model: + name: SequentialModel + params: + layers: [5000, 10] + dropout: 0 + weights_init_zero: False + + model_train_config: + resume_from_checkpoint: null + + trainer: SimpleModelTrainer + + dataloader: + name: SimpleDataLoader + params: + batch_size: 15000 + + optimizer: + name: Adam + params: + lr: 1.0e-3 + weight_decay: 0 + + loss: + name: CrossEntropyLoss + + epochs: 100 + + callbacks: + - name: TensorboardLogger + - name: EarlyStopping + params: + patience: 3 + min_delta: 1.0e-4 + - name: ModelCheckpoint + params: + interval: 5 +analysis: + + model_checkpoint: '' + + dataloader: + name: SimpleDataLoader + params: + batch_size: 15000 + + gene_analysis: + scoring_config: + name: LinearScorer + + features_selector: + name: ClasswisePromoters + params: + k: 100 + test_samples_downstream_analysis: + - name: GeneRecallCurve + params: + reference_genes_path: 'scaLR/tutorials/pipeline/grc_reference_gene.csv' + top_K: 300 + plots_per_row: 3 + features_selector: + name: ClasswiseAbs + params: {} + - name: Heatmap + params: {} + - name: RocAucCurve + params: {} +``` + +### Config for clinical condition-specific biomarker identification and DGE analysis + +An example configuration file (`scaLR/tutorials/pipeline/config_clinical.yaml`). Update the device as CUDA or CPU as per the requirement. + +- Experiment Config + - Make sure to change the exp_run number if you have an experiment with the same number earlier related to cell classification. As we have done one experiment earlier, we'll change the number now to '1'. +- Data Config + - The full_datapath remains the same as above. + - Change the target to disease (this column contains data for clinical conditions, COVID-19/normal). +- Feature Selection + - Update the model layers to [5000, 2], as there are only two types of clinical conditions. + - epoch as 10. +- Final Model Training + - Update the model layers to the same as for feature selection: [5000, 2]. + - epoch as 100. +- Analysis + - Downstream Analysis + - Uncomment the full_samples_downstream_analysis section for example config file. + - We are not performing the 'gene_recall_curve' analysis in this case. It can be performed if the COVID-19/normal specific genes are available, but there are many possibilities of genes in the case of normal conditions. + - There are two options to perform differential gene expression (DGE) analysis: **DgePseudoBulk and DgeLMEM**. The parameters are updated as follows. Note that DgeLMEM may take a bit more time, as the multiprocessing is not very efficient with only 2 CPUs in the current Colab runtime. + - Refer to the section below: + +```yaml +analysis: + + model_checkpoint: '' + + dataloader: + name: SimpleDataLoader + params: + batch_size: 15000 + + gene_analysis: + scoring_config: + name: LinearScorer + + features_selector: + name: ClasswisePromoters + params: + k: 100 + full_samples_downstream_analysis: + - name: Heatmap + params: + top_n_genes: 100 + - name: RocAucCurve + params: {} + - name: DgePseudoBulk + params: + celltype_column: 'cell_type' + design_factor: 'disease' + factor_categories: ['COVID-19', 'normal'] + sum_column: 'donor_id' + cell_subsets: ['conventional dendritic cell', 'natural killer cell'] + - name: DgeLMEM + params: + fixed_effect_column: 'disease' + fixed_effect_factors: ['COVID-19', 'normal'] + group: 'donor_id' + celltype_column: 'cell_type' + cell_subsets: ['conventional dendritic cell'] + gene_batch_size: 1000 + coef_threshold: 0.1 +``` + +## Experiment output structure + +- **pipeline.py**: The main script that performs an end-to-end run. + - `exp_dir`: root experiment directory for the storage of all step outputs of the platform specified in the config. + - `config.yml`: copy of config file to reproduce the user-defined experiment. + +- **data_ingestion**: Reads the data and splits it into Train/Validation/Test sets for the pipeline. Then, it performs sample-wise normalization on the data. + - `exp_dir` + - `data` + - `train_val_test_split.json`: contains sample indices for train/validation/test splits. + - `label_mappings.json`: contains mappings of all metadata columns between labels and IDs. + - `train_val_test_split`: directory containing the train, validation, and test samples and data files. + +- **feature_extraction**: Performs feature selection and extraction of new datasets containing a subset of features. + - `exp_dir` + - `feature_extraction` + - `chunked_models`: contains weights of each model trained on feature subset data (refer to feature subsetting algorithm). + - `feature_subset_data`: directory containing the new feature-subsetted train, val, and test samples anndatas. + - `score_matrix.csv`: combined scores of all individual models for each feature and class. shape: n_classes X n_features. + - `top_features.json`: a file containing a list of top features selected / to be subsetted from total features. + +- **final_model_training**: Trains a final model based on `train_datapath` and `val_datapath` in config. + - `exp_dir` + - `model` + - `logs`: directory containing Tensorboard Logs for the training of the model. + - `checkpoints`: directory containing model weights checkpointed at every interval specified in config. + - `best_model`: the best model checkpoint contains information to use model for inference/resume training. + - `model_config.yaml`: config file containing model parameters. + - `mappings.json`: contains mapping of class_names to class_ids used by model during training. + - `model.pt`: contains model weights. + +- **eval_and_analysis**: Performs evaluation of best model trained on user-defined metrics on the test set. Also performs various downstream tasks. + - `exp_dir` + - `analysis` + - `classification_report.csv`: contains classification report showing Precision, Recall, F1, and accuracy metrics for each class on the test set. + - `gene_analysis` + - `score_matrix.csv`: score of the final model, for each feature and class. shape: n_classes X n_features. + - `top_features.json`: a file containing a list of selected top features/biomarkers. + - `test_samples/full_samples` + - `heatmaps` + - `class_name.svg`: heatmap for top genes of a particular class w.r.t those genes association in other classes. E.g., B.svg, C.svg, etc. + - `roc_auc.svg`: contains ROC-AUC plot for all classes. + - `gene_recall_curve.svg`: contains gene recall curve plots. + - `gene_recall_curve_info.json`: contains reference genes list which are present in top_K ranked genes per class for each model. + - `pseudobulk_dge_result` + - `pbkDGE_celltype_factor_categories_0_vs_factor_categories_1.csv`: contains Pseudobulk DGE results between selected factor categories for a celltype. + - `pbkDGE_celltype_factor_categories_0_vs_factor_categories_1.svg`: volcano plot of Log2Foldchange vs -log10(p-value) of genes. + - `lmem_dge_result` + - `lmemDGE_celltype.csv`: contains LMEM DGE results between selected factor categories for a celltype. + - `lmemDGE_fixed_effect_factor_X.svg`: volcano plot of coefficient vs -log10(p-value) of genes. diff --git a/docs/guide_simple_api.md b/docs/guide_simple_api.md new file mode 100644 index 0000000..a157ad6 --- /dev/null +++ b/docs/guide_simple_api.md @@ -0,0 +1,134 @@ +# Simple API guide + +This guide covers the config-free `scalr.train`/`scalr.annotate`/`scalr.load_model` API in depth. For the 3-line hook, see the [Quickstart](https://github.com/infocusp/scaLR/blob/main/README.md#quickstart-annotate-cells-in-3-lines) in the main README; for the configuration-driven research pipeline, see the [Pipeline guide](guide_pipeline.md). + +## Validate your data before training or annotating + +```python +report = scalr.validate(adata, labels_key="cell_type", group_key="donor_id") +print(report) +``` + +`scalr.validate` (see [scalr/validation.py](https://github.com/infocusp/scaLR/blob/main/scalr/validation.py)) checks AnnData structure (empty/duplicate cells or genes), numerical issues (NaN/Inf, zero-count or unusually large-library-size cells), label/group columns, and reports the detected normalization state (raw counts, log1p-normalized, or scaled) — surfacing problems as an explicit report instead of a downstream stack trace. + +## Gene alignment + +Query data does not need to share the model's gene set or gene order. `model.predict`/`scalr.annotate` call [scalr/genes.py](https://github.com/infocusp/scaLR/blob/main/scalr/genes.py) internally to reorder/subset the query to the model's expected feature list, zero-filling any missing genes and reporting coverage (`result.metadata["gene_coverage"]`, `result.metadata["missing_features"]`). You can also run this alignment directly: + +```python +aligned_adata, report = scalr.align_genes(adata, reference_features, min_feature_overlap=0.1) +print(report) +``` + +## Streaming / bounded-memory prediction + +```python +result = model.predict("large_query.h5ad", streaming=True, chunk_size=20000) +# or, via scalr.annotate: +result = scalr.annotate(adata, model="models/pbmc_v1", streaming=True, chunk_size=20000) +``` + +Passing a file path to `model.predict`/`scalr.annotate` streams automatically; `streaming=True` on an in-memory `AnnData` also works. Gene alignment and scoring happen `chunk_size` cells at a time instead of materializing one aligned matrix for the whole dataset, bounding peak alignment/scoring memory regardless of total cell count. For out-of-core reads too (never loading the full raw dataset into memory), point it at a directory of chunked `.h5ad` files rather than a single large `.h5ad` file — see [scalr/artifact.py](https://github.com/infocusp/scaLR/blob/main/scalr/artifact.py). + +## Self-contained model artifacts + +`model.save("models/pbmc_v1")` writes a versioned, portable directory containing everything needed to reload and run the model — no external config required: `manifest.json`, `model.pt`, `model_config.json`, `label_mapping.json`, `features.json`, `preprocessing.json`, `calibration.json`, and `metrics.json`. `scalr.load_model("models/pbmc_v1")` reads it back into an `AnnotationModel` — see [scalr/artifact.py](https://github.com/infocusp/scaLR/blob/main/scalr/artifact.py). + +## Leakage-safe evaluation & class imbalance + +When `group_key` (e.g. `donor_id`) is passed to `scalr.train`, splitting keeps every group confined to a single split and checks for leakage across splits, warning if any group value still appears in more than one ([scalr/leakage.py](https://github.com/infocusp/scaLR/blob/main/scalr/leakage.py)). Training also reports class-size imbalance and uses inverse-frequency class weighting by default, and evaluates on the held-out test split with macro-F1, weighted-F1, balanced accuracy and a per-class report rather than relying on overall accuracy alone ([scalr/metrics.py](https://github.com/infocusp/scaLR/blob/main/scalr/metrics.py)). + +## Hierarchical (coarse-to-fine) annotation + +Pass a `taxonomy` mapping every fine-grained label to a broad label at training time, then request either level at inference time — without retraining: + +```python +taxonomy = {"CD4_T": "T_cell", "CD8_T": "T_cell", "B_cell": "B_cell", "DC": "Myeloid"} +model = scalr.train(adata, labels_key="cell_type", group_key="donor_id", taxonomy=taxonomy) + +fine_result = model.predict(adata, level="fine") # default +broad_result = model.predict(adata, level="broad") # aggregated via the taxonomy +``` + +See [scalr/hierarchy.py](https://github.com/infocusp/scaLR/blob/main/scalr/hierarchy.py); broad-class probabilities are the sum of the underlying fine-class probabilities. + +## Cluster-aware refinement & doublet/mixed-cell screening + +```python +result = scalr.annotate( + adata, + model="models/pbmc_v1", + cluster_refinement="auto", # or an adata.obs column name + flag_doublets=True, +) +``` + +`cluster_refinement` relabels cells to their cluster's confidence-weighted majority class ([scalr/refinement.py](https://github.com/infocusp/scaLR/blob/main/scalr/refinement.py)) — the raw, unrefined labels are always kept in `result.metadata["raw_labels"]`, never silently discarded. `flag_doublets` sets `result.is_possible_doublet` for cells whose top-two class probabilities are both substantial and close together ([scalr/doublet.py](https://github.com/infocusp/scaLR/blob/main/scalr/doublet.py)) — a heuristic screening signal, not a validated doublet call. + +## Explainability + +```python +explanations = model.explain(adata, indices=[10, 20, 30], top_k=20) +print(explanations[0]) + +supporting, contradictory = model.explain_class("T_cell", top_k=50) +``` + +`model.explain` reports, per requested cell, the genes whose expression most supported or contradicted the model's prediction, via a Grad x Input gene-attribution score. `model.explain_class` does the same for a class in general, independent of any specific cell. This is a first-order sensitivity explanation, not a validated biological one — treat it as a screening/interpretation aid, like the doublet/refinement heuristics above, not ground truth ([scalr/explain.py](https://github.com/infocusp/scaLR/blob/main/scalr/explain.py)). + +## Dataset drift detection + +`is_unknown`/`is_possible_doublet` flag individual out-of-distribution cells, but can miss a whole cohort that has drifted (a different tissue, protocol, or batch) if most cells still score confidently. `scalr.train` records per-gene and per-cell (library size, detected-gene-count) statistics from the training split as a `drift_reference`; `scalr.detect_drift` compares a new dataset against it before you trust its predictions: + +```python +report = scalr.detect_drift(query_adata, model="models/pbmc_v1") +print(report) +if report.is_drifted: + ... # treat predictions on this dataset with caution +``` + +A gene counts as drifted when its mean-expression shift exceeds `z_threshold` reference standard deviations (default 2); the dataset is flagged `is_drifted` when either too large a fraction of genes drift (`drift_fraction_threshold`, default 10%) or the library-size/detected-gene-count shift is extreme on its own (`severity_z_threshold`, default 5) — the latter catches e.g. a sequencing-depth change a per-gene test can miss. This is a univariate summary-statistics check, not a full distributional test — see [scalr/drift.py](https://github.com/infocusp/scaLR/blob/main/scalr/drift.py). Also available as `scalr drift --input query.h5ad --model models/pbmc_v1 [--fail-on-drift]` from the CLI, for use as a pre-flight/CI gate. + +## Feature-selection stability + +```python +report = scalr.compute_feature_stability(adata, labels_key="cell_type", top_k=50, n_runs=20) +print(report) +print(report.top_stable_features(min_runs=15)) +``` + +Repeats feature selection over random subsamples and reports each gene's selection frequency plus the mean pairwise Jaccard similarity of the selected sets, to help distinguish a stable biomarker signal from a split-specific artifact ([scalr/feature_stability.py](https://github.com/infocusp/scaLR/blob/main/scalr/feature_stability.py)). + +## Model cards, provenance & a model registry/hub + +Every `model.save(...)` call also writes a `README.md` model card (training data, held-out metrics, calibration, known limitations, license) into the artifact directory ([scalr/model_card.py](https://github.com/infocusp/scaLR/blob/main/scalr/model_card.py)). Models can be registered and resolved by name via a local registry ([scalr/models.py](https://github.com/infocusp/scaLR/blob/main/scalr/models.py)): + +```python +scalr.models.register("models/pbmc_v1", "human_pbmc") +scalr.models.list() # -> ['human_pbmc'] +scalr.models.info("human_pbmc") + +model = scalr.load_model("human_pbmc") # resolves the registered name +``` + +Models can also be published to, and pulled from, the [Hugging Face Hub](https://huggingface.co/models) — requires the optional `hub` extra (`pip install "pyscaLR[hub]"`): + +```python +scalr.models.push("models/pbmc_v1", "your-org/human-pbmc") # publish + +model = scalr.load_model("your-org/human-pbmc") # auto-downloads on first use, then cached locally +``` + +`scalr.load_model` tries an on-disk path, then the local registry, and only falls back to a Hub download for `"/"`-shaped names it can't resolve locally; once downloaded, a model is registered under that same name so later loads are local and offline. Use `scalr.models.download(repo_id, revision=..., token=...)` directly for a specific revision or a private repo. + +## Benchmark suite + +```python +df = scalr.run_benchmark( + {"small": small_adata, "medium": medium_adata}, + labels_key="cell_type", + group_key="donor_id", +) +``` + +Reports scientific quality (macro-F1, balanced accuracy) alongside computational cost (wall time, cells/sec, peak RSS) per dataset in one table, so quality and resource cost are always reported together rather than a one-off number ([scalr/benchmark.py](https://github.com/infocusp/scaLR/blob/main/scalr/benchmark.py)). diff --git a/docs/index.rst b/docs/index.rst index 7d3eaca..bb0ec10 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -14,6 +14,13 @@ scaLR Documentation :parser: myst_parser.sphinx_ :start-line: 26 +.. toctree:: + :maxdepth: 10 + :caption: scaLR Simple API + :hidden: + + scalr.simple_api + .. toctree:: :maxdepth: 10 :caption: scaLR @@ -25,13 +32,19 @@ scaLR Documentation scalr.nn scalr.utils +.. toctree:: + :maxdepth: 10 + :caption: Guides + :hidden: + + guide_simple_api + guide_cli + guide_pipeline + tutorials + .. toctree:: :maxdepth: 10 :caption: scaLR Pipelines :hidden: - + pipeline - scalr.data_ingestion_pipeline - scalr.feature_extraction_pipeline - scalr.model_training_pipeline - scalr.eval_and_analysis_pipeline diff --git a/docs/modules.rst b/docs/modules.rst deleted file mode 100644 index 878499e..0000000 --- a/docs/modules.rst +++ /dev/null @@ -1,7 +0,0 @@ -scaLR -===== - -.. toctree:: - :maxdepth: 4 - - scalr diff --git a/docs/scalr.rst b/docs/scalr.rst deleted file mode 100644 index d569948..0000000 --- a/docs/scalr.rst +++ /dev/null @@ -1,11 +0,0 @@ -scalr -===== - -.. toctree:: - :maxdepth: 10 - - scalr.analysis - scalr.data - scalr.feature - scalr.nn - scalr.utils diff --git a/docs/scalr.simple_api.rst b/docs/scalr.simple_api.rst new file mode 100644 index 0000000..b1e97ed --- /dev/null +++ b/docs/scalr.simple_api.rst @@ -0,0 +1,165 @@ +Simple API +========== + +api +--- + +.. automodule:: scalr.api + :members: + :undoc-members: + :show-inheritance: + :private-members: + +artifact +-------- + +.. automodule:: scalr.artifact + :members: + :undoc-members: + :show-inheritance: + :private-members: + +benchmark +--------- + +.. automodule:: scalr.benchmark + :members: + :undoc-members: + :show-inheritance: + :private-members: + +calibration +----------- + +.. automodule:: scalr.calibration + :members: + :undoc-members: + :show-inheritance: + :private-members: + +cli +--- + +.. automodule:: scalr.cli + :members: + :undoc-members: + :show-inheritance: + :private-members: + +doublet +------- + +.. automodule:: scalr.doublet + :members: + :undoc-members: + :show-inheritance: + :private-members: + +drift +----- + +.. automodule:: scalr.drift + :members: + :undoc-members: + :show-inheritance: + :private-members: + +explain +------- + +.. automodule:: scalr.explain + :members: + :undoc-members: + :show-inheritance: + :private-members: + +feature\_stability +------------------- + +.. automodule:: scalr.feature_stability + :members: + :undoc-members: + :show-inheritance: + :private-members: + +genes +----- + +.. automodule:: scalr.genes + :members: + :undoc-members: + :show-inheritance: + :private-members: + +hierarchy +--------- + +.. automodule:: scalr.hierarchy + :members: + :undoc-members: + :show-inheritance: + :private-members: + +leakage +------- + +.. automodule:: scalr.leakage + :members: + :undoc-members: + :show-inheritance: + :private-members: + +metrics +------- + +.. automodule:: scalr.metrics + :members: + :undoc-members: + :show-inheritance: + :private-members: + +model\_card +----------- + +.. automodule:: scalr.model_card + :members: + :undoc-members: + :show-inheritance: + :private-members: + +models +------ + +.. automodule:: scalr.models + :members: + :undoc-members: + :show-inheritance: + :private-members: + +refinement +---------- + +.. automodule:: scalr.refinement + :members: + :undoc-members: + :show-inheritance: + :private-members: + +result +------ + +.. automodule:: scalr.result + :members: + :undoc-members: + :show-inheritance: + :private-members: + :exclude-members: obs_names, labels, probabilities, class_names, confidence, entropy, margin, is_unknown, top_k, ood_score, explanations, metadata, is_possible_doublet + +validation +---------- + +.. automodule:: scalr.validation + :members: + :undoc-members: + :show-inheritance: + :private-members: diff --git a/docs/tutorials.md b/docs/tutorials.md new file mode 100644 index 0000000..1b9a68c --- /dev/null +++ b/docs/tutorials.md @@ -0,0 +1,12 @@ +# Tutorials + +Detailed tutorials have been made on how to use some pipeline functionalities as a scaLR library. Find the links below. + +- **scaLR pipeline** [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infocusp/scaLR/blob/main/tutorials/pipeline/scalr_pipeline.ipynb) +- **Differential gene expression analysis** [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infocusp/scaLR/blob/main/tutorials/analysis/differential_gene_expression/dge.ipynb) +- **Gene recall curve** [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infocusp/scaLR/blob/main/tutorials/analysis/gene_recall_curve/gene_recall_curve.ipynb) +- **SHAP analysis** [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infocusp/scaLR/blob/main/tutorials/analysis/shap_analysis/shap_heatmap.ipynb) +- **Normalization** [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infocusp/scaLR/blob/main/tutorials/preprocessing/normalization.ipynb) +- **Batch correction** [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/infocusp/scaLR/blob/main/tutorials/preprocessing/batch_correction.ipynb) + +- **An example of jupyter notebook to [run scaLR in local machine](https://github.com/infocusp/scaLR/blob/main/tutorials/pipeline/scalr_pipeline_local_run.ipynb)**. diff --git a/models/pbmc_v1/README.md b/models/pbmc_v1/README.md new file mode 100644 index 0000000..9e7fd72 --- /dev/null +++ b/models/pbmc_v1/README.md @@ -0,0 +1,37 @@ +# Model Card: cell_type + +## Summary +- **Task**: cell_type_annotation +- **Model version**: 1.0.0 +- **scaLR version**: 2.0.0.dev0 +- **Species**: not specified +- **Random seed**: 42 + +## Training data +- **Labels column**: cell_type +- **Grouping column (leakage-safe split)**: donor_id +- **Split ratio (train/val/test)**: [0.7, 0.1, 0.2] +- **Feature count**: 30695 +- **Classes (10)**: CD16-negative, CD56-bright natural killer cell, human, CD16-positive, CD56-dim natural killer cell, human, classical monocyte, conventional dendritic cell, granulocyte, intermediate monocyte, natural killer cell, non-classical monocyte, plasmacytoid dendritic cell, platelet + +## Performance (held-out test split) +- macro-F1: 0.6885 +- weighted-F1: 0.9281 +- balanced accuracy: 0.6822 +- test cells: 23,492 + +## Confidence & abstention +- Calibration: temperature scaling (T=2.546) +- Open-set thresholds: {'min_confidence': 0.8577296197414398, 'max_entropy': 0.20046073645353296, 'min_margin': 0.7288531720638275} + +## Known limitations +- Evaluated only on the training-distribution data described above; performance on + out-of-distribution donors, tissues, or sequencing technologies is not validated. +- Open-set (`is_unknown`) thresholds are fit on a single validation split and are a + heuristic screening signal, not a guarantee of correct abstention. +- Gene/feature coverage should be checked (`scalr.validate`) on every new query dataset. + +## License +- Source code license: see repository LICENSE. +- This model artifact's weights/data license: not specified by this template; + set explicitly by the model's author/distributor before sharing. diff --git a/pyproject.toml b/pyproject.toml index b3dd290..80a6fbf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,20 +5,32 @@ build-backend = "hatchling.build" [project] name = "pyscaLR" -version = "1.1.0" -requires-python = ">=3.10.20" +version = "1.2.0" +requires-python = ">=3.11,<3.15" authors = [ { name="Infocusp", email="saurabh@infocusp.com" }, ] description = "scaLR: Single cell analysis using low resource." readme = "README.md" classifiers = [ - "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", "Operating System :: OS Independent", - "Intended Audience :: Science/Research" + "Intended Audience :: Science/Research", + "License :: OSI Approved :: GNU General Public License v3 (GPLv3)" ] dynamic = ["dependencies"] + +[project.optional-dependencies] +dev = ["pytest>=8.3,<9", "pre_commit>=4,<5", "isort>=5.13,<7", "yapf>=0.40,<1"] +hub = ["huggingface_hub>=0.24,<1"] + +[project.scripts] +scalr = "scalr.cli:main" + [tool.hatch.build.targets.wheel] packages = ["scalr"] @@ -30,6 +42,8 @@ exclude = [ "docs/*", "tutorials/*", "tests/*", + "scalr/**/test_*.py", + "scaLR_v2_implementation_blueprint.md", ] license = {file = "LICENSE"} @@ -38,3 +52,5 @@ license = {file = "LICENSE"} Repository = "https://github.com/infocusp/scaLR.git" Homepage = "https://github.com/infocusp/scaLR" Issues = "https://github.com/infocusp/scaLR/issues" +Documentation = "https://infocusp.github.io/scaLR/" +Changelog = "https://github.com/infocusp/scaLR/blob/main/CHANGELOG.md" diff --git a/requirements.txt b/requirements.txt index c3a2e56..172e5ee 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,18 +1,16 @@ -anndata>=0.11.2,<0.12 -isort==5.13.2 -loky==3.4.1 -memory-profiler==0.61.0 -pillow==10.4.0 -pre_commit==4.0.1 -pydeseq2==0.4.11 -pyparsing==3.2.0 -pytest==8.3.3 -PyYAML==6.0.2 -scanpy==1.10.3 -scikit-learn==1.5.2 -shap==0.46.0 -tensorboard==2.17.0 -toml==0.10.2 -torch==2.4.1 --index-url https://download.pytorch.org/whl/cu118 -tqdm==4.66.5 -yapf==0.40.2 \ No newline at end of file +anndata>=0.12,<0.13 +loky>=3.4,<4 +memory-profiler>=0.61,<1 +numpy>=1.26,<3 +pandas>=2.2,<4 +pillow>=10.4,<12 +pydeseq2>=0.4.11,<1 +pyparsing>=3.2,<4 +PyYAML>=6,<7 +scanpy>=1.10.3,<2 +scikit-learn>=1.5,<2 +shap>=0.46,<1 +tensorboard>=2.17,<3 +toml>=0.10,<1 +torch>=2.4,<3 +tqdm>=4.66,<5 \ No newline at end of file diff --git a/scaLR_v2_implementation_blueprint.md b/scaLR_v2_implementation_blueprint.md new file mode 100644 index 0000000..f0253b9 --- /dev/null +++ b/scaLR_v2_implementation_blueprint.md @@ -0,0 +1,2036 @@ +# scaLR 2.0 — Comprehensive Implementation Blueprint + +**Date:** 2026-09-21 +**Repository:** https://github.com/infocusp/scaLR +**Paper:** https://doi.org/10.1093/bib/bbaf243 +**Scope:** API redesign, scientific reliability, scalability, reproducibility, packaging, model distribution, visualization, benchmarking, and ecosystem integration. + +--- + +## 0. Executive summary + +scaLR's strongest technical identity is **low-resource, large-scale single-cell analysis using DNN-based feature selection, classification, and biomarker analysis**. The published work focuses on reducing computational requirements while maintaining competitive classification performance at large cell counts. citeturn794974search2 + +The current repository is an end-to-end, configuration-driven pipeline. It expects `.h5ad` input, documents normalized `adata.X`, uses a YAML configuration for pipeline execution, and currently exposes the main pipeline stages rather than a concise inference-first API. citeturn799194search2 + +### Recommended v2 positioning + +Do **not** turn scaLR into another generic deep-learning framework. Instead evolve it from: + +> **pipeline → library + pipeline** + +and position v2 around four properties: + +1. **Easy:** `scalr.annotate(adata, model="human_pbmc")` +2. **Reliable:** validation, calibration, abstention, leakage-safe evaluation, rare-cell metrics, open-set detection. +3. **Scalable:** sparse processing, bounded-memory inference, chunking/streaming, CPU/GPU auto-selection. +4. **Interpretable/reproducible:** feature stability, explanations, model cards, self-contained artifacts, run metadata. + +The result should be useful both to a beginner who wants cell-type labels and to an expert running million- or multi-million-cell analyses. + +--- + +# 1. Current-state assessment + +The public repository currently contains `scalr/`, `pipeline.py`, configuration, documentation and tutorials. The README describes `.h5ad` as the current input format and expects normalized expression values in `adata.X`; the example configuration uses a `GroupSplitter` stratified by donor. citeturn799194search2 + +The repository also has open enhancement items around configuration validation, unified logging, abstract base classes, more tests/CI, and Conda installation. These should be folded into the v2 foundation rather than treated as isolated cleanup tasks. + +### Current architectural strengths + +- Existing separation between ingestion, feature extraction, model training, and evaluation/analysis. +- Explicit chunk-size/resource configuration. +- Group-based data splitting is already demonstrated in the example configuration. citeturn799194search2 +- Feature-selection and biomarker-analysis capabilities align with the paper's contribution. citeturn794974search2 +- Existing CLI/pipeline can remain as the advanced/research workflow. + +### Current barriers to adoption + +- Pipeline-first workflow rather than inference-first library API. +- Users need to understand YAML/configuration before getting a simple prediction. +- `.h5ad` is the primary documented input path and preprocessing expectations are strict. citeturn799194search2 +- Installation is more constrained than ideal for CPU/GPU/HPC environments. +- No clear first-class model artifact/model-hub contract. +- No first-class confidence/abstention/open-set workflow. +- No explicit contract for gene alignment and preprocessing transfer at inference time. +- Limited guardrails against leakage, rare-cell failure, mixed/doublet cells, and distribution shift. +- Benchmarking can be strengthened to cover calibration, unknown-cell detection, scaling curves, and reproducibility. + +--- + +# 2. Product architecture for scaLR 2.0 + +Use two user layers over the same internal engine. + +```text + scaLR 2.0 + | + +----------------+----------------+ + | | + Simple API Advanced API + | | + annotate() YAML pipeline + train() Pipeline classes + load_model() Custom components + explain() Research configs + | | + +----------------+----------------+ + | + Shared core engine + | + +------------------+------------------+ + | | | + Data/Preprocess Model/Inference Evaluation + | | | + AnnData/genes DNN/selection metrics + normalization uncertainty benchmarks + validation explanation reporting +``` + +### Design principle + +Do not rewrite the current system in one large refactor. Introduce the v2 public API as a thin, stable layer over the current components, then refactor internals incrementally. + +--- + +# 3. Public API + +## 3.1 One-line annotation + +```python +import scalr + +result = scalr.annotate( + adata, + model="human_pbmc", + device="auto", + preprocess="auto", +) + +adata.obs["scalr_cell_type"] = result.labels +adata.obs["scalr_confidence"] = result.confidence +adata.obs["scalr_unknown"] = result.is_unknown +``` + +## 3.2 Training + +```python +model = scalr.train( + adata, + labels_key="cell_type", + group_key="donor_id", + device="auto", +) + +model.save("models/pbmc_v1") +``` + +## 3.3 Loading and inference + +```python +model = scalr.load_model("models/pbmc_v1") +result = model.predict(adata) +``` + +## 3.4 Explanation + +```python +explanation = model.explain( + adata, + indices=[10, 20, 30], + top_k=20, +) +``` + +## 3.5 Validation + +```python +report = scalr.validate( + adata, + model="human_pbmc", +) +``` + +--- + +# 4. Prediction result contract + +Return a typed result object rather than a bare array. + +```python +result.labels +result.probabilities +result.top_k +result.confidence +result.entropy +result.is_unknown +result.ood_score +result.margin +result.explanations +result.metadata +``` + +Recommended conceptual structure: + +```text +PredictionResult +├── labels +├── probabilities +├── confidence +├── entropy +├── top_k +├── margin +├── is_unknown +├── ood_score +├── explanations +└── metadata +``` + +The result must preserve row/cell identity so predictions cannot silently become misaligned with `adata.obs_names`. + +--- + +# 5. AnnData-first integration + +AnnData should become the first-class Python interface. + +### Standard output fields + +```python +adata.obs["scalr_pred"] +adata.obs["scalr_confidence"] +adata.obs["scalr_unknown"] +``` + +Optional: + +```python +adata.obsm["scalr_probabilities"] +adata.obsm["scalr_embeddings"] +adata.uns["scalr"] +``` + +`adata.uns["scalr"]` should contain enough metadata to reconstruct the model/preprocessing contract used for that annotation. + +### Preserve compatibility + +Keep the existing `.h5ad` workflow. Add support rather than breaking it. + +Later adapters can support: + +- Zarr +- Seurat via conversion +- SingleCellExperiment/Bioconductor via a documented bridge +- streaming readers + +--- + +# 6. Input validation subsystem + +Create a formal validation layer instead of letting downstream PyTorch, NumPy or scikit-learn exceptions reveal implementation details. + +## 6.1 Validation checks + +### Data structure + +- AnnData object or supported path. +- Non-empty observations and variables. +- Unique cell IDs. +- Unique gene IDs. +- Valid matrix dimensions. +- Supported sparse formats. + +### Numerical checks + +- NaN/Inf detection. +- Negative values handled explicitly. +- Zero-count cells reported. +- Extremely large library sizes reported. +- Optional mitochondrial/QC metadata checks. + +### Metadata checks + +During training: + +- label column exists. +- grouping column exists when required. +- missing labels are reported. +- class sizes are reported. + +### Model compatibility + +- species. +- expected gene identifier type. +- expected preprocessing. +- required feature coverage. +- label compatibility. + +## 6.2 User-facing report + +```text +scaLR validation +---------------- +ERROR + None + +WARNING + 179/5,000 model genes are missing from query. + +INFO + Input: 125,431 cells × 18,642 genes + Sparse: yes + Detected expression state: log1p-normalized + Model: human_pbmc_v1 + Gene coverage: 96.4% + +Status: usable with warning +``` + +## 6.3 No silent transformations + +`preprocess="auto"` should detect and explain the state. + +`preprocess="force"` may apply an explicit documented transformation. + +Ambiguous inputs should trigger an informative warning/error rather than guessing. + +--- + +# 7. Preprocessing contract + +Separate **fit-time preprocessing** from **inference-time application**. + +At training time, store: + +- normalization method. +- layer used. +- scaling parameters where applicable. +- selected genes and ordering. +- preprocessing version. +- optional QC thresholds. + +At inference time, apply exactly the stored contract. + +### Important rule + +Never calculate training-dependent statistics from the query set when the model was trained on a reference dataset. + +This prevents data leakage and inconsistent predictions. + +--- + +# 8. Gene identifier normalization and alignment + +Gene alignment should be an explicit subsystem. + +Support, where appropriate: + +- gene symbols. +- Ensembl IDs. +- species-specific normalization. +- case normalization. +- documented mapping tables. + +Example: + +```python +aligned = scalr.align_genes( + query, + reference_model, + species="human", +) +``` + +Output: + +```text +Reference features: 5,000 +Query genes: 18,642 +Matched features: 4,821 (96.4%) +Missing features: 179 +Extra query genes: 13,821 + +Status: usable with warning +``` + +Add configurable `min_feature_overlap` and fail early when coverage is inadequate. + +--- + +# 9. Reference / query model + +Make the distinction explicit. + +## Training mode + +```text +Labeled reference + | + v +Feature selection + | + v +DNN training + | + v +Calibration / evaluation + | + v +Versioned model artifact +``` + +## Annotation mode + +```text +Versioned reference model + + + Query h5ad + | + v +Gene alignment + | + v +Preprocessing contract + | + v +Batch/chunk inference + | + v +Labels + confidence + uncertainty +``` + +The second path should not retrain the model unless explicitly requested. + +--- + +# 10. Confidence calibration and uncertainty + +A softmax maximum should not automatically be described as a biologically calibrated confidence value. + +Implement: + +- temperature scaling first. +- calibration dataset/holdout handling. +- Expected Calibration Error. +- Brier score. +- reliability diagrams. + +Return: + +- top-1 label. +- calibrated confidence. +- top-k probabilities. +- predictive entropy. +- top-1/top-2 margin. +- optional OOD score. + +Example: + +```text +Prediction: T cell +Calibrated confidence: 0.78 +Entropy: 0.82 +Top-2 margin: 0.11 +Status: UNCERTAIN +``` + +--- + +# 11. Open-set / unknown-cell detection + +Closed-set classification forces every cell into a known class. scaLR 2.0 should support abstention. + +```python +result = model.predict( + adata, + open_set=True, +) +``` + +Possible output: + +```text +label: UNKNOWN +confidence: 0.38 +ood_score: high +``` + +Initial implementation: + +- configurable confidence threshold. +- entropy threshold. +- top-1/top-2 margin. + +Later evaluation/experimentation: + +- energy score. +- Mahalanobis distance. +- ensemble disagreement. + +Benchmark against held-out novel cell types. + +--- + +# 12. Rare-cell and class-imbalance support + +Do not rely on overall accuracy. + +Add options for: + +- weighted cross entropy. +- class-balanced sampling. +- focal loss as an optional method. +- minimum class-size warnings. + +Report: + +- macro-F1. +- weighted-F1. +- balanced accuracy. +- per-class precision/recall/F1. +- rare-class recall. +- confusion matrix. + +Example warning: + +```text +Class imbalance detected + +T cell: 500,000 +B cell: 200,000 +DC: 800 + +Recommendation: enable class-balanced sampling. +``` + +--- + +# 13. Leakage-safe evaluation + +The current examples already demonstrate donor-grouped splitting; make this a first-class safety feature rather than only a configuration option. citeturn799194search2 + +Support: + +```python +scalr.train( + adata, + labels_key="cell_type", + group_key="donor_id", + split_strategy="group", +) +``` + +Detect likely leakage when the same donor/patient/sample/batch appears across train/test. + +Warnings should identify the suspected grouping variable and recommend grouped splitting. + +### Default philosophy + +For reference datasets, prefer group-safe splits when donor/patient/sample identifiers are available. + +--- + +# 14. Hierarchical cell-type annotation + +Add optional coarse-to-fine classification. + +```text +Immune +├── Lymphoid +│ ├── T cell +│ │ ├── CD4 +│ │ └── CD8 +│ ├── B cell +│ └── NK +└── Myeloid + ├── Monocyte + └── DC +``` + +API: + +```python +result = model.predict( + adata, + level="broad", +) + +result = model.predict( + adata, + level="subtype", +) +``` + +This allows scaLR to represent the difference between evidence for a broad type and evidence for a specific subtype. + +--- + +# 15. Cluster-aware refinement + +Provide optional post-processing using neighborhood/cluster information. + +```python +result = model.predict( + adata, + cluster_refinement="auto", +) +``` + +Always preserve both: + +```text +raw_prediction +refined_prediction +``` + +Never overwrite the raw model output without explicit user request. + +--- + +# 16. Doublet / mixed-cell handling + +A classifier should not silently force an obviously mixed cell into a single lineage. + +Add an optional mixed-cell signal using: + +- high predictive entropy. +- conflicting marker programs. +- top-two class margin. +- optional integration with a dedicated doublet detector. + +Possible output: + +```text +Prediction: T cell +Status: POSSIBLE_DOUBLET +T cell: 0.48 +B cell: 0.43 +``` + +The doublet detection feature should be labeled as a screening signal, not a definitive biological diagnosis, unless validated against a specific detector/dataset. + +--- + +# 17. Distribution-shift robustness + +Evaluate reference/query differences in: + +- donor. +- tissue. +- disease state. +- sequencing technology. +- sequencing depth. +- batch. +- species, when supported. + +Expose a compatibility report: + +```text +Reference: human PBMC / 10x +Query: human PBMC / Smart-seq2 + +Gene coverage: 92.1% +Distribution shift: detected +Recommendation: validate on domain-matched reference before production use. +``` + +Do not claim universal cross-domain robustness without benchmark evidence. + +--- + +# 18. Feature-selection stability + +Feature selection is a central scientific feature of scaLR. The paper describes subset-based feature analysis and final feature selection as part of the workflow. citeturn794974search2 + +Add stability analysis: + +```text +Run 1 → top 500 +Run 2 → top 500 +Run 3 → top 500 +... +Run 100 → top 500 +``` + +Compute: + +- selection frequency. +- rank stability. +- Jaccard similarity. +- class-specific stability. + +Example: + +```text +Gene Selected runs / 100 +CD3D 98 +CD3E 97 +TRBC2 91 +IL7R 88 +``` + +This makes feature selection more defensible as a biomarker-discovery aid. + +--- + +# 19. Separate predictive features from biological biomarkers + +Do not equate model importance with biological causality or biomarker specificity. + +Expose separate reports: + +```text +1. Predictive model features +2. Class-specific features +3. Differentially expressed genes +4. Known/reference markers +5. Feature-selection stability +``` + +Naming should make this distinction explicit. + +--- + +# 20. Explainability API + +Turn the paper's feature-importance capability into an accessible user feature. citeturn794974search2 + +```python +explanation = model.explain( + adata, + indices=[10, 20, 30], + top_k=20, +) +``` + +For each cell: + +```text +cell_10392 +Prediction: T cell +Confidence: 0.94 + +Supporting genes +CD3D +0.42 +CD3E +0.37 +TRBC2 +0.29 +IL7R +0.21 + +Contradictory signals +MS4A1 -0.02 +``` + +Also provide: + +```python +model.explain_class("T_cell", top_k=50) +``` + +--- + +# 21. Marker sanity checks + +Optionally compare predictions to expected marker programs. + +Example: + +```text +Predicted: CD4 T cell + +Supporting +CD3D ✓ +CD3E ✓ +TRBC1 ✓ +IL7R ✓ + +Conflicting +MS4A1 ✗ +NKG7 low +``` + +This is an explanation/sanity check, not a substitute for biological validation. + +--- + +# 22. Batch correction safety + +If batch correction is supported, provide diagnostics instead of exposing it as a silent switch. + +Report before/after measures for: + +- batch mixing. +- cell-type conservation. +- neighborhood structure. +- cluster preservation. + +The goal is not merely to reduce batch separation; it is to reduce technical separation while preserving biological structure. + +--- + +# 23. Model artifact format + +Every model should be self-contained and versioned. + +Recommended format: + +```text +model/ +├── manifest.json +├── model.safetensors +├── model_config.json +├── label_mapping.json +├── features.json +├── preprocessing.json +├── calibration.json +├── metrics.json +└── README.md +``` + +### `manifest.json` + +```json +{ + "format_version": 1, + "scalr_version": "2.x", + "model_name": "human_pbmc", + "model_version": "1.0.0", + "task": "cell_type_annotation", + "species": "human", + "feature_count": 5000, + "labels": ["B_cell", "T_cell", "NK_cell"], + "normalization": "log1p_cp10k", + "calibration": "temperature_scaling", + "random_seed": 42 +} +``` + +Prefer a portable state/weights representation such as `safetensors` where compatible rather than relying on arbitrary Python object deserialization. + +--- + +# 24. Model provenance and model cards + +Every published/pretrained model should have a model card containing: + +- training/reference datasets. +- species. +- tissue. +- technology. +- number of cells. +- donor count. +- cell types. +- normalization assumptions. +- input gene coverage. +- performance metrics. +- unknown-cell behavior. +- known limitations. +- model version. +- code version. +- model/data license information. + +Separate: + +```text +source-code license +model license +training-data license +paper license +``` + +The current repository identifies the project as GPL-3.0, while the journal article has its own publication license terms; future model distribution should make those boundaries explicit. citeturn799194search2turn794974search2 + +--- + +# 25. Model registry / model hub + +Start with a local registry API: + +```python +scalr.models.list() +scalr.models.info("human_pbmc") +scalr.models.download("human_pbmc") +``` + +Possible model families: + +```text +Human +├── PBMC +├── Lung +├── Brain +├── Liver +└── Pan-tissue + +Mouse +├── Brain +├── Immune +└── Development + +Disease +├── Cancer +├── Autoimmune +└── Other validated cohorts +``` + +Do not populate models merely for completeness. Every model should have a reproducible provenance record and benchmark/model card. + +--- + +# 26. Automatic model selection + +Provide: + +```python +result = scalr.auto_annotate(adata) +``` + +The selector can inspect: + +- species metadata. +- tissue metadata where available. +- gene overlap. +- supported labels. +- model provenance. +- compatibility constraints. + +It should explain the selected model: + +```text +Selected: human_lung_v2 + +Reason: +✓ human gene identifiers +✓ high feature overlap +✓ lung-associated model +✓ compatible normalization +``` + +Do not make hidden model choices. The selected model and reason must be recorded. + +--- + +# 27. Device / resource auto-selection + +Defaults: + +```python +device="auto" +batch_size="auto" +num_workers="auto" +``` + +Inspect: + +- CUDA availability. +- GPU VRAM. +- system RAM. +- number of cells. +- number of selected genes. +- sparse/dense input. + +Log the chosen resources before execution. + +Allow full override for reproducibility: + +```python +device="cpu" +batch_size=4096 +num_workers=4 +``` + +--- + +# 28. Sparse and bounded-memory inference + +Preserve sparse matrices until conversion to dense tensors is unavoidable. + +Avoid unnecessary `.toarray()` / `.A` / dense materialization. + +Provide: + +```python +model.predict( + "large_dataset.h5ad", + streaming=True, + chunk_size=20000, +) +``` + +A bounded-memory test must prove that peak RAM does not grow linearly with total cell count for the intended streaming workflow. + +--- + +# 29. Out-of-core and Zarr/cloud roadmap + +Later support: + +- H5AD chunked reading. +- Zarr. +- object storage where supported. +- remote references. + +Example future API: + +```python +model.predict( + "s3://bucket/atlas.zarr", + streaming=True, +) +``` + +Do not advertise cloud execution before it is tested end-to-end. + +--- + +# 30. Lightweight and low-resource model variants + +Make resource efficiency an explicit product feature. + +Model families could include: + +```text +human_pbmc_large +human_pbmc_medium +human_pbmc_cpu +human_pbmc_lite +``` + +Potential methods: + +- fewer features. +- smaller hidden layers. +- quantization where appropriate. +- distilled models if validated. + +Publish resource/quality trade-offs rather than claiming one model is universally preferable. + +--- + +# 31. Hierarchical / ensemble fallback architecture + +For taxonomies with many classes, support coarse-to-fine models: + +```text +Broad classifier + | + Immune / Stromal / Epithelial + | + v +Lineage classifier + | + v +Subtype classifier +``` + +This can make model size, explainability, and failure analysis more manageable. + +--- + +# 32. Visualization + +Provide optional plotting helpers: + +```python +scalr.pl.prediction_umap(adata) +scalr.pl.confidence_umap(adata) +scalr.pl.uncertainty_umap(adata) +scalr.pl.marker_importance(result) +scalr.pl.confusion_matrix(metrics) +scalr.pl.calibration_curve(metrics) +``` + +Keep plotting dependencies optional so the core package remains lightweight. + +--- + +# 33. CLI redesign + +Add a real console entry point: + +```bash +scalr validate --input data.h5ad +scalr annotate --input data.h5ad --model human_pbmc --output annotated.h5ad +scalr train --input train.h5ad --labels cell_type --group donor_id --output models/pbmc_v1 +scalr explain --input annotated.h5ad --model models/pbmc_v1 --output explanations/ +scalr evaluate --input test.h5ad --model models/pbmc_v1 --labels cell_type +scalr benchmark --config benchmark.yaml +scalr models list +scalr models info human_pbmc +``` + +Keep `python pipeline.py --config ...` working for backward compatibility during the v2 transition. + +CLI requirements: + +- helpful `--help`. +- clear error messages. +- non-zero exit codes on validation failure. +- machine-readable JSON output where useful. +- quiet/verbose modes. +- deterministic logging. + +--- + +# 34. Installation and packaging redesign + +The current repository recommends a specific Python environment and the installation documentation says users who install via pip may still need to obtain `pipeline.py` and config files for the complete pipeline. citeturn799194search2 + +The v2 package should make the common library path self-contained. + +### Target installation + +```bash +pip install pyscalr +``` + +Optional dependencies: + +```bash +pip install pyscalr[analysis] +pip install pyscalr[explain] +pip install pyscalr[dev] +``` + +GPU installation should be documented separately instead of forcing one CUDA-specific wheel index into the base requirements. + +### Packaging goals + +- clear Python compatibility range. +- CPU path works from a clean environment. +- GPU instructions are explicit and versioned. +- no need to manually copy `pipeline.py`/YAML just to use the packaged API. +- support Conda installation. +- packaging smoke test in CI. + +--- + +# 35. R / Seurat ecosystem bridge + +Do not rewrite the engine in R. + +Provide an interoperability layer: + +```text +Seurat / SingleCellExperiment + | + v + interoperable format + | + v + scaLR + | + v + predictions / metadata +``` + +This can start as documented AnnData conversion and later become a lightweight R wrapper. + +--- + +# 36. Experiment tracking and reproducibility + +Every run should create a structured artifact directory: + +```text +experiment/ +├── manifest.json +├── config.yaml +├── metrics.json +├── model/ +├── feature_selection/ +├── predictions/ +├── plots/ +└── logs/ +``` + +Record: + +- scaLR version. +- model version. +- Python version. +- OS. +- CPU. +- GPU. +- CUDA. +- dependency versions. +- random seeds. +- input hash. +- model hash. +- config hash. +- preprocessing hash. +- runtime. +- peak RAM/VRAM. + +--- + +# 37. Benchmark suite + +The paper establishes low-resource/scaling as a central contribution; v2 should make this a continuously reproducible benchmark rather than a one-time paper experiment. citeturn794974search2 + +Create: + +```text +benchmarks/ +├── small/ +├── medium/ +├── million_cells/ +├── multi_million_cells/ +├── cross_donor/ +├── cross_platform/ +├── cross_tissue/ +├── rare_cell/ +├── unknown_cell/ +└── reproducibility/ +``` + +Report: + +### Quality + +- accuracy. +- macro-F1. +- weighted-F1. +- balanced accuracy. +- per-class recall. +- AUROC where appropriate. +- confusion matrix. +- rare-cell recall. + +### Reliability + +- Expected Calibration Error. +- Brier score. +- unknown-cell AUROC. +- FPR95/related OOD metrics where appropriate. + +### Resource usage + +- wall time. +- cells/sec. +- peak RAM. +- peak VRAM. +- CPU utilization. +- GPU utilization. +- model size. + +### Scaling dimensions + +Vary: + +- number of cells. +- number of genes. +- feature count. +- number of classes. + +Publish scaling curves, not only one aggregate table. + +--- + +# 38. Testing strategy + +## Unit tests + +- configuration schema. +- cross-field validation. +- gene normalization. +- gene alignment. +- preprocessing detection. +- preprocessing fit/apply separation. +- label mapping. +- calibration. +- unknown detection. +- model save/load. +- sparse matrix behavior. +- AnnData round-trip. +- resource planner. + +## Integration tests + +- toy end-to-end training. +- toy feature selection. +- prediction. +- explanation. +- calibration. +- artifact loading. +- CLI commands. +- checkpoint resume. + +## Regression tests + +Use one small frozen dataset and validate: + +- prediction outputs. +- selected features. +- artifact metadata. +- basic memory budget. + +## CI + +GitHub Actions should run: + +- lint/format. +- import smoke test. +- unit tests. +- CPU integration test. +- packaging/install test. +- documentation build where applicable. + +GPU tests may be a separate workflow. + +--- + +# 39. Logging redesign + +Replace fragmented logging with one structured logger. + +Every message should have: + +```text +run_id +stage +severity +elapsed_time +resource_snapshot (optional) +``` + +Example: + +```text +2026-09-21 14:10:32 | INFO | run=abc123 | stage=inference +cells=1,200,000 genes=5,000 batch=20,000 device=cpu +``` + +Provide console and file handlers. + +--- + +# 40. Abstract interfaces + +Introduce explicit abstract interfaces for components likely to be swapped: + +```python +class BaseModel(ABC): + ... + +class BaseFeatureSelector(ABC): + ... + +class BaseTrainer(ABC): + ... + +class BasePreprocessor(ABC): + ... + +class BaseCalibrator(ABC): + ... +``` + +Each should define minimal stable methods and avoid over-abstracting code that has only one implementation. + +Keep backwards compatibility while refactoring. + +--- + +# 41. Data quality and QC recommendations + +Provide: + +```python +scalr.qc(adata) +``` + +that reports, without silently modifying the data: + +- low-count cells. +- high mitochondrial percentage where metadata permits. +- extreme library sizes. +- high/low feature counts. +- potential doublet signals if enabled. + +Example: + +```text +12,452 cells detected +3.2% low-count +1.1% high mitochondrial +0.4% unusual transcript counts + +No cells were removed automatically. +``` + +The user decides on filtering thresholds. + +--- + +# 42. Documentation and onboarding + +Create three paths. + +## Beginner + +```text +10-minute quickstart +``` + +Only: + +```python +import scalr +result = scalr.annotate(adata, model="human_pbmc") +``` + +## Research user + +- custom training. +- feature selection. +- model explanation. +- benchmarking. +- grouped splits. + +## Infrastructure/HPC user + +- streaming. +- resource control. +- CLI. +- reproducible runs. +- profiling. + +Add a "When should I use scaLR?" page that explains strengths and limitations honestly. + +--- + +# 43. Clear scope: what scaLR should not try to become immediately + +Do not make these the first v2 milestone: + +- a universal foundation model. +- an all-purpose batch-integration package. +- a replacement for Scanpy/Seurat. +- a general-purpose probabilistic single-cell framework. + +Foundation-model or generative-model integrations can be future adapters after the public API, artifact format, uncertainty layer, and benchmark framework are stable. + +This keeps the project aligned with its resource-efficiency and interpretable classification identity. + +--- + +# 44. Competitive positioning + +The aim is not to claim that scaLR is universally better than every other single-cell tool. + +Instead define its niche explicitly: + +| Capability | scaLR 2.0 target | +|---|---| +| Low-resource inference | Core strength | +| Very large datasets | Core strength | +| Supervised annotation | Core workflow | +| Feature selection | Core workflow | +| Biomarker-oriented explanation | Core workflow | +| Uncertainty/abstention | First-class | +| Open-set detection | First-class | +| AnnData integration | First-class | +| Pretrained models | Planned/first-class | +| Generative representation learning | Not core | +| Full ecosystem replacement | Not goal | + +The benchmark suite should compare documented capabilities and measured metrics, not produce a universal winner/ranking. + +--- + +# 45. Proposed package structure + +```text +scalr/ +├── __init__.py +├── api.py +├── cli.py +├── result.py +│ +├── config/ +│ ├── schema.py +│ └── validation.py +│ +├── io/ +│ ├── anndata.py +│ ├── h5ad.py +│ └── streaming.py +│ +├── preprocessing/ +│ ├── detect.py +│ ├── normalize.py +│ ├── qc.py +│ ├── align.py +│ └── contract.py +│ +├── models/ +│ ├── base.py +│ ├── registry.py +│ ├── artifact.py +│ ├── inference.py +│ └── hierarchy.py +│ +├── training/ +│ ├── trainer.py +│ ├── sampling.py +│ ├── losses.py +│ └── splitting.py +│ +├── uncertainty/ +│ ├── calibration.py +│ ├── open_set.py +│ └── metrics.py +│ +├── explain/ +│ ├── importance.py +│ ├── markers.py +│ └── stability.py +│ +├── analysis/ +├── data/ +├── feature/ +├── nn/ +│ +├── visualization/ +│ ├── prediction.py +│ ├── calibration.py +│ └── explanation.py +│ +├── benchmark/ +│ ├── runner.py +│ ├── metrics.py +│ └── resources.py +│ +└── utils/ + ├── logging.py + ├── reproducibility.py + └── device.py +``` + +Do not force all existing code into this structure immediately. Migrate modules behind stable interfaces over time. + +--- + +# 46. GitHub issues / PR roadmap + +## Existing issues to retain and expand + +### #2 — CI and end-to-end tests + +**Goal:** establish a reliable test foundation. + +Acceptance criteria: + +- GitHub Actions workflow. +- unit test suite. +- toy end-to-end dataset. +- packaging/import test. +- CPU inference regression test. +- lint/format checks. + +### #3 — Conda installation + +Acceptance criteria: + +- clean `environment.yml` or documented Conda install. +- CPU instructions. +- GPU instructions. +- clean-environment CI test. + +### #4 — Config validation + +Acceptance criteria: + +- typed schema. +- cross-field validation. +- AnnData metadata validation. +- actionable errors. + +### #5 — Unified logging + +Acceptance criteria: + +- one logger. +- configurable handlers. +- run/stage IDs. +- consistent severity levels. + +### #6 — Abstract base classes + +Acceptance criteria: + +- `BaseModel`. +- `BaseTrainer`. +- `BaseFeatureSelector`. +- minimal stable interfaces. +- backwards-compatible adapters. + +## New issues + +### #7 — Public inference API + +Implement: + +- `annotate`. +- `load_model`. +- typed `PredictionResult`. + +Acceptance: + +- <15-line quickstart. +- AnnData support. +- CPU path. +- `device="auto"`. +- labels/confidence/probabilities. + +### #8 — AnnData validation and gene alignment + +Acceptance: + +- duplicate/invalid genes handled. +- normalization state reported. +- feature overlap report. +- minimum-overlap threshold. + +### #9 — Versioned model artifact format + +Acceptance: + +- manifest. +- model weights. +- preprocessing metadata. +- features. +- labels. +- calibration. +- reproducible load/save test. + +### #10 — Confidence calibration and open-set detection + +Acceptance: + +- calibrated confidence. +- entropy. +- unknown status. +- OOD score interface. +- novel-class benchmark. + +### #11 — Model registry + +Acceptance: + +- local list/info/download. +- metadata schema. +- compatibility validation. + +### #12 — CLI + +Acceptance: + +- validate. +- annotate. +- train. +- evaluate. +- explain. +- benchmark. +- models. + +### #13 — Sparse and bounded-memory inference + +Acceptance: + +- sparse test. +- chunked prediction. +- peak-memory benchmark. +- no unnecessary densification. + +### #14 — Visualization and reporting + +Acceptance: + +- prediction UMAP. +- confidence UMAP. +- uncertainty UMAP. +- marker/explanation plot. +- calibration plot. + +### #15 — Reproducible benchmark suite + +Acceptance: + +- frozen benchmark configs. +- quality metrics. +- resource metrics. +- scaling curves. +- environment metadata. + +### #16 — Rare-class training and evaluation + +Acceptance: + +- class-weighted loss. +- balanced sampling. +- macro-F1/balanced accuracy. +- per-class reporting. + +### #17 — Leakage-safe split and dataset checks + +Acceptance: + +- group-based split API. +- donor/patient/sample leakage warning. +- train/query preprocessing separation. + +### #18 — Hierarchical annotation + +Acceptance: + +- taxonomy representation. +- coarse-to-fine prediction. +- broad/subtype levels. + +### #19 — Doublet/mixed-cell screening + +Acceptance: + +- uncertainty/mixed-cell signal. +- explicit status. +- validation benchmark. + +### #20 — Feature-selection stability + +Acceptance: + +- repeated feature-selection runner. +- Jaccard/rank stability. +- selection frequency report. + +### #21 — Model cards and provenance + +Acceptance: + +- standard model-card template. +- data/model/code versioning. +- license metadata. + +### #22 — Lightweight model variants + +Acceptance: + +- at least one lightweight model. +- measured resource/quality trade-off. + +### #23 — Seurat/R interoperability + +Acceptance: + +- documented conversion workflow. +- round-trip labels/metadata test. + +### #24 — Zarr/streaming data support + +Acceptance: + +- tested chunked format. +- bounded-memory benchmark. + +--- + +# 47. Phased implementation plan + +## Phase 0 — Foundation / correctness + +1. #2 CI/tests +2. #3 Conda/install +3. #4 config validation +4. #5 logging +5. #6 abstract interfaces +6. #17 leakage-safe split and preprocessing separation + +### Exit criteria + +- clean install from scratch. +- deterministic toy training/inference. +- clear validation errors. +- CI green. + +--- + +## Phase 1 — Public usability + +7. #7 public API +8. #8 AnnData validation + gene alignment +9. #9 model artifact format +10. #12 CLI +11. AnnData output integration + +### Exit criteria + +```python +result = scalr.annotate( + adata, + model="human_pbmc", +) +``` + +works without YAML. + +--- + +## Phase 2 — Scientific reliability + +12. #10 confidence/calibration/open-set +13. #16 rare-class handling +14. #19 mixed/doublet screening +15. #18 hierarchical annotation +16. marker sanity checks +17. feature-selection stability #20 + +### Exit criteria + +Results include: + +- confidence. +- uncertainty. +- unknown status. +- per-class metrics. +- leakage-safe evaluation. + +--- + +## Phase 3 — Models and distribution + +18. #11 model registry +19. #21 model cards +20. automatic model selection +21. self-contained model packages +22. lightweight model variants #22 + +### Exit criteria + +A user can download a compatible model, validate it against AnnData, run inference, inspect provenance, and reproduce the result. + +--- + +## Phase 4 — Scale and performance + +23. #13 sparse/bounded-memory inference +24. #24 Zarr/streaming +25. device/batch auto-selection +26. benchmark suite #15 +27. resource profiling + +### Exit criteria + +The project can demonstrate scaling behavior with increasing cell counts and clear peak-RAM/VRAM measurements. + +--- + +## Phase 5 — Ecosystem and adoption + +28. #14 visualization/reporting +29. #23 R/Seurat bridge +30. optional notebook dashboard / lightweight web UI +31. tutorials and model cookbook + +### Exit criteria + +A Scanpy user, a CLI/HPC user, and an R/Seurat user each have a documented entry path. + +--- + +# 48. Definition of done for scaLR 2.0 + +A new user should be able to: + +```python +import scalr +result = scalr.annotate(adata, model="human_pbmc") +``` + +without writing YAML. + +An advanced user should still be able to use the configuration-driven pipeline. + +Every result should include: + +- prediction. +- confidence. +- uncertainty. +- model/version metadata. +- preprocessing metadata. + +Every model should be: + +- portable. +- self-describing. +- versioned. +- provenance-aware. + +Every benchmark should report both: + +- scientific quality. +- computational cost. + +Every installation path should work from a clean environment. + +The CPU workflow must work without a CUDA-specific installation requirement. + +--- + +# 49. Minimum viable v2 vs full v2 + +## MVP v2 + +Must-have: + +- public API. +- AnnData validation. +- gene alignment. +- model artifact format. +- CPU/GPU auto-device. +- CLI. +- confidence + abstention. +- CI. +- reproducibility metadata. +- basic benchmark suite. + +## Full v2 + +Adds: + +- model hub. +- hierarchical annotation. +- cluster refinement. +- doublet screening. +- feature stability. +- lightweight models. +- Zarr/cloud. +- R bridge. +- richer dashboards. + +This separation prevents the first v2 release from becoming too large. + +--- + +# 50. Recommended first 10 PRs + +### PR 1 — Packaging + CI baseline + +Clean installation, test matrix, linting, CPU smoke test. + +### PR 2 — Validation subsystem + +`ValidationReport`, config schema, AnnData checks. + +### PR 3 — Public API + +`annotate`, `train`, `load_model`, `PredictionResult`. + +### PR 4 — Model artifact + +Self-contained manifest/weights/preprocessing/labels. + +### PR 5 — Gene alignment + preprocessing contract + +Reference/query compatibility. + +### PR 6 — CLI + +`validate`, `annotate`, `train`, `evaluate`, `explain`. + +### PR 7 — Uncertainty + +Calibration, entropy, margin, unknown status. + +### PR 8 — Scientific evaluation + +Leakage-safe splitting, class imbalance, per-class metrics. + +### PR 9 — Explainability + feature stability + +Cell/class explanations, repeated feature-selection stability. + +### PR 10 — Benchmark harness + +Resource + quality + scaling curves. + +After these ten PRs, scaLR would have a materially stronger foundation before adding a model hub or cloud-scale data sources. + +--- + +# 51. Suggested quickstart for the future README + +```python +import scanpy as sc +import scalr + +adata = sc.read_h5ad("pbmc.h5ad") + +result = scalr.annotate( + adata, + model="human_pbmc", + device="auto", +) + +adata.obs["cell_type"] = result.labels +adata.obs["cell_type_confidence"] = result.confidence +adata.obs["cell_type_unknown"] = result.is_unknown + +sc.pl.umap(adata, color="cell_type") +``` + +That should be the first code users see, followed by the advanced YAML pipeline. + +--- + +# 52. Final strategic direction + +The most valuable evolution is not to add more DNN architectures for the sake of novelty. + +The opportunity is to make scaLR the following: + +```text + scaLR 2.0 + | + +------------+------------+ + | | | + Easy Reliable Scalable + | | | + AnnData uncertainty millions+ cells + Model Hub open-set low RAM + CLI calibration streaming + | | | + +------------+------------+ + | + Explainable + | + feature importance + biomarker analysis + stability reporting + | + Reproducible + | + model cards + artifact metadata + benchmarks +``` + +### Core proposition + +> **scaLR is a resource-efficient, interpretable, uncertainty-aware single-cell annotation and biomarker-analysis platform designed to scale from ordinary datasets to atlas-scale workloads.** + +That proposition preserves the technical contribution of the published work while addressing the practical barriers that currently prevent the project from being a broadly usable library. citeturn794974search2turn799194search2 + +--- + +# 53. Reference links + +- scaLR repository: https://github.com/infocusp/scaLR citeturn799194search2 +- scaLR paper: https://academic.oup.com/bib/article/26/3/bbaf243/8152766 citeturn794974search2 +- CellTypist documentation (useful ecosystem reference for model-based annotation UX): https://www.celltypist.org/tutorials/onlineguide + +--- + +## One-sentence roadmap + +**First make scaLR easy to install and call, then make predictions scientifically safer, then package/distribute models, and only after that push deeper into atlas-scale streaming and ecosystem integrations.** \ No newline at end of file diff --git a/scalr/__init__.py b/scalr/__init__.py index 1222f77..cf0fa29 100644 --- a/scalr/__init__.py +++ b/scalr/__init__.py @@ -1,9 +1,28 @@ from . import analysis from . import data from . import feature +from . import models from . import nn from . import utils +from .api import annotate +from .api import detect_drift +from .api import load_model +from .api import train +from .artifact import AnnotationModel +from .benchmark import run_benchmark from .data_ingestion_pipeline import DataIngestionPipeline +from .doublet import flag_possible_doublets +from .drift import DriftReport from .eval_and_analysis_pipeline import EvalAndAnalysisPipeline +from .explain import CellExplanation from .feature_extraction_pipeline import FeatureExtractionPipeline +from .feature_stability import compute_feature_stability +from .genes import align_genes +from .hierarchy import aggregate_to_broad +from .model_card import render_model_card from .model_training_pipeline import ModelTrainingPipeline +from .refinement import refine_with_clusters +from .result import PredictionResult +from .validation import validate + +__version__ = '2.0.0.dev0' diff --git a/scalr/analysis/__init__.py b/scalr/analysis/__init__.py index 1e03566..7e28f69 100644 --- a/scalr/analysis/__init__.py +++ b/scalr/analysis/__init__.py @@ -1,9 +1,35 @@ +"""`scalr.analysis`'s heavier submodules (scanpy/statsmodels/matplotlib-backed) +are imported lazily on first attribute access (PEP 562), so importing `scalr` +for something unrelated to analysis (e.g. the `scalr` CLI's `models`/`train`/ +`annotate` commands) doesn't pay their import cost. `getattr(module, name)` +(the `name`/`params` builder pattern used by `build_analyser`) transparently +triggers this, same as an eager import would. +""" + +import importlib + from ._analyser import AnalysisBase from ._analyser import build_analyser -from .dge_lmem import DgeLMEM -from .dge_pseudobulk import DgePseudoBulk -from .evaluation import generate_and_save_classification_report -from .evaluation import get_accuracy -from .gene_recall_curve import GeneRecallCurve -from .heatmap import Heatmap -from .roc_auc import RocAucCurve + +_LAZY_SUBMODULES = { + 'DgeLMEM': '.dge_lmem', + 'DgePseudoBulk': '.dge_pseudobulk', + 'generate_and_save_classification_report': '.evaluation', + 'get_accuracy': '.evaluation', + 'GeneRecallCurve': '.gene_recall_curve', + 'Heatmap': '.heatmap', + 'RocAucCurve': '.roc_auc', +} + + +def __getattr__(name): + module_name = _LAZY_SUBMODULES.get(name) + if module_name is None: + raise AttributeError(f'module {__name__!r} has no attribute {name!r}') + value = getattr(importlib.import_module(module_name, __name__), name) + globals()[name] = value + return value + + +def __dir__(): + return sorted(set(globals()) | set(_LAZY_SUBMODULES)) diff --git a/scalr/analysis/dge_pseudobulk.py b/scalr/analysis/dge_pseudobulk.py index d919727..a58ecec 100644 --- a/scalr/analysis/dge_pseudobulk.py +++ b/scalr/analysis/dge_pseudobulk.py @@ -94,7 +94,7 @@ def _make_design_matrix(self, adata: AnnData, cell_type: str): sum_subset = condition_subset[condition_subset.obs[ self.sum_column] == sum_sample] subdata = AnnData( - X=sum_subset[:].X.sum(axis=0).reshape( + X=np.asarray(sum_subset[:].X.sum(axis=0)).reshape( 1, len(sum_subset.var_names)), var=DataFrame(index=sum_subset.var_names), obs=DataFrame(index=[f'{sum_sample}_{condition}'])) diff --git a/scalr/api.py b/scalr/api.py new file mode 100644 index 0000000..bb46f85 --- /dev/null +++ b/scalr/api.py @@ -0,0 +1,444 @@ +"""This file implements the scaLR 2.0 simple, inference-first public API. + + import scalr + + model = scalr.train(adata, labels_key="cell_type", group_key="donor_id") + model.save("models/pbmc_v1") + + model = scalr.load_model("models/pbmc_v1") + result = model.predict(adata) + + result = scalr.annotate(adata, model="models/pbmc_v1") +""" + +import os +from typing import Optional, Union + +from anndata import AnnData +import numpy as np +import torch +from torch import nn + +import scalr +from scalr import models +from scalr.artifact import AnnotationModel +from scalr.artifact import load_model as _load_model +from scalr.calibration import OpenSetThresholds +from scalr.calibration import TemperatureScaler +from scalr.drift import compute_reference_stats +from scalr.drift import DriftReport +from scalr.hierarchy import validate_taxonomy +from scalr.leakage import check_group_leakage +from scalr.leakage import detect_likely_grouping_column +from scalr.leakage import group_safe_split +from scalr.metrics import check_class_imbalance +from scalr.metrics import compute_class_weights +from scalr.metrics import compute_classification_metrics +from scalr.nn.model import build_model +from scalr.result import PredictionResult +from scalr.utils import FlowLogger +from scalr.utils import set_seed +from scalr.validation import detect_normalization_state +from scalr.validation import validate + + +def _select_device(device: str) -> str: + if device == 'auto': + return 'cuda' if torch.cuda.is_available() else 'cpu' + return device + + +def _dense_batch(adata: AnnData, indices: np.ndarray, start: int, end: int): + chunk = adata[indices[start:end]].X + if not isinstance(chunk, np.ndarray): + chunk = chunk.toarray() + return np.asarray(chunk, dtype=np.float32) + + +def _iterate_batches(adata: AnnData, + indices: np.ndarray, + label_ids: np.ndarray, + batch_size: int, + shuffle: bool = False): + order = indices.copy() + if shuffle: + np.random.shuffle(order) + for start in range(0, len(order), batch_size): + batch_idx = order[start:start + batch_size] + x = _dense_batch(adata, batch_idx, 0, len(batch_idx)) + y = label_ids[batch_idx] + yield torch.as_tensor(x), torch.as_tensor(y, dtype=torch.long) + + +def train( + adata: AnnData, + labels_key: str, + group_key: Optional[str] = None, + features: Optional[list[str]] = None, + hidden_layers: tuple = (256, 64), + epochs: int = 15, + batch_size: int = 256, + lr: float = 1e-3, + device: str = 'auto', + split_ratio: tuple = (0.7, 0.1, 0.2), + class_balanced_loss: bool = True, + open_set: bool = True, + taxonomy: Optional[dict[str, str]] = None, + seed: int = 42, + verbose: bool = True, +) -> AnnotationModel: + """Train a self-contained scaLR annotation model from an AnnData object. + + Performs leakage-safe splitting (grouped by `group_key` when provided), + class-weighted training, temperature-scaling calibration, and evaluation + with macro/weighted-F1 and balanced accuracy on a held-out test split. + + Args: + adata: Training AnnData; `adata.X` should already be normalized + (the scaLR convention). `adata.obs[labels_key]` holds cell-type labels. + labels_key: Column in `adata.obs` with classification labels. + group_key: Column in `adata.obs` identifying donor/patient/sample. + When provided, splitting keeps every group in a single split, + preventing train/test leakage. Strongly recommended for reference + datasets with donor/patient structure. + features: Ordered gene list to use as model input. Defaults to + `adata.var_names`. + hidden_layers: Sizes of hidden layers of the classification network. + epochs: Number of training epochs. + batch_size: Training batch size. + lr: Adam learning rate. + device: 'auto', 'cpu' or 'cuda'. + split_ratio: (train, val, test) ratio for the internal split. + class_balanced_loss: When True, use inverse-frequency class weights in + the training loss, improving rare-cell performance. + open_set: When True, fit open-set (unknown-cell) abstention thresholds + from the validation split's confidence/entropy/margin distribution. + taxonomy: Optional mapping from every fine-grained class name (i.e. + every value in `adata.obs[labels_key]`) to a broad class name, + enabling `model.predict(adata, level="broad")` for hierarchical + (coarse-to-fine) annotation. + seed: Random seed for reproducibility. + verbose: Log progress via scaLR's FlowLogger. + + Returns: + A trained, calibrated, self-contained `AnnotationModel`. + """ + logger = FlowLogger('scalr.train') + set_seed(seed) + + report = validate(adata, labels_key=labels_key, group_key=group_key) + report.raise_if_unusable() + if verbose: + for w in report.warnings: + logger.warning(w) + + if group_key is None: + detected = detect_likely_grouping_column(adata.obs) + if detected: + logger.warning( + f'No group_key given, but column(s) {detected} look like a ' + 'donor/patient/sample identifier. Consider passing ' + 'group_key to avoid train/test leakage.') + + imbalance = check_class_imbalance(adata.obs[labels_key]) + if imbalance.small_classes and verbose: + logger.warning(str(imbalance)) + + features = list(features) if features is not None else list(adata.var_names) + class_names = sorted(adata.obs[labels_key].astype(str).unique().tolist()) + label2id = {c: i for i, c in enumerate(class_names)} + label_ids = adata.obs[labels_key].astype(str).map(label2id).values + + if taxonomy is not None: + validate_taxonomy(class_names, taxonomy) + + split = group_safe_split(adata.obs, labels_key, group_key, split_ratio, + seed) + if group_key is not None: + leakage = check_group_leakage( + { + name: adata.obs.iloc[idx] for name, idx in split.items() + }, group_key) + if leakage.has_leakage: + logger.warning(str(leakage)) + + resolved_device = _select_device(device) + model_config = { + 'name': 'SequentialModel', + 'params': { + 'layers': [len(features), *hidden_layers, + len(class_names)], + }, + } + model, model_config = build_model(model_config) + model.to(resolved_device) + + weights = None + if class_balanced_loss: + class_weights = compute_class_weights(label_ids[split['train']], + len(class_names)) + weights = torch.as_tensor(class_weights, + dtype=torch.float32).to(resolved_device) + loss_fn = nn.CrossEntropyLoss(weight=weights) + optimizer = torch.optim.Adam(model.parameters(), lr=lr) + + for epoch in range(epochs): + model.train() + epoch_loss = 0.0 + n_batches = 0 + for x, y in _iterate_batches(adata, + split['train'], + label_ids, + batch_size, + shuffle=True): + x, y = x.to(resolved_device), y.to(resolved_device) + optimizer.zero_grad() + out = model(x)['cls_output'] + loss = loss_fn(out, y) + loss.backward() + optimizer.step() + epoch_loss += loss.item() + n_batches += 1 + + if verbose: + logger.info( + f'epoch={epoch + 1}/{epochs} train_loss={epoch_loss / max(n_batches, 1):.4f}' + ) + + # --- Calibration on the validation split ------------------------------- + model.eval() + val_logits, val_labels = [], [] + with torch.no_grad(): + for x, y in _iterate_batches(adata, split['val'], label_ids, + batch_size): + x = x.to(resolved_device) + out = model(x)['cls_output'] + val_logits.append(out.cpu()) + val_labels.append(y) + val_logits = torch.cat(val_logits, dim=0) + val_labels = torch.cat(val_labels, dim=0) + + calibrator = TemperatureScaler() + if len(val_labels) > 0: + calibrator.fit(val_logits, val_labels) + + val_probs = calibrator.calibrate_probabilities(val_logits.numpy()) + open_set_thresholds = OpenSetThresholds() + if open_set and len(val_labels) > 0: + from scalr.calibration import predictive_entropy + from scalr.calibration import top1_top2_margin + confidence = val_probs.max(axis=1) + entropy = predictive_entropy(val_probs) + margin = top1_top2_margin(val_probs) + open_set_thresholds = OpenSetThresholds( + min_confidence=float(np.quantile(confidence, 0.05)), + max_entropy=float(np.quantile(entropy, 0.95)), + min_margin=float(np.quantile(margin, 0.05)), + ) + + # --- Evaluation on the held-out test split ------------------------------ + test_logits, test_labels = [], [] + with torch.no_grad(): + for x, y in _iterate_batches(adata, split['test'], label_ids, + batch_size): + x = x.to(resolved_device) + out = model(x)['cls_output'] + test_logits.append(out.cpu()) + test_labels.append(y) + + metrics = {} + if test_logits: + test_probs = calibrator.calibrate_probabilities( + torch.cat(test_logits, dim=0).numpy()) + test_preds = test_probs.argmax(axis=1) + test_labels_np = torch.cat(test_labels, dim=0).numpy() + class_metrics = compute_classification_metrics(test_labels_np, + test_preds, class_names) + metrics = { + 'macro_f1': + class_metrics.macro_f1, + 'weighted_f1': + class_metrics.weighted_f1, + 'balanced_accuracy': + class_metrics.balanced_accuracy, + 'per_class': + class_metrics.per_class.reset_index().to_dict(orient='records'), + 'n_test_cells': + int(len(test_labels_np)), + } + if verbose: + logger.info(str(class_metrics)) + + norm_state = detect_normalization_state(adata) + preprocessing = { + 'normalization': norm_state, + 'feature_count': len(features), + } + drift_reference = compute_reference_stats(adata[split['train']]) + + annotation_model = AnnotationModel( + model=model, + model_config=model_config, + class_names=class_names, + features=features, + preprocessing=preprocessing, + calibrator=calibrator, + open_set_thresholds=open_set_thresholds, + metrics=metrics, + taxonomy=taxonomy, + drift_reference=drift_reference, + metadata={ + 'scalr_version': getattr(scalr, '__version__', 'unknown'), + 'model_version': '1.0.0', + 'species': None, + 'labels_key': labels_key, + 'group_key': group_key, + 'random_seed': seed, + 'split_ratio': list(split_ratio), + }, + ) + return annotation_model + + +def load_model(dirpath_or_name: str) -> AnnotationModel: + """Load a self-contained scaLR model artifact. + + Args: + dirpath_or_name: A path to a saved model artifact directory, the name + of a model registered locally via `scalr.models.register` + (resolved via the local model registry, see `scalr/models.py`), + or a Hugging Face Hub repo id (`"/"`) — pulled + via `scalr.models.download` and cached in the local registry on + first use. Hub access requires `pip install "pyscaLR[hub]"`. + """ + if os.path.exists(os.path.join(dirpath_or_name, 'manifest.json')): + return _load_model(dirpath_or_name) + try: + return _load_model(models.resolve_path(dirpath_or_name)) + except FileNotFoundError: + if '/' in dirpath_or_name and not os.path.exists(dirpath_or_name): + models.download(dirpath_or_name) + return _load_model(models.resolve_path(dirpath_or_name)) + raise + + +def annotate( + adata: Union[AnnData, str], + model: Union[str, AnnotationModel], + device: str = 'auto', + preprocess: str = 'auto', + open_set: bool = True, + top_k: int = 5, + min_feature_overlap: float = 0.1, + level: str = 'fine', + cluster_refinement: Optional[str] = None, + flag_doublets: bool = False, + streaming: bool = False, + chunk_size: int = 20000, +) -> PredictionResult: + """Annotate cells in `adata` using a trained scaLR model. + + Args: + adata: Query AnnData object. + model: Either a path to a saved model artifact directory, the name of + a model registered locally via `scalr.models.register`, or an + already-loaded `AnnotationModel`. + device: 'auto', 'cpu' or 'cuda'. + preprocess: 'auto' detects and reports the query's normalization state + versus the model's training-time contract, warning (not silently + transforming) on a mismatch. 'skip' disables this check. + open_set: When True, low-confidence/high-entropy/low-margin cells are + flagged `is_unknown=True` instead of forced into a class. + top_k: Number of top classes to report per cell. + min_feature_overlap: Minimum required gene-overlap fraction. + level: 'fine' (default) or 'broad' for hierarchical annotation, when + the model was trained with a `taxonomy`. + cluster_refinement: When given, relabel cells to their cluster's + confidence-weighted majority class (an `adata.obs` column name, + or 'auto' to detect one). `None` (default) disables this. + flag_doublets: When True, screen for possible mixed/doublet cells and + set `result.is_possible_doublet`. + streaming: When True, or implicitly when `adata` is a file/directory + path, gene-align and score the data `chunk_size` cells at a time + instead of materializing the whole aligned matrix at once. + chunk_size: Cells scored per chunk when streaming. + + Returns: + A `PredictionResult` aligned to `adata.obs_names`. + """ + if isinstance(model, str): + model = load_model(model) + + if preprocess == 'auto' and isinstance(adata, AnnData): + query_state = detect_normalization_state(adata) + expected_state = model.preprocessing.get('normalization') + if expected_state and query_state not in (expected_state, 'unknown'): + FlowLogger('scalr.annotate').warning( + f'Query data looks like "{query_state}" but the model was ' + f'trained on "{expected_state}" data. Predictions may be unreliable.' + ) + elif preprocess not in ('auto', 'skip'): + raise NotImplementedError( + f'preprocess="{preprocess}" is not supported yet; use "auto" or "skip".' + ) + + return model.predict(adata, + device=device, + open_set=open_set, + top_k=top_k, + min_feature_overlap=min_feature_overlap, + level=level, + cluster_refinement=cluster_refinement, + flag_doublets=flag_doublets, + streaming=streaming, + chunk_size=chunk_size) + + +def detect_drift( + adata: AnnData, + model: Union[str, AnnotationModel], + min_feature_overlap: float = 0.1, + z_threshold: float = 2.0, + drift_fraction_threshold: float = 0.1, + severity_z_threshold: float = 5.0, + top_k: int = 20, +) -> DriftReport: + """Check whether `adata` looks like it comes from a different + distribution than `model`'s training data. + + A per-cell `is_unknown`/`is_possible_doublet` flag (see `annotate`) can + miss a whole cohort that has drifted if most individual cells still score + confidently. This instead compares `adata`'s gene-expression and library- + size/detected-gene-count distributions against the training reference + recorded at `train()` time — useful as a pre-flight check before trusting + predictions on a new cohort/batch/tissue. + + Args: + adata: Query AnnData object. + model: Either a path to a saved model artifact directory, the name of + a model registered locally via `scalr.models.register`, or an + already-loaded `AnnotationModel`. + min_feature_overlap: Minimum required gene-overlap fraction. + z_threshold: Per-gene standardized mean shift above which a gene + counts as "drifted". + drift_fraction_threshold: Fraction of drifted genes above which the + dataset as a whole is flagged `is_drifted`. + severity_z_threshold: A library-size or detected-gene-count shift + this extreme (in reference std deviations) flags `is_drifted` on + its own, even if per-gene drift looks mild. + top_k: Number of top drifted genes to report. + + Returns: + A `DriftReport`. + """ + if isinstance(model, str): + model = load_model(model) + + return model.detect_drift( + adata, + min_feature_overlap=min_feature_overlap, + z_threshold=z_threshold, + drift_fraction_threshold=drift_fraction_threshold, + severity_z_threshold=severity_z_threshold, + top_k=top_k, + ) diff --git a/scalr/artifact.py b/scalr/artifact.py new file mode 100644 index 0000000..3e33c2f --- /dev/null +++ b/scalr/artifact.py @@ -0,0 +1,583 @@ +"""This file implements the self-contained, versioned model artifact format. + +A scaLR model artifact is a directory containing everything needed to +reproduce inference without the original training config: + + model/ + ├── manifest.json # format/version, task, species, labels, hashes + ├── model.pt # model weights (scalr.nn.model state dict) + ├── model_config.json # class name + params to rebuild the architecture + ├── label_mapping.json # id2label / label2id + ├── features.json # ordered gene/feature list the model expects + ├── preprocessing.json # fit-time preprocessing contract + ├── calibration.json # temperature scaling + open-set thresholds + ├── metrics.json # evaluation metrics recorded at training time + └── drift_reference.json # training-set gene/library-size stats for drift detection +""" + +from dataclasses import asdict +import json +import os +from os import path +from typing import Optional, Union + +from anndata import AnnData +from anndata.experimental import AnnCollection +import numpy as np +import torch + +import scalr +from scalr.calibration import OpenSetThresholds +from scalr.calibration import predictive_entropy +from scalr.calibration import TemperatureScaler +from scalr.calibration import top1_top2_margin +from scalr.doublet import flag_possible_doublets +from scalr.drift import detect_drift as _detect_drift +from scalr.drift import DriftReport +from scalr.explain import attribute_cells +from scalr.explain import attribute_class +from scalr.explain import CellExplanation +from scalr.genes import align_genes +from scalr.genes import GeneAlignmentReport +from scalr.hierarchy import aggregate_to_broad +from scalr.hierarchy import validate_taxonomy +from scalr.model_card import render_model_card +from scalr.nn.model import build_model +from scalr.refinement import refine_with_clusters +from scalr.result import PredictionResult +from scalr.utils import read_data +from scalr.utils import write_data +from scalr.validation import validate + +FORMAT_VERSION = 1 + + +def _select_device(device: str) -> str: + if device == 'auto': + return 'cuda' if torch.cuda.is_available() else 'cpu' + return device + + +class AnnotationModel: + """A trained, self-contained scaLR annotation model. + + Wraps a `scalr.nn.model` network together with its label mapping, feature + list, preprocessing contract and confidence calibrator, so it can be + saved/loaded and used for inference without any external config. + """ + + def __init__( + self, + model: torch.nn.Module, + model_config: dict, + class_names: list[str], + features: list[str], + preprocessing: Optional[dict] = None, + calibrator: Optional[TemperatureScaler] = None, + open_set_thresholds: Optional[OpenSetThresholds] = None, + metrics: Optional[dict] = None, + metadata: Optional[dict] = None, + taxonomy: Optional[dict[str, str]] = None, + drift_reference: Optional[dict] = None, + ): + self.model = model + self.model_config = model_config + self.class_names = class_names + self.features = features + self.preprocessing = preprocessing or {} + self.calibrator = calibrator or TemperatureScaler(temperature=1.0) + self.open_set_thresholds = open_set_thresholds or OpenSetThresholds() + self.metrics = metrics or {} + self.metadata = metadata or {} + self.drift_reference = drift_reference + if taxonomy is not None: + validate_taxonomy(class_names, taxonomy) + self.taxonomy = taxonomy + + # ------------------------------------------------------------------ # + # Inference + # ------------------------------------------------------------------ # + def _score_aligned(self, aligned: AnnData, batch_size: int, + device: str) -> np.ndarray: + """Run the network over an already gene-aligned AnnData in batches + and return calibrated probabilities. Never densifies more than one + `batch_size` chunk of `aligned.X` at a time.""" + n = aligned.shape[0] + all_logits = [] + with torch.no_grad(): + for start in range(0, n, batch_size): + chunk = aligned[start:start + batch_size].X + if not isinstance(chunk, np.ndarray): + chunk = chunk.toarray() + x = torch.as_tensor(chunk, dtype=torch.float32).to(device) + out = self.model(x)['cls_output'] + all_logits.append(out.cpu()) + logits = torch.cat(all_logits, dim=0) + return self.calibrator.calibrate_probabilities(logits.numpy()) + + def _predict_streaming( + self, + source: Union[AnnData, AnnCollection], + chunk_size: int, + min_feature_overlap: float, + batch_size: int, + device: str, + ) -> tuple[list[str], np.ndarray, GeneAlignmentReport]: + """Score `source` `chunk_size` cells at a time, only ever + gene-aligning one chunk in memory at once, so peak alignment/scoring + memory does not grow with total cell count. + + For a single backed `.h5ad` file, `source` is first fully loaded + (backed row-slicing of a single file cannot be done lazily here); to + avoid ever holding the whole dataset in memory, pass a directory of + chunked `.h5ad` files (read as an `AnnCollection`), which is sliced + and loaded one chunk at a time. + """ + if isinstance(source, AnnData) and source.isbacked: + source = source.to_memory(copy=True) + + n_total = source.shape[0] + obs_names: list[str] = [] + probability_chunks = [] + gene_report = None + + for start in range(0, n_total, chunk_size): + end = min(start + chunk_size, n_total) + chunk = source[start:end] + if not isinstance(chunk, AnnData): + chunk = chunk.to_adata() + if hasattr(chunk, 'to_memory'): + chunk = chunk.to_memory(copy=True) + + if gene_report is None: + report = validate(chunk, + model_features=self.features, + min_feature_overlap=min_feature_overlap) + report.raise_if_unusable() + + aligned_chunk, chunk_gene_report = align_genes( + chunk, self.features, min_feature_overlap) + if not chunk_gene_report.is_usable: + raise ValueError('Gene coverage too low for reliable ' + f'prediction:\n{chunk_gene_report}') + if gene_report is None: + gene_report = chunk_gene_report + + obs_names.extend(list(chunk.obs_names)) + probability_chunks.append( + self._score_aligned(aligned_chunk, batch_size, device)) + + return obs_names, np.concatenate(probability_chunks, + axis=0), gene_report + + def predict( + self, + adata: Union[AnnData, AnnCollection, str], + device: str = 'auto', + batch_size: int = 4096, + open_set: bool = True, + top_k: int = 5, + min_feature_overlap: float = 0.1, + level: str = 'fine', + cluster_refinement: Optional[str] = None, + flag_doublets: bool = False, + streaming: bool = False, + chunk_size: int = 20000, + ) -> PredictionResult: + """Run gene-aligned, calibrated inference on an AnnData object. + + Args: + adata: Query AnnData/AnnCollection object, or a path to an + `.h5ad` file or directory of chunked `.h5ad` files (read via + `scalr.utils.read_data`). Does not need to share the model's + gene set or gene order; alignment is performed automatically. + device: 'auto', 'cpu' or 'cuda'. + batch_size: Number of cells scored per forward pass. + open_set: When True, cells failing the model's confidence/entropy/ + margin thresholds are flagged `is_unknown=True`. + top_k: Number of top classes to report per cell. + min_feature_overlap: Minimum required gene-overlap fraction; + raises if coverage falls below this. + level: 'fine' (default, the model's native classes) or 'broad' to + aggregate predictions to broad classes using `self.taxonomy` + (raises if the model has no taxonomy). + cluster_refinement: When given, relabel cells to their cluster's + confidence-weighted majority class. Either an `adata.obs` + column name, or 'auto' to detect a 'leiden'/'louvain'/ + 'cluster'/'clusters' column. `None` (default) disables this; + the raw, unrefined labels are always kept in + `result.metadata['raw_labels']` when refinement is applied. + flag_doublets: When True, screen for cells whose top-two class + probabilities are both substantial and close together, and + set `result.is_possible_doublet` — a heuristic screening + signal, not a validated doublet call. + streaming: When True, or implicitly when `adata` is a path, gene- + align and score the data `chunk_size` cells at a time instead + of materializing the whole aligned matrix at once, bounding + peak alignment/scoring memory to roughly + `chunk_size x n_features` regardless of total dataset size. + For out-of-core reads too (never loading the full raw dataset + into memory), pass a directory of chunked `.h5ad` files + rather than a single large `.h5ad` file. + chunk_size: Cells scored per chunk when `streaming`/path input is + used. + + Returns: + A `PredictionResult` aligned to `adata.obs_names`. + """ + if level == 'broad' and not self.taxonomy: + raise ValueError( + 'level="broad" requires a taxonomy; this model has none. ' + 'Pass `taxonomy` when training/constructing the model.') + + if isinstance(adata, str): + adata = read_data(adata, backed='r') + streaming = True + + resolved_device = _select_device(device) + self.model.to(resolved_device) + self.model.eval() + + if streaming or not isinstance(adata, AnnData): + obs_names, probabilities, gene_report = self._predict_streaming( + adata, chunk_size, min_feature_overlap, batch_size, + resolved_device) + else: + report = validate(adata, + model_features=self.features, + min_feature_overlap=min_feature_overlap) + report.raise_if_unusable() + + aligned, gene_report = align_genes(adata, self.features, + min_feature_overlap) + if not gene_report.is_usable: + raise ValueError('Gene coverage too low for reliable ' + f'prediction:\n{gene_report}') + obs_names = list(adata.obs_names) + probabilities = self._score_aligned(aligned, batch_size, + resolved_device) + + n = len(obs_names) + top1_ids = probabilities.argmax(axis=1) + labels = [self.class_names[i] for i in top1_ids] + confidence = probabilities.max(axis=1) + entropy = predictive_entropy(probabilities) + margin = top1_top2_margin(probabilities) + + if open_set: + is_unknown = self.open_set_thresholds.flag_unknown( + confidence, entropy, margin) + else: + is_unknown = np.zeros(n, dtype=bool) + + k = min(top_k, len(self.class_names)) + top_k_out = [] + top_k_idx = np.argsort(-probabilities, axis=1)[:, :k] + for row_idx, row in enumerate(top_k_idx): + top_k_out.append([(self.class_names[c], + float(probabilities[row_idx, c])) for c in row]) + + is_possible_doublet = None + if flag_doublets: + is_possible_doublet = flag_possible_doublets( + probabilities, entropy, margin) + + result = PredictionResult( + obs_names=obs_names, + labels=labels, + probabilities=probabilities, + class_names=self.class_names, + confidence=confidence, + entropy=entropy, + margin=margin, + is_unknown=is_unknown, + top_k=top_k_out, + is_possible_doublet=is_possible_doublet, + metadata={ + 'model_version': self.metadata.get('model_version'), + 'scalr_version': self.metadata.get('scalr_version'), + 'temperature': self.calibrator.temperature, + 'gene_coverage': gene_report.coverage, + 'missing_features': gene_report.missing_features, + 'open_set_thresholds': self.open_set_thresholds.to_dict(), + 'device': resolved_device, + }, + ) + + if level == 'broad': + result = aggregate_to_broad(result, self.taxonomy) + + if cluster_refinement is not None: + result = refine_with_clusters(result, + adata, + cluster_key=cluster_refinement) + + return result + + # ------------------------------------------------------------------ # + # Explainability + # ------------------------------------------------------------------ # + def explain( + self, + adata: AnnData, + indices: list[int], + top_k: int = 20, + device: str = 'auto', + min_feature_overlap: float = 0.1, + ) -> list[CellExplanation]: + """Explain individual cell predictions via Grad x Input gene attribution. + + For each requested cell, reports the model's prediction and the genes + whose expression most supported (positive attribution) or contradicted + (negative attribution) that prediction. This is a first-order + sensitivity explanation, not a validated biological one — see + `scalr.explain` for the method. + + Args: + adata: Query AnnData object; gene-aligned automatically. + indices: Positional row indices into `adata` to explain. + top_k: Number of supporting and contradictory genes to report per cell. + device: 'auto', 'cpu' or 'cuda'. + min_feature_overlap: Minimum required gene-overlap fraction. + + Returns: + One `CellExplanation` per entry in `indices`, in the same order. + """ + resolved_device = _select_device(device) + self.model.to(resolved_device) + self.model.eval() + + aligned, gene_report = align_genes(adata, self.features, + min_feature_overlap) + if not gene_report.is_usable: + raise ValueError('Gene coverage too low for reliable ' + f'explanation:\n{gene_report}') + + subset = aligned[indices] + chunk = subset.X + if not isinstance(chunk, np.ndarray): + chunk = chunk.toarray() + x = torch.as_tensor(chunk, dtype=torch.float32).to(resolved_device) + + with torch.no_grad(): + logits = self.model(x)['cls_output'] + probabilities = self.calibrator.calibrate_probabilities( + logits.cpu().numpy()) + predicted_ids = probabilities.argmax(axis=1) + confidence = probabilities.max(axis=1) + obs_names = [str(adata.obs_names[i]) for i in indices] + + return attribute_cells(self.model, + x, + self.features, + self.class_names, + predicted_ids, + confidence, + obs_names, + top_k=top_k) + + def explain_class( + self, + class_name: str, + top_k: int = 50, + device: str = 'auto', + ) -> tuple[list[tuple[str, float]], list[tuple[str, float]]]: + """Explain a class in general, via gradient at a zero-expression baseline. + + Returns: + `(supporting_genes, contradictory_genes)`, each a list of + `(gene, score)` pairs. See `scalr.explain.attribute_class`. + """ + resolved_device = _select_device(device) + self.model.to(resolved_device) + self.model.eval() + return attribute_class(self.model, + self.features, + self.class_names, + class_name, + top_k=top_k, + device=resolved_device) + + # ------------------------------------------------------------------ # + # Drift detection + # ------------------------------------------------------------------ # + def detect_drift( + self, + adata: AnnData, + min_feature_overlap: float = 0.1, + z_threshold: float = 2.0, + drift_fraction_threshold: float = 0.1, + severity_z_threshold: float = 5.0, + top_k: int = 20, + ) -> DriftReport: + """Check whether `adata` looks like it comes from a different + distribution than this model's training data. + + A per-cell `is_unknown`/`is_possible_doublet` flag can miss a whole + cohort that has drifted (e.g. a different tissue, protocol, or batch) + if most individual cells still score confidently; this instead + compares the query dataset's gene-expression and library-size/ + detected-gene-count distributions against the training reference + recorded at `train()` time. + + Args: + adata: Query AnnData object; gene-aligned automatically. + min_feature_overlap: Minimum required gene-overlap fraction. + z_threshold: Per-gene standardized mean shift above which a gene + counts as "drifted". + drift_fraction_threshold: Fraction of drifted genes above which + the dataset as a whole is flagged `is_drifted`. + severity_z_threshold: A library-size or detected-gene-count shift + this extreme (in reference std deviations) flags + `is_drifted` on its own. + top_k: Number of top drifted genes to report. + + Returns: + A `DriftReport`. + """ + if self.drift_reference is None: + raise ValueError( + 'This model has no drift reference (it was trained with an ' + 'older scaLR version, or `preprocessing`/`drift_reference` ' + 'was not recorded). Drift detection is unavailable; retrain ' + 'to enable it.') + + aligned, gene_report = align_genes(adata, self.features, + min_feature_overlap) + if not gene_report.is_usable: + raise ValueError('Gene coverage too low for reliable drift ' + f'detection:\n{gene_report}') + + return _detect_drift( + aligned, + self.drift_reference, + self.features, + z_threshold=z_threshold, + drift_fraction_threshold=drift_fraction_threshold, + severity_z_threshold=severity_z_threshold, + top_k=top_k, + ) + + # ------------------------------------------------------------------ # + # Persistence + # ------------------------------------------------------------------ # + def save(self, dirpath: str) -> None: + """Save this model as a self-contained artifact directory.""" + os.makedirs(dirpath, exist_ok=True) + + self.model.save_weights(path.join(dirpath, 'model.pt')) + write_data(self.model_config, path.join(dirpath, 'model_config.json')) + write_data( + { + 'label2id': { + c: i for i, c in enumerate(self.class_names) + }, + 'id2label': { + i: c for i, c in enumerate(self.class_names) + }, + }, path.join(dirpath, 'label_mapping.json')) + write_data({'features': self.features}, + path.join(dirpath, 'features.json')) + write_data(self.preprocessing, path.join(dirpath, 'preprocessing.json')) + write_data( + { + 'calibration': self.calibrator.to_dict(), + 'open_set_thresholds': self.open_set_thresholds.to_dict(), + }, path.join(dirpath, 'calibration.json')) + write_data(self.metrics, path.join(dirpath, 'metrics.json')) + + if self.taxonomy is not None: + write_data({'taxonomy': self.taxonomy}, + path.join(dirpath, 'taxonomy.json')) + + if self.drift_reference is not None: + write_data(self.drift_reference, + path.join(dirpath, 'drift_reference.json')) + + manifest = { + 'format_version': FORMAT_VERSION, + 'scalr_version': getattr(scalr, '__version__', 'unknown'), + 'task': 'cell_type_annotation', + 'feature_count': len(self.features), + 'labels': self.class_names, + **{ + k: v for k, v in self.metadata.items() if k not in ('scalr_version',) + }, + 'temperature': self.calibrator.temperature, + 'open_set_thresholds': self.open_set_thresholds.to_dict(), + 'has_taxonomy': self.taxonomy is not None, + 'has_drift_reference': self.drift_reference is not None, + } + write_data(manifest, path.join(dirpath, 'manifest.json')) + + with open(path.join(dirpath, 'README.md'), 'w') as fh: + fh.write(render_model_card(manifest, self.metrics)) + + @classmethod + def load(cls, dirpath: str) -> 'AnnotationModel': + """Load a self-contained model artifact directory.""" + manifest = read_data(path.join(dirpath, 'manifest.json')) + if manifest.get('format_version', 1) > FORMAT_VERSION: + raise ValueError( + f'Model artifact format_version={manifest["format_version"]} is ' + f'newer than the format this scaLR version supports ' + f'({FORMAT_VERSION}). Please upgrade scaLR.') + + model_config = read_data(path.join(dirpath, 'model_config.json')) + model, model_config = build_model(model_config) + model.load_weights(path.join(dirpath, 'model.pt')) + + label_mapping = read_data(path.join(dirpath, 'label_mapping.json')) + id2label = label_mapping['id2label'] + class_names = [id2label[str(i)] for i in range(len(id2label))] + + features = read_data(path.join(dirpath, 'features.json'))['features'] + + preprocessing = {} + if path.exists(path.join(dirpath, 'preprocessing.json')): + preprocessing = read_data(path.join(dirpath, 'preprocessing.json')) + + calibrator = TemperatureScaler() + open_set_thresholds = OpenSetThresholds() + calib_path = path.join(dirpath, 'calibration.json') + if path.exists(calib_path): + calib = read_data(calib_path) + if calib.get('calibration'): + calibrator = TemperatureScaler.from_dict(calib['calibration']) + if calib.get('open_set_thresholds'): + open_set_thresholds = OpenSetThresholds.from_dict( + calib['open_set_thresholds']) + + metrics = {} + metrics_path = path.join(dirpath, 'metrics.json') + if path.exists(metrics_path): + metrics = read_data(metrics_path) + + taxonomy = None + taxonomy_path = path.join(dirpath, 'taxonomy.json') + if path.exists(taxonomy_path): + taxonomy = read_data(taxonomy_path)['taxonomy'] + + drift_reference = None + drift_reference_path = path.join(dirpath, 'drift_reference.json') + if path.exists(drift_reference_path): + drift_reference = read_data(drift_reference_path) + + return cls( + model=model, + model_config=model_config, + taxonomy=taxonomy, + class_names=class_names, + features=features, + preprocessing=preprocessing, + calibrator=calibrator, + open_set_thresholds=open_set_thresholds, + metrics=metrics, + metadata=manifest, + drift_reference=drift_reference, + ) + + +def load_model(dirpath: str) -> AnnotationModel: + """Load a self-contained scaLR model artifact from `dirpath`.""" + return AnnotationModel.load(dirpath) diff --git a/scalr/benchmark.py b/scalr/benchmark.py new file mode 100644 index 0000000..5ff164c --- /dev/null +++ b/scalr/benchmark.py @@ -0,0 +1,112 @@ +"""This file implements a lightweight, reproducible benchmark harness. + +Reports scientific quality (macro-F1, balanced accuracy) alongside +computational cost (wall time, throughput, peak memory) for `scalr.train` + +`predict` runs across one or more datasets, plus the environment metadata +needed to make a run reproducible. This is a starting point for the "P1 +benchmark suite" — enough to produce scaling-curve-able numbers without a +separate configuration format. +""" + +from dataclasses import asdict +from dataclasses import dataclass +import platform +import resource +import sys +import time +from typing import Optional + +from anndata import AnnData +import pandas as pd +import torch + +import scalr + + +def peak_rss_mb() -> float: + """Peak resident set size of this process so far, in MB.""" + peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + # ru_maxrss is kilobytes on Linux, bytes on macOS. + return peak / 1024 if sys.platform != 'darwin' else peak / (1024 * 1024) + + +def capture_environment_metadata() -> dict: + """Capture environment metadata for reproducible benchmark reporting.""" + return { + 'scalr_version': getattr(scalr, '__version__', 'unknown'), + 'python_version': platform.python_version(), + 'torch_version': torch.__version__, + 'cuda_available': torch.cuda.is_available(), + 'platform': platform.platform(), + } + + +@dataclass +class BenchmarkResult: + """A single benchmark run's quality + resource metrics.""" + + name: str + n_cells: int + n_genes: int + n_classes: int + train_seconds: float + predict_seconds: float + cells_per_second_predict: float + peak_rss_mb: float + macro_f1: Optional[float] + balanced_accuracy: Optional[float] + + +def run_benchmark( + datasets: dict[str, AnnData], + labels_key: str, + group_key: Optional[str] = None, + train_kwargs: Optional[dict] = None, +) -> pd.DataFrame: + """Benchmark scaLR train+predict quality and resource cost across datasets. + + Args: + datasets: Mapping of benchmark name (e.g. 'small', 'medium', 'large') + to an AnnData for that configuration. + labels_key: Classification label column, shared across all datasets. + group_key: Optional donor/patient/sample column for leakage-safe splitting. + train_kwargs: Extra kwargs forwarded to `scalr.train` (e.g. `epochs`, `device`). + + Returns: + A DataFrame indexed by dataset name, with one row of quality + + resource-cost metrics per dataset — enough rows to plot a scaling + curve across dataset sizes. + """ + train_kwargs = dict(train_kwargs or {}) + rows = [] + + for name, adata in datasets.items(): + start = time.perf_counter() + model = scalr.train(adata, + labels_key=labels_key, + group_key=group_key, + verbose=False, + **train_kwargs) + train_seconds = time.perf_counter() - start + + start = time.perf_counter() + result = model.predict(adata, device=train_kwargs.get('device', 'auto')) + predict_seconds = time.perf_counter() - start + + rows.append( + BenchmarkResult( + name=name, + n_cells=adata.shape[0], + n_genes=adata.shape[1], + n_classes=len(model.class_names), + train_seconds=train_seconds, + predict_seconds=predict_seconds, + cells_per_second_predict=(adata.shape[0] / + predict_seconds if predict_seconds > 0 + else float('inf')), + peak_rss_mb=peak_rss_mb(), + macro_f1=model.metrics.get('macro_f1'), + balanced_accuracy=model.metrics.get('balanced_accuracy'), + )) + + return pd.DataFrame([asdict(r) for r in rows]).set_index('name') diff --git a/scalr/calibration.py b/scalr/calibration.py new file mode 100644 index 0000000..ea7efc5 --- /dev/null +++ b/scalr/calibration.py @@ -0,0 +1,137 @@ +"""This file implements confidence calibration and open-set (unknown-cell) detection. + +A raw softmax maximum is not a calibrated probability. This module fits a +temperature-scaling calibrator on held-out logits/labels, and derives +confidence, entropy, margin and an "unknown" abstention flag from calibrated +probabilities. +""" + +from dataclasses import dataclass + +import numpy as np +import torch +from torch import nn +from torch import Tensor + + +class TemperatureScaler: + """Single-parameter temperature scaling for multi-class calibration. + + Reference: Guo et al., 2017, "On Calibration of Modern Neural Networks". + """ + + def __init__(self, temperature: float = 1.0): + self.temperature = float(temperature) + + def fit(self, + logits: Tensor, + labels: Tensor, + lr: float = 0.01, + max_iter: int = 100) -> 'TemperatureScaler': + """Fit the scalar temperature on held-out (logits, labels) via NLL minimization.""" + logits = logits.detach().float() + labels = labels.detach().long() + + log_temperature = torch.zeros(1, requires_grad=True) + optimizer = torch.optim.LBFGS([log_temperature], + lr=lr, + max_iter=max_iter) + nll = nn.CrossEntropyLoss() + + def closure(): + optimizer.zero_grad() + temperature = torch.exp(log_temperature) + loss = nll(logits / temperature, labels) + loss.backward() + return loss + + optimizer.step(closure) + self.temperature = float(torch.exp(log_temperature).item()) + return self + + def calibrate_probabilities(self, logits: np.ndarray) -> np.ndarray: + """Apply the fitted temperature and return softmax probabilities.""" + logits_t = torch.as_tensor(logits, dtype=torch.float32) + probs = torch.softmax(logits_t / self.temperature, dim=-1) + return probs.numpy() + + def to_dict(self) -> dict: + return { + 'method': 'temperature_scaling', + 'temperature': self.temperature + } + + @classmethod + def from_dict(cls, d: dict) -> 'TemperatureScaler': + return cls(temperature=d.get('temperature', 1.0)) + + +def predictive_entropy(probabilities: np.ndarray) -> np.ndarray: + """Normalized (0-1) Shannon entropy of each row's class distribution.""" + n_classes = probabilities.shape[1] + eps = 1e-12 + raw_entropy = -(probabilities * np.log(probabilities + eps)).sum(axis=1) + max_entropy = np.log(n_classes) if n_classes > 1 else 1.0 + return raw_entropy / max_entropy if max_entropy > 0 else raw_entropy + + +def top1_top2_margin(probabilities: np.ndarray) -> np.ndarray: + """Difference between the top-1 and top-2 class probabilities per row.""" + sorted_probs = np.sort(probabilities, axis=1) + if probabilities.shape[1] < 2: + return sorted_probs[:, -1] + return sorted_probs[:, -1] - sorted_probs[:, -2] + + +def expected_calibration_error(probabilities: np.ndarray, + labels: np.ndarray, + n_bins: int = 15) -> float: + """Expected Calibration Error (ECE) using the top-1 confidence/correctness pairs.""" + confidences = probabilities.max(axis=1) + predictions = probabilities.argmax(axis=1) + accuracies = (predictions == labels).astype(np.float64) + + bin_edges = np.linspace(0, 1, n_bins + 1) + ece = 0.0 + n = len(labels) + for lo, hi in zip(bin_edges[:-1], bin_edges[1:]): + in_bin = (confidences > lo) & (confidences <= hi) + if not in_bin.any(): + continue + bin_conf = confidences[in_bin].mean() + bin_acc = accuracies[in_bin].mean() + ece += (in_bin.sum() / n) * abs(bin_conf - bin_acc) + return float(ece) + + +def brier_score(probabilities: np.ndarray, labels: np.ndarray) -> float: + """Multi-class Brier score: mean squared error between one-hot labels and probabilities.""" + n_classes = probabilities.shape[1] + one_hot = np.eye(n_classes)[labels] + return float(np.mean(np.sum((probabilities - one_hot)**2, axis=1))) + + +@dataclass +class OpenSetThresholds: + """Abstention thresholds used to flag predictions as `UNKNOWN`.""" + + min_confidence: float = 0.5 + max_entropy: float = 0.7 + min_margin: float = 0.0 + + def flag_unknown(self, confidence: np.ndarray, entropy: np.ndarray, + margin: np.ndarray) -> np.ndarray: + """Return a boolean array: True where a cell should be abstained on.""" + return ((confidence < self.min_confidence) | + (entropy > self.max_entropy) | (margin < self.min_margin)) + + def to_dict(self) -> dict: + return { + 'min_confidence': self.min_confidence, + 'max_entropy': self.max_entropy, + 'min_margin': self.min_margin, + } + + @classmethod + def from_dict(cls, d: dict) -> 'OpenSetThresholds': + return cls(**d) diff --git a/scalr/cli.py b/scalr/cli.py new file mode 100644 index 0000000..821ac40 --- /dev/null +++ b/scalr/cli.py @@ -0,0 +1,447 @@ +"""This file implements the `scalr` command-line interface for the simple public API. + + scalr validate --input data.h5ad [--labels-key cell_type] [--group-key donor_id] + scalr train --input train.h5ad --labels-key cell_type [--group-key donor_id] --output models/pbmc_v1 + scalr annotate --input data.h5ad --model models/pbmc_v1 --output annotated.h5ad + scalr explain --input data.h5ad --model models/pbmc_v1 --indices 10 20 30 + scalr explain --model models/pbmc_v1 --class-name T_cell + scalr evaluate --input test.h5ad --model models/pbmc_v1 --labels-key cell_type + scalr drift --input query.h5ad --model models/pbmc_v1 + scalr models list + scalr models info + scalr models push models/pbmc_v1 your-org/human-pbmc + scalr models pull your-org/human-pbmc + +This CLI wraps the in-memory simple API (`scalr.train`/`annotate`/`validate`); +the configuration-driven, chunked/streaming pipeline remains available via +`python pipeline.py --config ...` for large-scale/advanced runs. +""" + +import argparse +import json +from pathlib import Path +import subprocess +import sys + +import scalr +from scalr.metrics import compute_classification_metrics +from scalr.utils import read_data +from scalr.utils import write_data + + +def _cmd_validate(args) -> int: + adata = read_data(args.input, backed=None) + report = scalr.validate(adata, + labels_key=args.labels_key, + group_key=args.group_key, + min_feature_overlap=args.min_feature_overlap) + print(report) + return 0 if report.is_usable else 1 + + +def _cmd_train(args) -> int: + adata = read_data(args.input, backed=None) + features = None + if args.features_file: + with open(args.features_file, encoding='utf-8') as handle: + features = json.load(handle) + taxonomy = None + if args.taxonomy_file: + with open(args.taxonomy_file, encoding='utf-8') as handle: + taxonomy = json.load(handle) + model = scalr.train(adata, + labels_key=args.labels_key, + group_key=args.group_key, + features=features, + hidden_layers=tuple(args.hidden_layers), + epochs=args.epochs, + batch_size=args.batch_size, + lr=args.lr, + device=args.device, + split_ratio=tuple(args.split_ratio), + class_balanced_loss=not args.no_class_balanced_loss, + open_set=not args.no_open_set, + taxonomy=taxonomy, + seed=args.seed, + verbose=not args.quiet) + model.save(args.output) + print(f'Model saved to {args.output}') + if model.metrics: + print(f"macro-F1: {model.metrics.get('macro_f1', float('nan')):.4f} " + f"balanced accuracy: " + f"{model.metrics.get('balanced_accuracy', float('nan')):.4f}") + return 0 + + +def _cmd_annotate(args) -> int: + adata = read_data(args.input, backed=None) + result = scalr.annotate(adata, + model=args.model, + device=args.device, + open_set=not args.no_open_set, + min_feature_overlap=args.min_feature_overlap, + flag_doublets=args.flag_doublets, + streaming=args.streaming, + chunk_size=args.chunk_size) + result.write_to_adata(adata) + write_data(adata, args.output) + + n_unknown = int(result.is_unknown.sum()) + print(f'Annotated {len(result)} cells ({n_unknown} flagged unknown). ' + f'Written to {args.output}') + return 0 + + +def _cmd_explain(args) -> int: + model = scalr.load_model(args.model) + + if not args.class_name and not args.indices: + print('ERROR: pass --indices or --class-name.', file=sys.stderr) + return 1 + + if args.class_name: + supporting, contradictory = model.explain_class(args.class_name, + top_k=args.top_k, + device=args.device) + print(f'Class: {args.class_name}') + print() + print('Supporting genes') + for gene, score in supporting: + print(f'{gene} {score:+.4f}') + print() + print('Contradictory signals') + for gene, score in contradictory: + print(f'{gene} {score:+.4f}') + return 0 + + if not args.input: + print('ERROR: --input is required when using --indices.', + file=sys.stderr) + return 1 + + adata = read_data(args.input, backed=None) + explanations = model.explain(adata, + indices=args.indices, + top_k=args.top_k, + device=args.device) + for explanation in explanations: + print(explanation) + print() + return 0 + + +def _cmd_evaluate(args) -> int: + adata = read_data(args.input, backed=None) + model = scalr.load_model(args.model) + result = model.predict(adata, device=args.device) + + true_labels = adata.obs[args.labels_key].astype(str).values + label2id = {c: i for i, c in enumerate(model.class_names)} + unseen_labels = sorted(set(true_labels) - set(label2id)) + if unseen_labels: + print( + f'WARNING: {len(unseen_labels)} true label(s) not known to the ' + f'model will be excluded from evaluation: {unseen_labels}', + file=sys.stderr) + + keep_mask = [label in label2id for label in true_labels] + y_true = [ + label2id[label] for label, keep in zip(true_labels, keep_mask) if keep + ] + y_pred = [ + label2id[pred] for pred, keep in zip(result.labels, keep_mask) if keep + ] + + metrics = compute_classification_metrics(y_true, y_pred, model.class_names) + print(metrics) + print() + print(metrics.per_class) + return 0 + + +def _cmd_drift(args) -> int: + adata = read_data(args.input, backed=None) + model = scalr.load_model(args.model) + report = model.detect_drift( + adata, + min_feature_overlap=args.min_feature_overlap, + z_threshold=args.z_threshold, + drift_fraction_threshold=args.drift_fraction_threshold, + ) + print(report) + if args.fail_on_drift and report.is_drifted: + return 1 + return 0 + + +def _cmd_models_list(args) -> int: + names = scalr.models.list() + if not names: + print('No models registered locally.') + return 0 + for name in names: + print(name) + return 0 + + +def _cmd_models_info(args) -> int: + manifest = scalr.models.info(args.name) + print(json.dumps(manifest, indent=2)) + return 0 + + +def _cmd_models_push(args) -> int: + url = scalr.models.push(args.artifact_dir, + args.repo_id, + token=args.token, + private=args.private) + print(f'Pushed "{args.artifact_dir}" to {url}') + return 0 + + +def _cmd_models_pull(args) -> int: + local_path = scalr.models.download(args.repo_id, + revision=args.revision, + token=args.token, + name=args.name) + print(f'Pulled "{args.repo_id}" and registered locally as ' + f'"{args.name or args.repo_id}" at {local_path}') + return 0 + + +def _cmd_analyze(args) -> int: + """Run configured downstream analyses through the pipeline runner.""" + pipeline_path = Path(__file__).resolve().parent.parent / 'pipeline.py' + command = [sys.executable, str(pipeline_path), '--config', args.config] + if args.log: + command.append('--log') + if args.level: + command.extend(['--level', args.level]) + if args.logpath: + command.extend(['--logpath', args.logpath]) + if args.memoryprofiler: + command.append('--memoryprofiler') + return subprocess.run(command, check=False).returncode + + +def build_parser() -> argparse.ArgumentParser: + """Build the `scalr` CLI argument parser.""" + parser = argparse.ArgumentParser( + prog='scalr', description='scaLR 2.0 command-line interface.') + subparsers = parser.add_subparsers(dest='command', required=True) + + p = subparsers.add_parser( + 'validate', help='Validate an AnnData file for use with scaLR.') + p.add_argument('--input', required=True, help='Path to an .h5ad file.') + p.add_argument('--labels-key', default=None) + p.add_argument('--group-key', default=None) + p.add_argument('--min-feature-overlap', type=float, default=0.1) + p.set_defaults(func=_cmd_validate) + + p = subparsers.add_parser('train', help='Train a scaLR annotation model.') + p.add_argument('--input', + required=True, + help='Path to a training .h5ad file.') + p.add_argument('--labels-key', required=True) + p.add_argument('--group-key', default=None) + p.add_argument('--features-file', + default=None, + help='JSON file containing the ordered input gene list.') + p.add_argument('--hidden-layers', + nargs='+', + type=int, + default=[256, 64], + help='Hidden layer sizes. Default: 256 64.') + p.add_argument('--output', + required=True, + help='Directory to save the model artifact to.') + p.add_argument('--epochs', type=int, default=15) + p.add_argument('--batch-size', type=int, default=256) + p.add_argument('--lr', + type=float, + default=1e-3, + help='Adam learning rate. Default: 0.001.') + p.add_argument('--device', default='auto') + p.add_argument('--split-ratio', + nargs=3, + type=float, + default=[0.7, 0.1, 0.2], + metavar=('TRAIN', 'VAL', 'TEST')) + p.add_argument('--no-class-balanced-loss', + action='store_true', + help='Disable inverse-frequency class weighting.') + p.add_argument('--no-open-set', + action='store_true', + help='Disable unknown-cell abstention thresholds.') + p.add_argument('--taxonomy-file', + default=None, + help='JSON file mapping fine labels to broad labels.') + p.add_argument('--seed', type=int, default=42) + p.add_argument('--quiet', action='store_true') + p.set_defaults(func=_cmd_train) + + p = subparsers.add_parser( + 'annotate', help='Annotate cells using a trained scaLR model.') + p.add_argument('--input', required=True, help='Path to a query .h5ad file.') + p.add_argument( + '--model', + required=True, + help='Model artifact directory, or a locally registered model name.') + p.add_argument('--output', + required=True, + help='Path to write the annotated .h5ad to.') + p.add_argument('--device', default='auto') + p.add_argument('--no-open-set', + action='store_true', + help='Disable open-set abstention flagging.') + p.add_argument('--flag-doublets', + action='store_true', + help='Screen for possible mixed/doublet cells.') + p.add_argument('--min-feature-overlap', type=float, default=0.1) + p.add_argument( + '--streaming', + action='store_true', + help='Gene-align and score in chunks instead of all at once, ' + 'bounding peak memory for large inputs.') + p.add_argument('--chunk-size', + type=int, + default=20000, + help='Cells scored per chunk when --streaming is set.') + p.set_defaults(func=_cmd_annotate) + + p = subparsers.add_parser( + 'explain', + help="Explain a trained model's predictions via gene attribution.") + p.add_argument( + '--input', + default=None, + help='Path to a query .h5ad file (required unless --class-name is used).' + ) + p.add_argument( + '--model', + required=True, + help='Model artifact directory, or a locally registered model name.') + p.add_argument('--indices', + nargs='+', + type=int, + default=None, + help='Positional row indices of cells to explain.') + p.add_argument( + '--class-name', + default=None, + help='Explain this class in general instead of specific cells ' + '(ignores --input/--indices).') + p.add_argument('--top-k', type=int, default=20) + p.add_argument('--device', default='auto') + p.set_defaults(func=_cmd_explain) + + p = subparsers.add_parser( + 'evaluate', help='Evaluate a trained model against labeled test data.') + p.add_argument('--input', + required=True, + help='Path to a labeled test .h5ad file.') + p.add_argument('--model', required=True) + p.add_argument('--labels-key', required=True) + p.add_argument('--device', default='auto') + p.set_defaults(func=_cmd_evaluate) + + p = subparsers.add_parser( + 'drift', + help="Check whether a dataset's distribution has drifted from a " + 'model\'s training data.') + p.add_argument('--input', required=True, help='Path to a query .h5ad file.') + p.add_argument( + '--model', + required=True, + help='Model artifact directory, or a locally registered model name.') + p.add_argument('--min-feature-overlap', type=float, default=0.1) + p.add_argument( + '--z-threshold', + type=float, + default=2.0, + help='Per-gene standardized mean shift above which a gene counts as ' + 'drifted. Default: 2.0.') + p.add_argument( + '--drift-fraction-threshold', + type=float, + default=0.1, + help='Fraction of drifted genes above which the dataset is flagged ' + 'as drifted. Default: 0.1.') + p.add_argument( + '--fail-on-drift', + action='store_true', + help='Exit with a non-zero status if drift is detected, for use as ' + 'a CI/pipeline gate.') + p.set_defaults(func=_cmd_drift) + + p_models = subparsers.add_parser('models', + help='Manage locally registered models.') + models_sub = p_models.add_subparsers(dest='models_command', required=True) + p_list = models_sub.add_parser('list', + help='List locally registered models.') + p_list.set_defaults(func=_cmd_models_list) + p_info = models_sub.add_parser('info', + help="Show a registered model's manifest.") + p_info.add_argument('name') + p_info.set_defaults(func=_cmd_models_info) + p_push = models_sub.add_parser( + 'push', + help='Publish a saved model artifact to the Hugging Face Hub. ' + 'Requires: pip install "pyscaLR[hub]".') + p_push.add_argument('artifact_dir', + help='Path to a saved model artifact directory.') + p_push.add_argument( + 'repo_id', + help='Hugging Face repo id to publish to, e.g. org/model-name.') + p_push.add_argument( + '--private', + action='store_true', + help='Create the repo as private if it does not already exist.') + p_push.add_argument( + '--token', + default=None, + help='Hugging Face auth token (defaults to HF_TOKEN/huggingface-cli ' + 'login).') + p_push.set_defaults(func=_cmd_models_push) + p_pull = models_sub.add_parser( + 'pull', + help='Download a model artifact from the Hugging Face Hub and ' + 'register it locally. Requires: pip install "pyscaLR[hub]".') + p_pull.add_argument( + 'repo_id', help='Hugging Face repo id to pull, e.g. org/model-name.') + p_pull.add_argument('--revision', + default=None, + help='Hub revision (branch/tag/commit) to pull.') + p_pull.add_argument('--token', + default=None, + help='Hugging Face auth token, for private repos.') + p_pull.add_argument( + '--name', + default=None, + help='Local registry name to register under (default: repo_id).') + p_pull.set_defaults(func=_cmd_models_pull) + + p = subparsers.add_parser( + 'analyze', + help='Run downstream analyses configured in a YAML pipeline config.') + p.add_argument('--config', required=True, help='Path to config.yaml.') + p.add_argument('--log', action='store_true', help='Save experiment logs.') + p.add_argument('--level', default=None, help='Logging level, e.g. INFO.') + p.add_argument('--logpath', default=None, help='Path to the log file.') + p.add_argument('--memoryprofiler', + action='store_true', + help='Record peak memory usage.') + p.set_defaults(func=_cmd_analyze) + + return parser + + +def main(argv: list = None) -> int: + """Entry point for the `scalr` console script.""" + parser = build_parser() + args = parser.parse_args(argv) + return args.func(args) + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/scalr/data_ingestion_pipeline.py b/scalr/data_ingestion_pipeline.py index 7f1d4fe..43ab6c3 100644 --- a/scalr/data_ingestion_pipeline.py +++ b/scalr/data_ingestion_pipeline.py @@ -42,29 +42,43 @@ def __init__(self, data_config: dict, dirpath: str = '.'): self.datadir = dirpath - def generate_train_val_test_split(self): - """A function to split data into train, validation and test sets.""" + def _validate_train_val_test_config(self): + """Validate the `train_val_test` config section before it is used. - # TODO: Move to config validation - if self.data_config['train_val_test'].get( - 'full_datapath', - False) and self.data_config['train_val_test'].get( - 'split_datapaths', False): + Raises: + ValueError: If pipeline-generated paths are supplied by the user, + or neither `full_datapath` nor `split_datapaths` is given. + """ + train_val_test_config = self.data_config['train_val_test'] + + if train_val_test_config.get('full_datapath', + False) and train_val_test_config.get( + 'split_datapaths', False): self.flow_logger.warning( '''`full_datapath` and `split_datapaths` are both provided in the config. Pipeline will use `full_datapath` and it will overwrite the `split_datapaths`. ''') - # TODO: Move to config validation - if self.data_config['train_val_test'].get( - 'feature_subset_datapaths' - ) or self.data_config['train_val_test'].get('final_datapaths'): + if train_val_test_config.get( + 'feature_subset_datapaths') or train_val_test_config.get( + 'final_datapaths'): raise ValueError( '''`final_datapaths` or `feature_subset_datapaths` can not be provided by the user in the config. These paths are generated by the pipeline! ''') + if not (train_val_test_config.get('full_datapath') or + train_val_test_config.get('split_datapaths')): + raise ValueError( + 'No Data Provided. Please provide `full_datapath` or `split_datapaths`!' + ) + + def generate_train_val_test_split(self): + """A function to split data into train, validation and test sets.""" + + self._validate_train_val_test_config() + if self.data_config['train_val_test'].get('full_datapath'): self.flow_logger.info('Generating Train, Validation and Test sets') if not self.target: @@ -109,16 +123,10 @@ def generate_train_val_test_split(self): self.data_config['train_val_test'][ 'split_datapaths'] = train_val_test_split_dirpath - elif self.data_config['train_val_test'].get('split_datapaths'): + else: self.flow_logger.info( 'Reading Train, Validation and Test sets from config') - # TODO: Move to config validation - else: - raise ValueError( - 'No Data Provided. Please provide `full_datapath` or `split_datapaths`!' - ) - def preprocess_data(self): """A function to apply preprocessing on data splits.""" diff --git a/scalr/doublet.py b/scalr/doublet.py new file mode 100644 index 0000000..2b0985b --- /dev/null +++ b/scalr/doublet.py @@ -0,0 +1,46 @@ +"""This file implements a doublet/mixed-cell screening signal. + +A classifier should not silently force an ambiguous, likely-mixed cell into a +single lineage. This module flags cells whose top-two class probabilities are +both substantial and close together as `POSSIBLE_DOUBLET` candidates. + +This is a screening signal derived from prediction uncertainty, not a +validated doublet caller (e.g. Scrublet/DoubletFinder) — treat it as "worth a +second look", not a definitive biological diagnosis. +""" + +import numpy as np + + +def flag_possible_doublets( + probabilities: np.ndarray, + entropy: np.ndarray, + margin: np.ndarray, + max_margin: float = 0.15, + min_entropy: float = 0.5, + min_second_probability: float = 0.25, +) -> np.ndarray: + """Flag cells with two closely-competing, substantial class probabilities. + + Args: + probabilities: [n_cells, n_classes] calibrated class probabilities. + entropy: Normalized predictive entropy per cell (see `scalr.calibration`). + margin: Top1-top2 probability margin per cell. + max_margin: A cell is only flagged if its margin is <= this. + min_entropy: A cell is only flagged if its entropy is >= this. + min_second_probability: A cell is only flagged if its second-highest + class probability is >= this (i.e. a real competing signal, not + just low confidence in a single class). + + Returns: + Boolean array, True where the prediction looks like a possible + mixed/doublet cell. + """ + if probabilities.shape[1] < 2: + return np.zeros(probabilities.shape[0], dtype=bool) + + sorted_probs = np.sort(probabilities, axis=1) + second_probability = sorted_probs[:, -2] + + return ((margin <= max_margin) & (entropy >= min_entropy) & + (second_probability >= min_second_probability)) diff --git a/scalr/drift.py b/scalr/drift.py new file mode 100644 index 0000000..a3bda29 --- /dev/null +++ b/scalr/drift.py @@ -0,0 +1,192 @@ +"""This file implements dataset-level drift detection: whether a whole new +(query) dataset's gene-expression distribution looks meaningfully different +from the data a model was trained on. + +Per-cell signals (`is_unknown` open-set abstention, `is_possible_doublet`, +see `scalr.calibration`/`scalr.doublet`) catch individual out-of-distribution +cells. This catches the case a per-cell signal can miss: a whole new cohort, +batch, tissue, or protocol that looks systematically different even if most +individual cells still score confidently, which matters most for a low- +resource classifier deployed on new, unseen cohorts. + +Reference statistics (per-gene mean/std, per-cell library size and detected- +gene-count distributions) are computed once at training time from the +training split and stored in the model artifact (`drift_reference.json`), so +drift detection at inference time never needs the original training data. +""" + +from dataclasses import dataclass +from dataclasses import field + +from anndata import AnnData +import numpy as np +from scipy import sparse + + +def _row_stats(X) -> tuple[np.ndarray, np.ndarray]: + """Per-cell (library_size, n_genes_detected), without densifying `X`.""" + if sparse.issparse(X): + library_size = np.asarray(X.sum(axis=1)).ravel() + n_genes_detected = np.asarray((X > 0).sum(axis=1)).ravel() + else: + X = np.asarray(X) + library_size = X.sum(axis=1) + n_genes_detected = (X > 0).sum(axis=1) + return library_size.astype(np.float64), n_genes_detected.astype(np.float64) + + +def _column_mean_std(X) -> tuple[np.ndarray, np.ndarray]: + """Per-gene (mean, std), without densifying `X`.""" + if sparse.issparse(X): + mean = np.asarray(X.mean(axis=0)).ravel() + mean_sq = np.asarray(X.multiply(X).mean(axis=0)).ravel() + else: + X = np.asarray(X) + mean = X.mean(axis=0) + mean_sq = (X**2).mean(axis=0) + std = np.sqrt(np.maximum(mean_sq - mean**2, 0.0)) + return mean, std + + +def compute_reference_stats(adata: AnnData) -> dict: + """Compute per-gene and per-cell reference statistics from `adata`. + + `adata` is expected to already be restricted to the model's training + cells and aligned to the model's feature order (e.g. the training split, + in `features` order), so `gene_mean`/`gene_std` line up 1:1 with + `AnnotationModel.features` at drift-detection time. + + Returns a JSON-serializable dict, suitable for `drift_reference.json` in + a model artifact. + """ + gene_mean, gene_std = _column_mean_std(adata.X) + library_size, n_genes_detected = _row_stats(adata.X) + + return { + 'n_cells': int(adata.shape[0]), + 'gene_mean': gene_mean.tolist(), + 'gene_std': gene_std.tolist(), + 'library_size_mean': float(library_size.mean()), + 'library_size_std': float(library_size.std()), + 'n_genes_detected_mean': float(n_genes_detected.mean()), + 'n_genes_detected_std': float(n_genes_detected.std()), + } + + +@dataclass +class DriftReport: + """Comparison of a query dataset's distribution against a model's + training reference statistics.""" + + n_query_cells: int + n_reference_cells: int + n_features: int + drifted_gene_fraction: float + top_drifted_genes: list[tuple[str, float]] = field(default_factory=list) + library_size_z: float = 0.0 + n_genes_detected_z: float = 0.0 + drift_fraction_threshold: float = 0.1 + severity_z_threshold: float = 5.0 + + @property + def is_drifted(self) -> bool: + return (self.drifted_gene_fraction >= self.drift_fraction_threshold or + abs(self.library_size_z) >= self.severity_z_threshold or + abs(self.n_genes_detected_z) >= self.severity_z_threshold) + + def __str__(self) -> str: + lines = [ + 'scaLR drift report', + '-' * 19, + f'Query cells: {self.n_query_cells:,}', + f'Reference (training) cells: {self.n_reference_cells:,}', + f'Drifted genes: ' + f'{self.drifted_gene_fraction * 100:.1f}% of {self.n_features:,} ' + f'(threshold {self.drift_fraction_threshold * 100:.0f}%)', + f'Library size shift (z): {self.library_size_z:+.2f}', + f'Genes-detected shift (z): {self.n_genes_detected_z:+.2f}', + '', + 'Top drifted genes (|z-score| of mean expression shift):', + ] + if self.top_drifted_genes: + lines += [ + f' {gene}: z={z:+.2f}' for gene, z in self.top_drifted_genes + ] + else: + lines.append(' None') + lines.append('') + lines.append(f'Status: {"DRIFTED" if self.is_drifted else "OK"}') + return '\n'.join(lines) + + +def detect_drift( + query: AnnData, + reference: dict, + features: list[str], + z_threshold: float = 2.0, + drift_fraction_threshold: float = 0.1, + severity_z_threshold: float = 5.0, + top_k: int = 20, +) -> DriftReport: + """Compare a gene-aligned query dataset against a model's training + reference statistics. + + Args: + query: Query AnnData, already gene-aligned to `features` (e.g. via + `scalr.align_genes`), so its columns line up 1:1 with + `reference['gene_mean']`/`reference['gene_std']`. + reference: Reference stats dict, as produced by + `compute_reference_stats` (stored in a model artifact as + `drift_reference.json`). + features: The model's ordered feature list, for labeling drifted + genes by name. + z_threshold: Per-gene standardized mean shift above which a gene + counts as "drifted". + drift_fraction_threshold: Fraction of drifted genes above which the + dataset as a whole is flagged `is_drifted`. + severity_z_threshold: A library-size or detected-gene-count shift + this extreme (in reference std deviations) flags `is_drifted` on + its own, even if per-gene drift looks mild — catches e.g. a + sequencing-depth/protocol change a per-gene test can miss. + top_k: Number of top drifted genes to report. + + Returns: + A `DriftReport`. + """ + ref_mean = np.asarray(reference['gene_mean'], dtype=np.float64) + ref_std = np.asarray(reference['gene_std'], dtype=np.float64) + if len(ref_mean) != len(features) or query.shape[1] != len(features): + raise ValueError( + 'Reference stats/features/query column count mismatch ' + f'({len(ref_mean)} vs {len(features)} vs {query.shape[1]}); ' + 'query must already be gene-aligned to `features`.') + + query_mean, _ = _column_mean_std(query.X) + safe_std = np.where(ref_std > 1e-8, ref_std, 1e-8) + gene_z = (query_mean - ref_mean) / safe_std + + drifted_mask = np.abs(gene_z) >= z_threshold + drifted_fraction = float(drifted_mask.mean()) if len(gene_z) else 0.0 + + order = np.argsort(-np.abs(gene_z))[:top_k] + top_drifted = [(features[i], float(gene_z[i])) for i in order] + + library_size, n_genes_detected = _row_stats(query.X) + ref_lib_std = reference.get('library_size_std') or 1e-8 + ref_genes_std = reference.get('n_genes_detected_std') or 1e-8 + library_size_z = (library_size.mean() - + reference.get('library_size_mean', 0.0)) / ref_lib_std + n_genes_detected_z = (n_genes_detected.mean() - reference.get( + 'n_genes_detected_mean', 0.0)) / ref_genes_std + + return DriftReport( + n_query_cells=int(query.shape[0]), + n_reference_cells=int(reference.get('n_cells', 0)), + n_features=len(features), + drifted_gene_fraction=drifted_fraction, + top_drifted_genes=top_drifted, + library_size_z=float(library_size_z), + n_genes_detected_z=float(n_genes_detected_z), + drift_fraction_threshold=drift_fraction_threshold, + severity_z_threshold=severity_z_threshold, + ) diff --git a/scalr/explain.py b/scalr/explain.py new file mode 100644 index 0000000..0d52372 --- /dev/null +++ b/scalr/explain.py @@ -0,0 +1,150 @@ +"""This file implements gradient-based gene attributions for individual cell +predictions and per-class signatures. + +This is a first-order sensitivity explanation (Grad x Input saliency), not a +causal or mechanistic one: it answers "which genes, in this cell's +expression profile, pushed the model's output most strongly toward (or away +from) the predicted class" via the gradient of the class logit with respect +to each input gene, scaled by that gene's expression. Like the doublet/ +refinement heuristics elsewhere in this layer, treat it as a screening/ +interpretation aid, not ground truth biology. +""" + +from dataclasses import dataclass + +import numpy as np +import torch +from torch import nn + + +@dataclass +class CellExplanation: + """Per-cell gene attribution for the model's top prediction.""" + + obs_name: str + prediction: str + confidence: float + supporting_genes: list[tuple[str, float]] + contradictory_genes: list[tuple[str, float]] + + def __str__(self) -> str: + lines = [ + self.obs_name, + f'Prediction: {self.prediction}', + f'Confidence: {self.confidence:.2f}', + '', + 'Supporting genes', + ] + lines += [f'{g} {s:+.2f}' for g, s in self.supporting_genes] + lines.append('') + lines.append('Contradictory signals') + lines += [f'{g} {s:+.2f}' for g, s in self.contradictory_genes] + return '\n'.join(lines) + + +def _grad_x_input(model: nn.Module, x: torch.Tensor, + class_idx: torch.Tensor) -> np.ndarray: + """Gradient-of-logit x input attribution for a batch of cells.""" + x = x.clone().detach().requires_grad_(True) + logits = model(x)['cls_output'] + selected = logits[torch.arange(logits.shape[0]), class_idx] + model.zero_grad(set_to_none=True) + selected.sum().backward() + attributions = x.grad.detach() * x.detach() + return attributions.cpu().numpy() + + +def attribute_cells( + model: nn.Module, + x: torch.Tensor, + features: list[str], + class_names: list[str], + predicted_ids: np.ndarray, + confidence: np.ndarray, + obs_names: list[str], + top_k: int = 20, +) -> list[CellExplanation]: + """Explain each row of `x` in terms of its aligned gene attributions. + + Args: + model: The underlying network, in eval mode, on the same device as `x`. + x: Aligned, dense `[n_cells, n_features]` input tensor. + features: Gene names for each column of `x`, in the model's feature order. + class_names: Model's class label list. + predicted_ids: Per-cell predicted class index to explain. + confidence: Per-cell top-1 probability, carried through for display. + obs_names: Per-cell identifiers, aligned to `x`'s rows. + top_k: Number of supporting and contradictory genes to keep per cell. + + Returns: + One `CellExplanation` per row of `x`, in the same order. + """ + class_idx = torch.as_tensor(predicted_ids, + dtype=torch.long, + device=x.device) + attributions = _grad_x_input(model, x, class_idx) + + explanations = [] + for i in range(x.shape[0]): + scores = attributions[i] + descending = np.argsort(-scores) + supporting = [(features[j], float(scores[j])) + for j in descending[:top_k] + if scores[j] > 0] + ascending = descending[::-1] + contradictory = [(features[j], float(scores[j])) + for j in ascending[:top_k] + if scores[j] < 0] + explanations.append( + CellExplanation( + obs_name=obs_names[i], + prediction=class_names[predicted_ids[i]], + confidence=float(confidence[i]), + supporting_genes=supporting, + contradictory_genes=contradictory, + )) + return explanations + + +def attribute_class( + model: nn.Module, + features: list[str], + class_names: list[str], + class_name: str, + top_k: int = 50, + device: str = 'cpu', +) -> tuple[list[tuple[str, float]], list[tuple[str, float]]]: + """Explain a class in general, independent of any specific cell. + + Computes the gradient of the class logit with respect to a zero- + expression baseline input. This approximates the genes the model relies + on most for this class near a neutral expression profile; unlike + `attribute_cells`, it is not derived from any observed cell. + + Returns: + `(supporting_genes, contradictory_genes)`, each a list of + `(gene, score)` sorted by descending/ascending attribution. + """ + if class_name not in class_names: + raise ValueError(f'Unknown class {class_name!r}; expected one of ' + f'{class_names}.') + class_id = class_names.index(class_name) + + x = torch.zeros((1, len(features)), + dtype=torch.float32, + device=device, + requires_grad=True) + logits = model(x)['cls_output'] + model.zero_grad(set_to_none=True) + logits[0, class_id].backward() + scores = x.grad.detach().cpu().numpy()[0] + + descending = np.argsort(-scores) + supporting = [(features[j], float(scores[j])) + for j in descending[:top_k] + if scores[j] > 0] + ascending = descending[::-1] + contradictory = [(features[j], float(scores[j])) + for j in ascending[:top_k] + if scores[j] < 0] + return supporting, contradictory diff --git a/scalr/feature/feature_subsetting.py b/scalr/feature/feature_subsetting.py index 561d20e..0b7b8c9 100644 --- a/scalr/feature/feature_subsetting.py +++ b/scalr/feature/feature_subsetting.py @@ -155,20 +155,29 @@ def train_chunked_model(i, start): model_trainer.build_model_training_artifacts() best_model = model_trainer.train() - self.chunk_model_config, self.chunk_model_train_config = model_trainer.get_updated_config( + updated_model_config, updated_model_train_config = model_trainer.get_updated_config( ) - return i, best_model + return i, best_model, updated_model_config, updated_model_train_config parallel = Parallel(n_jobs=self.num_workers) - models = parallel( + results = parallel( delayed(train_chunked_model)(i, start) for i, (start) in enumerate( range(0, self.total_features, self.feature_subsetsize))) # parallel loop returns all models with the chunk number, which is used to sort models in order - # model[1] returns only the model, without the chunk number - models = sorted(models) - models = [model[1] for model in models] + results = sorted(results, key=lambda result: result[0]) + + # With a process-based backend (num_workers > 1), mutating `self` + # inside `train_chunked_model` above does not propagate back to this + # process, so the resolved configs must come from each worker's + # return value instead; every chunk resolves the same model/train + # config class, so the first result is representative. + if results: + self.chunk_model_config = results[0][2] + self.chunk_model_train_config = results[0][3] + + models = [result[1] for result in results] return models def get_updated_configs(self): diff --git a/scalr/feature/scoring/__init__.py b/scalr/feature/scoring/__init__.py index 6d6cee9..5b2fcf1 100644 --- a/scalr/feature/scoring/__init__.py +++ b/scalr/feature/scoring/__init__.py @@ -1,4 +1,29 @@ +"""`ShapScorer` pulls in the `shap` package, which is only needed when SHAP- +based scoring is actually configured; it's imported lazily on first attribute +access (PEP 562) so importing `scalr` doesn't pay that cost for unrelated +commands. `getattr(module, name)` (the `name`/`params` builder pattern used +by `build_scorer`) transparently triggers this, same as an eager import would. +""" + +import importlib + from ._scoring import build_scorer from ._scoring import ScoringBase from .linear_scorer import LinearScorer -from .shap_scorer import ShapScorer \ No newline at end of file + +_LAZY_SUBMODULES = { + 'ShapScorer': '.shap_scorer', +} + + +def __getattr__(name): + module_name = _LAZY_SUBMODULES.get(name) + if module_name is None: + raise AttributeError(f'module {__name__!r} has no attribute {name!r}') + value = getattr(importlib.import_module(module_name, __name__), name) + globals()[name] = value + return value + + +def __dir__(): + return sorted(set(globals()) | set(_LAZY_SUBMODULES)) diff --git a/scalr/feature/test_feature_subsetting.py b/scalr/feature/test_feature_subsetting.py new file mode 100644 index 0000000..d926213 --- /dev/null +++ b/scalr/feature/test_feature_subsetting.py @@ -0,0 +1,90 @@ +"""This is a test file for feature_subsetting.py""" + +from scalr.feature.feature_subsetting import FeatureSubsetting +from scalr.utils import generate_dummy_dge_anndata + +_MODEL_CONFIG = {'name': 'SequentialModel', 'params': {'layers': [4, 2]}} +_MODEL_TRAIN_CONFIG = { + 'trainer': 'SimpleModelTrainer', + 'dataloader': { + 'name': 'SimpleDataLoader', + 'params': { + 'batch_size': 8 + } + }, + 'optimizer': { + 'name': 'SGD', + 'params': { + 'lr': 0.01 + } + }, + 'loss': { + 'name': 'CrossEntropyLoss' + }, + 'epochs': 1, +} + + +def _toy_adata(): + return generate_dummy_dge_anndata(n_donors=4, + cell_type_list=['B_cell', 'T_cell'], + cell_replicate=5, + n_vars=4) + + +def _resolved_model_params(tmp_path, num_workers): + adata = _toy_adata() + mappings = { + 'cell_type': { + 'id2label': { + 0: 'B_cell', + 1: 'T_cell' + }, + 'label2id': { + 'B_cell': 0, + 'T_cell': 1 + }, + } + } + + trainer = FeatureSubsetting( + feature_subsetsize=4, + chunk_model_config=_MODEL_CONFIG, + chunk_model_train_config=_MODEL_TRAIN_CONFIG, + train_data=adata, + val_data=adata, + target='cell_type', + mappings=mappings, + dirpath=str(tmp_path), + device='cpu', + num_workers=num_workers, + sample_chunksize=100, + ) + if num_workers > 1: + trainer.write_feature_subsetted_data() + trainer.train_chunked_models() + model_config, _ = trainer.get_updated_configs() + return model_config['params'] + + +def test_train_chunked_models_resolves_default_params_single_worker(tmp_path): + """`get_updated_configs` should reflect resolved (default-filled) params + sourced from each worker's returned config, not just echo the unresolved + input config.""" + params = _resolved_model_params(tmp_path, num_workers=1) + + assert params['dropout'] == 0 + assert params['activation'] == 'ReLU' + assert params['weights_init_zero'] is False + + +def test_train_chunked_models_resolves_default_params_multiple_workers( + tmp_path): + """With num_workers>1 (process-based joblib backend), resolved params from a + worker must still propagate back to the parent process's config, not just + echo the unresolved input config.""" + params = _resolved_model_params(tmp_path, num_workers=2) + + assert params['dropout'] == 0 + assert params['activation'] == 'ReLU' + assert params['weights_init_zero'] is False diff --git a/scalr/feature_extraction_pipeline.py b/scalr/feature_extraction_pipeline.py index 4920731..4655c10 100644 --- a/scalr/feature_extraction_pipeline.py +++ b/scalr/feature_extraction_pipeline.py @@ -7,6 +7,8 @@ from anndata import AnnData from anndata.experimental import AnnCollection +from joblib import delayed +from joblib import Parallel import numpy as np import pandas as pd from torch import nn @@ -77,6 +79,7 @@ def set_data_and_targets(self, self.val_data = val_data self.target = target self.mappings = mappings + self.sample_chunksize = sample_chunksize def feature_subsetted_model_training(self) -> list[nn.Module]: """This function train models on subsetted data containing `feature_subsetsize` genes.""" @@ -124,12 +127,11 @@ def feature_scoring(self) -> pd.DataFrame: deepcopy(self.feature_selection_config.get('scoring_config'))) self.feature_selection_config['scoring_config'] = scorer_config - all_scores = [] if not getattr(self, 'feature_subsetsize', None): self.feature_subsetsize = self.train_data.shape[1] + num_workers = getattr(self, 'num_workers', None) or 1 - # TODO: Parallelize feature scoring - for i, (model) in enumerate(self.chunked_models): + def score_chunk(i: int, model: nn.Module) -> tuple[int, np.ndarray]: subset_train_data = self.train_data[:, i * self.feature_subsetsize:(i + 1) * @@ -140,8 +142,17 @@ def feature_scoring(self) -> pd.DataFrame: score = scorer.generate_scores(model, subset_train_data, subset_val_data, self.target, self.mappings) - - all_scores.append(score[:self.feature_subsetsize]) + return i, score[:self.feature_subsetsize] + + # Uses the threading backend (not process-based) so per-chunk scoring + # shares memory with the parent process: no pickling of the (often + # GPU-resident) models/data/loggers across processes, and numpy/torch + # release the GIL during their heavy lifting, so this still + # parallelizes real work rather than just adding process overhead. + results = Parallel(n_jobs=num_workers, backend='threading')( + delayed(score_chunk)(i, model) + for i, model in enumerate(self.chunked_models)) + all_scores = [score for _, score in sorted(results)] columns = self.train_data.var_names columns.name = "index" diff --git a/scalr/feature_stability.py b/scalr/feature_stability.py new file mode 100644 index 0000000..860487e --- /dev/null +++ b/scalr/feature_stability.py @@ -0,0 +1,116 @@ +"""This file implements feature-selection stability analysis. + +Feature selection (biomarker discovery) is only scientifically defensible if +the selected genes are stable across independent resamplings of the data, +not an artifact of one particular train/test split. This module repeats +feature selection on random subsamples and reports how consistently each +gene is selected, and how similar the selected sets are pairwise (Jaccard). +""" + +from dataclasses import dataclass +from dataclasses import field + +from anndata import AnnData +import numpy as np +import pandas as pd +from sklearn.linear_model import LogisticRegression + + +@dataclass +class FeatureStabilityReport: + """Stability of a repeated top-k feature-selection procedure.""" + + selection_frequency: pd.Series + n_runs: int + top_k: int + mean_jaccard: float + + def top_stable_features(self, min_runs: int = None) -> list[str]: + """Genes selected in at least `min_runs` of the `n_runs` runs (default: half).""" + min_runs = min_runs if min_runs is not None else max( + self.n_runs // 2, 1) + return self.selection_frequency[self.selection_frequency >= + min_runs].index.tolist() + + def __str__(self) -> str: + lines = [ + f'Feature-selection stability over {self.n_runs} runs (top-{self.top_k}):', + '', + f'{"Gene":<20}Selected runs / {self.n_runs}', + ] + for gene, count in self.selection_frequency.head(20).items(): + lines.append(f'{gene:<20}{count}') + lines.append('') + lines.append( + f'Mean pairwise Jaccard similarity: {self.mean_jaccard:.3f}') + return '\n'.join(lines) + + +def compute_feature_stability( + adata: AnnData, + labels_key: str, + top_k: int = 50, + n_runs: int = 20, + subsample_frac: float = 0.8, + seed: int = 42, +) -> FeatureStabilityReport: + """Run repeated feature selection on random subsamples and measure stability. + + Each run fits a multinomial logistic regression on a random + `subsample_frac` fraction of cells (sampled without replacement) and + scores genes by the sum of `|coefficient|` across classes; the top-`k` + genes by that score form the run's selected set. + + Args: + adata: AnnData with normalized `adata.X` and label column `labels_key`. + labels_key: Column in `adata.obs` with classification labels. + top_k: Number of top-scoring genes selected per run. + n_runs: Number of independent subsampling runs. + subsample_frac: Fraction of cells sampled (without replacement) per run. + seed: Base random seed; run `i` uses seed `seed + i`. + + Returns: + A `FeatureStabilityReport` with per-gene selection frequency and the + mean pairwise Jaccard similarity of the selected sets across runs. + """ + genes = np.asarray(adata.var_names) + n_cells = adata.shape[0] + labels = adata.obs[labels_key].astype(str).values + + X = adata.X + if not isinstance(X, np.ndarray): + X = X.toarray() + + selected_sets = [] + gene_counts = pd.Series(0, index=genes) + + sample_size = max(int(n_cells * subsample_frac), 2) + for run in range(n_runs): + run_rng = np.random.default_rng(seed + run) + idx = run_rng.choice(n_cells, size=sample_size, replace=False) + + clf = LogisticRegression(max_iter=200) + clf.fit(X[idx], labels[idx]) + + scores = np.abs(clf.coef_).sum(axis=0) + top_idx = np.argsort(-scores)[:top_k] + top_genes = set(genes[top_idx]) + selected_sets.append(top_genes) + gene_counts[list(top_genes)] += 1 + + gene_counts = gene_counts[gene_counts > 0].sort_values(ascending=False) + + jaccards = [] + for i in range(len(selected_sets)): + for j in range(i + 1, len(selected_sets)): + union = len(selected_sets[i] | selected_sets[j]) + intersection = len(selected_sets[i] & selected_sets[j]) + jaccards.append(intersection / union if union else 0.0) + mean_jaccard = float(np.mean(jaccards)) if jaccards else 1.0 + + return FeatureStabilityReport( + selection_frequency=gene_counts, + n_runs=n_runs, + top_k=top_k, + mean_jaccard=mean_jaccard, + ) diff --git a/scalr/genes.py b/scalr/genes.py new file mode 100644 index 0000000..61c5824 --- /dev/null +++ b/scalr/genes.py @@ -0,0 +1,109 @@ +"""This file implements gene identifier alignment between a reference model and query data.""" + +from dataclasses import dataclass +from dataclasses import field +from typing import Union + +from anndata import AnnData +import numpy as np +from scipy import sparse + + +@dataclass +class GeneAlignmentReport: + """Summary of aligning a query AnnData's genes onto a reference feature list.""" + + reference_features: int + query_genes: int + matched_features: int + missing_features: list[str] + extra_genes: int + + @property + def coverage(self) -> float: + return (self.matched_features / + self.reference_features) if self.reference_features else 0.0 + + @property + def is_usable(self) -> bool: + return self.coverage >= self._min_overlap + + def __str__(self) -> str: + lines = [ + f'Reference features: {self.reference_features:,}', + f'Query genes: {self.query_genes:,}', + f'Matched features: {self.matched_features:,} ' + f'({self.coverage * 100:.1f}%)', + f'Missing features: {len(self.missing_features):,}', + f'Extra query genes: {self.extra_genes:,}', + '', + f'Status: {"usable" if self.is_usable else "unusable"}' + f'{" with warning" if self.is_usable and self.missing_features else ""}', + ] + return '\n'.join(lines) + + +def align_genes( + query: AnnData, + reference_features: list[str], + min_feature_overlap: float = 0.1, +) -> tuple[AnnData, GeneAlignmentReport]: + """Align a query AnnData's genes onto an ordered reference feature list. + + Genes present in `reference_features` but missing from the query are + added as zero-filled columns. Genes in the query not present in + `reference_features` are dropped. The result has columns in exactly + `reference_features` order, so it can be fed directly to a model trained + on that feature list. + + Args: + query: Query AnnData object. + reference_features: Ordered list of gene names the reference model expects. + min_feature_overlap: Minimum required fraction of `reference_features` + present in the query for the result to be marked usable. + + Returns: + (aligned_adata, report): AnnData with `reference_features` columns (in + order), and a `GeneAlignmentReport` describing the alignment. + """ + query_genes = list(query.var_names) + # Keep the first occurrence of a duplicated gene name (a plain dict + # comprehension would keep the last), matching the "only the first + # occurrence will be used" contract from scalr.validation. + query_index = {} + for i, g in enumerate(query_genes): + query_index.setdefault(g, i) + + matched = [g for g in reference_features if g in query_index] + missing = [g for g in reference_features if g not in query_index] + + n_obs = query.shape[0] + n_ref = len(reference_features) + + col_map = {g: i for i, g in enumerate(reference_features)} + src_cols = [query_index[g] for g in matched] + dst_cols = [col_map[g] for g in matched] + + X_query = query.X + if not sparse.issparse(X_query): + X_query = sparse.csr_matrix(X_query) + X_query = X_query.tocsc() + + aligned = sparse.lil_matrix((n_obs, n_ref), dtype=X_query.dtype) + if src_cols: + aligned[:, dst_cols] = X_query[:, src_cols] + aligned = aligned.tocsr() + + aligned_adata = AnnData(X=aligned, obs=query.obs.copy()) + aligned_adata.var_names = reference_features + + report = GeneAlignmentReport( + reference_features=n_ref, + query_genes=len(query_genes), + matched_features=len(matched), + missing_features=missing, + extra_genes=len(query_genes) - len(matched), + ) + report._min_overlap = min_feature_overlap + + return aligned_adata, report diff --git a/scalr/hierarchy.py b/scalr/hierarchy.py new file mode 100644 index 0000000..7334e58 --- /dev/null +++ b/scalr/hierarchy.py @@ -0,0 +1,84 @@ +"""This file implements hierarchical (coarse-to-fine) cell-type annotation. + +Rather than training a separate broad-level model, a fine-grained +`PredictionResult` is aggregated to a broad level by summing calibrated +fine-class probabilities within each broad group, using a user-supplied +`taxonomy` mapping (fine label -> broad label). This lets a single trained +model represent both "this is definitely Immune" (broad, high confidence) +and "which lymphoid subtype is less certain" (fine, lower confidence). +""" + +import numpy as np + +from scalr.calibration import predictive_entropy +from scalr.calibration import top1_top2_margin +from scalr.result import PredictionResult + + +def validate_taxonomy(class_names: list[str], taxonomy: dict[str, str]) -> None: + """Raise if `taxonomy` does not cover every class in `class_names`.""" + missing = [c for c in class_names if c not in taxonomy] + if missing: + raise ValueError( + f'taxonomy is missing an entry for fine-grained classes: {missing}') + + +def aggregate_to_broad(result: PredictionResult, + taxonomy: dict[str, str]) -> PredictionResult: + """Aggregate a fine-grained `PredictionResult` to broad-level labels. + + Args: + result: A fine-grained prediction result (e.g. CD4 T cell, CD8 T cell, ...). + taxonomy: Mapping from every fine class name in `result.class_names` to + its broad class name (e.g. {'CD4 T cell': 'T cell', ...}). + + Returns: + A new `PredictionResult` at the broad level. Broad-class probabilities + are the sum of the fine-class probabilities within each broad group, + and confidence/entropy/margin are recomputed from them. `is_unknown` + is carried over unchanged from the fine-level result (not + recomputed): a fine-level abstention is still a lack of broad-level + evidence, so it stays flagged unknown at the broad level too. + """ + validate_taxonomy(result.class_names, taxonomy) + + broad_names = sorted(set(taxonomy.values())) + broad_index = {b: i for i, b in enumerate(broad_names)} + fine_to_broad_idx = np.array( + [broad_index[taxonomy[c]] for c in result.class_names]) + + n_cells = result.probabilities.shape[0] + broad_probs = np.zeros((n_cells, len(broad_names))) + for fine_idx, broad_idx in enumerate(fine_to_broad_idx): + broad_probs[:, broad_idx] += result.probabilities[:, fine_idx] + + top1 = broad_probs.argmax(axis=1) + labels = [broad_names[i] for i in top1] + confidence = broad_probs.max(axis=1) + entropy = predictive_entropy(broad_probs) + margin = top1_top2_margin(broad_probs) + + k = min(5, len(broad_names)) + top_k_idx = np.argsort(-broad_probs, axis=1)[:, :k] + top_k_out = [[(broad_names[c], float(broad_probs[row, c])) + for c in idx] + for row, idx in enumerate(top_k_idx)] + + metadata = dict(result.metadata) + metadata['level'] = 'broad' + metadata['fine_class_names'] = result.class_names + + return PredictionResult( + obs_names=result.obs_names, + labels=labels, + probabilities=broad_probs, + class_names=broad_names, + confidence=confidence, + entropy=entropy, + margin=margin, + # Cells flagged unknown at the fine level remain unknown at the broad + # level; a fine-level abstention is still a lack of broad-level evidence. + is_unknown=result.is_unknown.copy(), + top_k=top_k_out, + metadata=metadata, + ) diff --git a/scalr/leakage.py b/scalr/leakage.py new file mode 100644 index 0000000..c6ab9f5 --- /dev/null +++ b/scalr/leakage.py @@ -0,0 +1,138 @@ +"""This file implements leakage-safe train/val/test splitting and leakage detection. + +Splitting cells at random can put the same donor/patient/sample in both train +and test, letting a model "cheat" on donor-specific signal rather than +learning the classification task. This module makes group-safe splitting a +first-class, explicit operation. +""" + +from dataclasses import dataclass +from dataclasses import field + +import numpy as np +import pandas as pd +from sklearn.model_selection import GroupShuffleSplit +from sklearn.model_selection import StratifiedShuffleSplit + + +@dataclass +class LeakageReport: + """Result of checking for group leakage across data splits.""" + + group_key: str + leaked_groups: dict = field(default_factory=dict) + + @property + def has_leakage(self) -> bool: + return len(self.leaked_groups) > 0 + + def __str__(self) -> str: + if not self.has_leakage: + return f'No leakage detected across splits on "{self.group_key}".' + lines = [ + f'Leakage detected: the following "{self.group_key}" values ' + 'appear in more than one split:' + ] + for split_pair, groups in self.leaked_groups.items(): + lines.append(f' {split_pair}: {sorted(groups)[:10]}' + f'{"..." if len(groups) > 10 else ""}') + lines.append( + 'Recommendation: use group-safe splitting (split_strategy="group") ' + f'on column "{self.group_key}".') + return '\n'.join(lines) + + +def check_group_leakage(splits: dict[str, pd.DataFrame], + group_key: str) -> LeakageReport: + """Check whether any group value spans more than one split. + + Args: + splits: Mapping of split name (e.g. 'train', 'val', 'test') to that + split's `obs` DataFrame. + group_key: Column name identifying donor/patient/sample/batch. + + Returns: + A `LeakageReport`. + """ + split_groups = { + name: set(df[group_key].unique()) + for name, df in splits.items() + if group_key in df.columns + } + names = list(split_groups) + leaked = {} + for i in range(len(names)): + for j in range(i + 1, len(names)): + overlap = split_groups[names[i]] & split_groups[names[j]] + if overlap: + leaked[f'{names[i]} <-> {names[j]}'] = overlap + + return LeakageReport(group_key=group_key, leaked_groups=leaked) + + +def detect_likely_grouping_column( + obs: pd.DataFrame, + candidates: tuple[str, ...] = ('donor_id', 'donor', 'patient_id', 'patient', + 'sample_id', 'sample', 'subject_id', 'batch') +) -> list[str]: + """Detect columns in `obs` that look like a donor/patient/sample identifier. + + This is a name-based heuristic used to warn users who did not specify a + `group_key`, not a guarantee that grouped splitting is required. + """ + return [c for c in candidates if c in obs.columns] + + +def group_safe_split( + obs: pd.DataFrame, + labels_key: str, + group_key: str = None, + split_ratio: tuple[float, float, float] = (0.7, 0.1, 0.2), + seed: int = 42, +) -> dict[str, np.ndarray]: + """Split sample indices into train/val/test, keeping each group in a single split. + + Falls back to a stratified (non-grouped) split when `group_key` is None, + but this should only be done when no donor/patient/sample identifier is + available, since it risks train/test leakage. + + Args: + obs: `adata.obs` DataFrame. + labels_key: Column used for stratification. + group_key: Column identifying donor/patient/sample; when given, a + group never spans more than one split. + split_ratio: (train, val, test) ratios; need not sum to 1. + seed: Random seed for reproducibility. + + Returns: + Dict with 'train', 'val', 'test' -> integer position arrays into `obs`. + """ + n = len(obs) + indices = np.arange(n) + total = sum(split_ratio) + train_r, val_r, test_r = (r / total for r in split_ratio) + + def _split(idx, y, groups, test_size): + if groups is not None: + splitter = GroupShuffleSplit(test_size=test_size, + n_splits=1, + random_state=seed) + return next(splitter.split(idx, y[idx], groups=groups[idx])) + splitter = StratifiedShuffleSplit(test_size=test_size, + n_splits=1, + random_state=seed) + return next(splitter.split(idx, y[idx])) + + y = obs[labels_key].values + groups = obs[group_key].values if group_key else None + + train_val_pos, test_pos = _split(indices, y, groups, test_ratio := test_r) + train_val_idx = indices[train_val_pos] + test_idx = indices[test_pos] + + relative_val_ratio = val_r / (train_r + val_r) + train_pos, val_pos = _split(train_val_idx, y, groups, relative_val_ratio) + train_idx = train_val_idx[train_pos] + val_idx = train_val_idx[val_pos] + + return {'train': train_idx, 'val': val_idx, 'test': test_idx} diff --git a/scalr/metrics.py b/scalr/metrics.py new file mode 100644 index 0000000..51a7fbe --- /dev/null +++ b/scalr/metrics.py @@ -0,0 +1,149 @@ +"""This file implements class-imbalance-aware evaluation metrics.""" + +from dataclasses import dataclass +from dataclasses import field + +import numpy as np +import pandas as pd +from sklearn.metrics import balanced_accuracy_score +from sklearn.metrics import confusion_matrix +from sklearn.metrics import f1_score +from sklearn.metrics import precision_recall_fscore_support + + +@dataclass +class ClassificationMetrics: + """Class-imbalance-aware classification metrics. + + Overall accuracy is not reported as the headline metric because it is + dominated by majority classes; macro-F1 and balanced accuracy weight + every class equally so rare-cell performance is visible. + """ + + macro_f1: float + weighted_f1: float + balanced_accuracy: float + per_class: pd.DataFrame + confusion: np.ndarray + class_names: list[str] + + def rare_class_recall(self, min_support: int) -> pd.DataFrame: + """Recall restricted to classes with support < `min_support` in the evaluation set.""" + return self.per_class[self.per_class['support'] < min_support] + + def __str__(self) -> str: + lines = [ + f'macro-F1: {self.macro_f1:.4f}', + f'weighted-F1: {self.weighted_f1:.4f}', + f'balanced accuracy: {self.balanced_accuracy:.4f}', + ] + return '\n'.join(lines) + + +def compute_classification_metrics( + labels: np.ndarray, + predictions: np.ndarray, + class_names: list[str], +) -> ClassificationMetrics: + """Compute macro/weighted-F1, balanced accuracy, and a per-class report. + + Args: + labels: True integer class labels, shape [n_samples]. + predictions: Predicted integer class labels, shape [n_samples]. + class_names: Ordered class names corresponding to label ids [0..n_classes). + + Returns: + A `ClassificationMetrics` object. + """ + labels = np.asarray(labels) + predictions = np.asarray(predictions) + all_ids = list(range(len(class_names))) + + precision, recall, f1, support = precision_recall_fscore_support( + labels, predictions, labels=all_ids, zero_division=0) + + per_class = pd.DataFrame({ + 'class': class_names, + 'precision': precision, + 'recall': recall, + 'f1': f1, + 'support': support, + }).set_index('class') + + macro_f1 = f1_score(labels, + predictions, + labels=all_ids, + average='macro', + zero_division=0) + weighted_f1 = f1_score(labels, + predictions, + labels=all_ids, + average='weighted', + zero_division=0) + balanced_acc = balanced_accuracy_score(labels, predictions) + conf = confusion_matrix(labels, predictions, labels=all_ids) + + return ClassificationMetrics( + macro_f1=float(macro_f1), + weighted_f1=float(weighted_f1), + balanced_accuracy=float(balanced_acc), + per_class=per_class, + confusion=conf, + class_names=class_names, + ) + + +@dataclass +class ClassImbalanceReport: + """Report describing class-size imbalance in a labeled dataset.""" + + class_counts: pd.Series + imbalance_ratio: float + small_classes: list[str] = field(default_factory=list) + + def __str__(self) -> str: + lines = [ + 'Class imbalance detected' + if self.small_classes else 'Class distribution' + ] + lines.append('') + for cls, count in self.class_counts.items(): + lines.append(f'{cls}: {count:,}') + lines.append('') + if self.small_classes: + lines.append( + 'Recommendation: enable class-balanced sampling or class-weighted loss.' + ) + return '\n'.join(lines) + + +def check_class_imbalance( + labels: pd.Series, + min_class_size: int = 50, + imbalance_ratio_threshold: float = 20.0) -> ClassImbalanceReport: + """Report class sizes and flag rare classes / severe imbalance. + + Args: + labels: Categorical label series (e.g. `adata.obs[labels_key]`). + min_class_size: Classes with fewer samples than this are flagged as small. + imbalance_ratio_threshold: max_class_size / min_class_size above which the + dataset is considered severely imbalanced. + + Returns: + A `ClassImbalanceReport`. + """ + counts = labels.value_counts() + ratio = float(counts.max() / max(counts.min(), 1)) + small = counts[counts < min_class_size].index.tolist() + + return ClassImbalanceReport(class_counts=counts, + imbalance_ratio=ratio, + small_classes=small) + + +def compute_class_weights(labels: np.ndarray, n_classes: int) -> np.ndarray: + """Inverse-frequency class weights for a weighted cross-entropy loss.""" + counts = np.bincount(labels, minlength=n_classes).astype(np.float64) + counts[counts == 0] = 1.0 + weights = counts.sum() / (n_classes * counts) + return weights diff --git a/scalr/model_card.py b/scalr/model_card.py new file mode 100644 index 0000000..1f323d2 --- /dev/null +++ b/scalr/model_card.py @@ -0,0 +1,92 @@ +"""This file implements model card generation for provenance and reproducibility. + +Every self-contained model artifact should document what it was trained on, +how it performs, and its known limitations, separately from its raw +`manifest.json` metadata — so a person deciding whether to trust/reuse the +model does not have to reverse-engineer the training run. +""" + +_TEMPLATE = """# Model Card: {model_name} + +## Summary +- **Task**: {task} +- **Model version**: {model_version} +- **scaLR version**: {scalr_version} +- **Species**: {species} +- **Random seed**: {random_seed} + +## Training data +- **Labels column**: {labels_key} +- **Grouping column (leakage-safe split)**: {group_key} +- **Split ratio (train/val/test)**: {split_ratio} +- **Feature count**: {feature_count} +- **Classes ({n_classes})**: {class_list} + +## Performance (held-out test split) +{metrics_section} + +## Confidence & abstention +- Calibration: temperature scaling (T={temperature:.3f}) +- Open-set thresholds: {open_set_thresholds} + +## Known limitations +- Evaluated only on the training-distribution data described above; performance on + out-of-distribution donors, tissues, or sequencing technologies is not validated. +- Open-set (`is_unknown`) thresholds are fit on a single validation split and are a + heuristic screening signal, not a guarantee of correct abstention. +- Gene/feature coverage should be checked (`scalr.validate`) on every new query dataset. + +## License +- Source code license: see repository LICENSE. +- This model artifact's weights/data license: not specified by this template; + set explicitly by the model's author/distributor before sharing. +""" + + +def render_model_card(metadata: dict, metrics: dict) -> str: + """Render a Markdown model card from an `AnnotationModel`'s manifest metadata and metrics. + + Args: + metadata: The model's manifest dict (as written by `AnnotationModel.save`), + containing keys such as `labels`, `model_version`, `labels_key`, etc. + metrics: The model's recorded evaluation metrics dict (may be empty). + + Returns: + Markdown text suitable for writing to the artifact's `README.md`. + """ + class_names = metadata.get('labels', []) + + metrics_section = 'No held-out metrics recorded.' + if metrics: + lines = [] + if 'macro_f1' in metrics: + lines.append(f"- macro-F1: {metrics['macro_f1']:.4f}") + if 'weighted_f1' in metrics: + lines.append(f"- weighted-F1: {metrics['weighted_f1']:.4f}") + if 'balanced_accuracy' in metrics: + lines.append( + f"- balanced accuracy: {metrics['balanced_accuracy']:.4f}") + if 'n_test_cells' in metrics: + lines.append(f"- test cells: {metrics['n_test_cells']:,}") + if lines: + metrics_section = '\n'.join(lines) + + return _TEMPLATE.format( + model_name=metadata.get('model_name') or metadata.get('labels_key') or + 'unnamed model', + task=metadata.get('task', 'cell_type_annotation'), + model_version=metadata.get('model_version', 'unknown'), + scalr_version=metadata.get('scalr_version', 'unknown'), + species=metadata.get('species') or 'not specified', + random_seed=metadata.get('random_seed', 'unknown'), + labels_key=metadata.get('labels_key', 'unknown'), + group_key=metadata.get('group_key') or + 'none (no leakage-safe grouping used)', + split_ratio=metadata.get('split_ratio', 'unknown'), + feature_count=metadata.get('feature_count', 'unknown'), + n_classes=len(class_names), + class_list=', '.join(class_names) if class_names else 'unknown', + metrics_section=metrics_section, + temperature=metadata.get('temperature', 1.0), + open_set_thresholds=metadata.get('open_set_thresholds', {}), + ) diff --git a/scalr/models.py b/scalr/models.py new file mode 100644 index 0000000..3484175 --- /dev/null +++ b/scalr/models.py @@ -0,0 +1,218 @@ +"""This file implements a local model registry, plus push/pull of model +artifacts to the Hugging Face Hub, for the scaLR 2.0 "model hub". + +Hub access (`download`/`push`) requires the optional `huggingface_hub` +dependency (`pip install pyscaLR[hub]`) and is only imported lazily, so the +base install stays lightweight. Hub-hosted models are addressed by their HF +`repo_id` (`"/"`); local registry names are plain strings +with no `/`. +""" + +import os +from os import path +import shutil +from typing import Optional + +from scalr.utils import read_data + +DEFAULT_REGISTRY_DIR = path.expanduser('~/.scalr/models') + + +def _require_hub(): + try: + import huggingface_hub + except ImportError as e: + raise ImportError( + 'Hugging Face Hub access requires the optional "hub" extra. ' + 'Install it with: pip install "pyscaLR[hub]"') from e + return huggingface_hub + + +def _registry_dir(registry_dir: Optional[str] = None) -> str: + return registry_dir or os.environ.get('SCALR_MODEL_REGISTRY', + DEFAULT_REGISTRY_DIR) + + +def _resolve(name: str, registry_dir: Optional[str] = None) -> str: + target = path.join(_registry_dir(registry_dir), name) + pointer_file = path.join(target, '_pointer.txt') + if path.exists(pointer_file): + with open(pointer_file) as fh: + return fh.read().strip() + return target + + +def register( + artifact_dir: str, + name: str, + registry_dir: Optional[str] = None, + copy: bool = True, +) -> str: + """Register a saved model artifact directory under `name` in the local registry. + + Args: + artifact_dir: Path to a model artifact directory, as written by + `AnnotationModel.save`. + name: Name to register the model under (used by `list`/`info`/`load_model`). + registry_dir: Registry root directory. Defaults to the + `SCALR_MODEL_REGISTRY` environment variable, or `~/.scalr/models`. + copy: When True (default), copies the artifact into the registry. + When False, records a pointer to `artifact_dir` instead of + duplicating it on disk. + + Returns: + The path the model is registered/resolvable at. + """ + if not path.exists(path.join(artifact_dir, 'manifest.json')): + raise ValueError( + f'"{artifact_dir}" does not look like a scaLR model artifact ' + '(no manifest.json found).') + + reg_dir = _registry_dir(registry_dir) + os.makedirs(reg_dir, exist_ok=True) + target = path.join(reg_dir, name) + + if copy: + if path.exists(target): + shutil.rmtree(target) + shutil.copytree(artifact_dir, target) + else: + os.makedirs(target, exist_ok=True) + with open(path.join(target, '_pointer.txt'), 'w') as fh: + fh.write(path.abspath(artifact_dir)) + + return target + + +def _is_model_dir(model_dir: str) -> bool: + return path.exists(path.join(model_dir, 'manifest.json')) or path.exists( + path.join(model_dir, '_pointer.txt')) + + +def list_models(registry_dir: Optional[str] = None) -> list[str]: + """List model names registered in the local registry. + + Hub-style names (`"/"`, e.g. from `download()`) are + registered one directory level deeper than plain names, and are listed + the same way, e.g. `"org/some-model"`. + """ + reg_dir = _registry_dir(registry_dir) + if not path.isdir(reg_dir): + return [] + names = [] + for top in sorted(os.listdir(reg_dir)): + top_dir = path.join(reg_dir, top) + if not path.isdir(top_dir): + continue + if _is_model_dir(top_dir): + names.append(top) + continue + for sub in sorted(os.listdir(top_dir)): + sub_dir = path.join(top_dir, sub) + if path.isdir(sub_dir) and _is_model_dir(sub_dir): + names.append(f'{top}/{sub}') + return names + + +def info(name: str, registry_dir: Optional[str] = None) -> dict: + """Return a registered model's manifest dict.""" + return read_data( + path.join(resolve_path(name, registry_dir), 'manifest.json')) + + +def download( + repo_id: str, + registry_dir: Optional[str] = None, + revision: Optional[str] = None, + token: Optional[str] = None, + name: Optional[str] = None, +) -> str: + """Download a scaLR model artifact from the Hugging Face Hub and + register it locally. + + Args: + repo_id: Hugging Face repo id to pull, e.g. `"infocusp/pbmc-v1"`. + registry_dir: Registry root directory. Defaults to the + `SCALR_MODEL_REGISTRY` environment variable, or `~/.scalr/models`. + revision: Hub revision (branch/tag/commit) to pull. Defaults to the + repo's default branch. + token: Hugging Face auth token, for private repos. Defaults to the + token `huggingface_hub` picks up from `HF_TOKEN`/`huggingface-cli + login`. + name: Local registry name to register the model under. Defaults to + `repo_id` itself, so `scalr.load_model(repo_id)` resolves it + without a separate `register()` call. + + Returns: + The local path the model is registered/resolvable at. + """ + hub = _require_hub() + + cache_dir = hub.snapshot_download(repo_id=repo_id, + revision=revision, + token=token, + repo_type='model') + + if not path.exists(path.join(cache_dir, 'manifest.json')): + raise ValueError( + f'"{repo_id}" does not look like a scaLR model artifact ' + '(no manifest.json found in the downloaded repo).') + + return register(cache_dir, name or repo_id, registry_dir, copy=True) + + +def push( + artifact_dir: str, + repo_id: str, + token: Optional[str] = None, + private: bool = False, + commit_message: str = 'Upload scaLR model artifact', +) -> str: + """Upload a saved model artifact directory to the Hugging Face Hub. + + Args: + artifact_dir: Path to a model artifact directory, as written by + `AnnotationModel.save`. + repo_id: Hugging Face repo id to publish to, e.g. + `"infocusp/pbmc-v1"`. Created if it doesn't already exist. + token: Hugging Face auth token. Defaults to the token + `huggingface_hub` picks up from `HF_TOKEN`/`huggingface-cli + login`. + private: Whether to create the repo as private, if it doesn't + already exist. + commit_message: Commit message for the upload. + + Returns: + The URL of the pushed model on the Hub. + """ + if not path.exists(path.join(artifact_dir, 'manifest.json')): + raise ValueError( + f'"{artifact_dir}" does not look like a scaLR model artifact ' + '(no manifest.json found).') + + hub = _require_hub() + + api = hub.HfApi(token=token) + api.create_repo(repo_id, repo_type='model', private=private, exist_ok=True) + api.upload_folder(repo_id=repo_id, + folder_path=artifact_dir, + token=token, + commit_message=commit_message) + + return f'https://huggingface.co/{repo_id}' + + +def resolve_path(name: str, registry_dir: Optional[str] = None) -> str: + """Resolve a registered model name to its artifact directory path.""" + model_dir = _resolve(name, registry_dir) + if not path.exists(path.join(model_dir, 'manifest.json')): + raise FileNotFoundError( + f'No model named "{name}" found in registry ' + f'"{_registry_dir(registry_dir)}". Registered models: ' + f'{list_models(registry_dir)}') + return model_dir + + +# `list` matches the scaLR 2.0 blueprint's `scalr.models.list()` API; the +# builtin `list` is not otherwise used in this module. +list = list_models diff --git a/scalr/nn/loss/_loss.py b/scalr/nn/loss/_loss.py index 8969587..6811553 100644 --- a/scalr/nn/loss/_loss.py +++ b/scalr/nn/loss/_loss.py @@ -8,10 +8,17 @@ from torch import nn import scalr +from scalr.utils import build_object class CustomLossBase(nn.Module): - """Base class to implement custom loss functions.""" + """Base class to implement custom loss functions. + + Subclass this, set `self.criterion` in `__init__`, and add the subclass + next to this file (or import it into `scalr/nn/loss/__init__.py`) so + `build_loss_fn` can find it by name via the `name`/`params` config + pattern, the same way as models/preprocessors/splitters/etc. + """ def __init__(self): super().__init__() @@ -21,15 +28,27 @@ def forward(self, out, preds): """Returns loss betwen outputs and predictions.""" return self.criterion(out, preds) + @classmethod + def get_default_params(cls) -> dict: + """Class method to get default params for loss_config.""" + return dict() + def build_loss_fn(loss_config): - """Builder object to get Loss function, updated loss_config.""" + """Builder object to get Loss function, updated loss_config. + + Resolves `name` against custom loss classes registered under + `scalr.nn.loss` (subclasses of `CustomLossBase`) first, falling back to + `torch.nn` for built-in losses (e.g. `CrossEntropyLoss`). + """ name = loss_config.get('name') if not name: raise ValueError('Loss function not provided') - params = loss_config.get('params', dict()) + custom_cls = getattr(scalr.nn.loss, name, None) + if isinstance(custom_cls, type) and issubclass(custom_cls, CustomLossBase): + return build_object(scalr.nn.loss, loss_config) - # TODO: Add provision for custom loss object + params = loss_config.get('params', dict()) loss_fn = getattr(torch.nn, name)(**params) return loss_fn, loss_config diff --git a/scalr/nn/loss/test_loss.py b/scalr/nn/loss/test_loss.py new file mode 100644 index 0000000..4b0f649 --- /dev/null +++ b/scalr/nn/loss/test_loss.py @@ -0,0 +1,93 @@ +"""This is a test file for _loss.py""" + +import pytest +import torch + +from scalr.nn.loss import build_loss_fn +from scalr.nn.loss import CustomLossBase +import scalr.nn.loss as loss_module + + +def test_build_loss_fn_resolves_builtin_torch_loss(): + """A name matching a torch.nn class should build that built-in loss.""" + loss_fn, loss_config = build_loss_fn({'name': 'CrossEntropyLoss'}) + + assert isinstance(loss_fn, torch.nn.CrossEntropyLoss) + assert loss_config == {'name': 'CrossEntropyLoss'} + + +def test_build_loss_fn_passes_params_to_builtin_loss(): + """Params should be forwarded to the built-in loss class's constructor.""" + loss_fn, _ = build_loss_fn({ + 'name': 'CrossEntropyLoss', + 'params': { + 'label_smoothing': 0.1 + }, + }) + + assert loss_fn.label_smoothing == 0.1 + + +def test_build_loss_fn_resolves_custom_loss_registered_under_scalr_nn_loss( + monkeypatch): + """A custom CustomLossBase subclass registered under scalr.nn.loss should be resolved by name.""" + + class _DummyWeightedLoss(CustomLossBase): + + def __init__(self, gamma=1.0): + super().__init__() + self.gamma = gamma + self.criterion = torch.nn.CrossEntropyLoss() + + @classmethod + def get_default_params(cls): + return {'gamma': 1.0} + + monkeypatch.setattr(loss_module, + '_DummyWeightedLoss', + _DummyWeightedLoss, + raising=False) + + loss_fn, loss_config = build_loss_fn({ + 'name': '_DummyWeightedLoss', + 'params': { + 'gamma': 2.5 + }, + }) + + assert isinstance(loss_fn, _DummyWeightedLoss) + assert loss_fn.gamma == 2.5 + assert loss_config['params'] == {'gamma': 2.5} + + out = loss_fn(torch.randn(4, 3), torch.tensor([0, 1, 2, 1])) + assert out.ndim == 0 # scalar loss + + +def test_build_loss_fn_custom_loss_fills_in_default_params(monkeypatch): + """Unspecified params for a custom loss should fall back to get_default_params().""" + + class _DummyDefaultedLoss(CustomLossBase): + + def __init__(self, gamma=1.0): + super().__init__() + self.gamma = gamma + self.criterion = torch.nn.CrossEntropyLoss() + + @classmethod + def get_default_params(cls): + return {'gamma': 1.0} + + monkeypatch.setattr(loss_module, + '_DummyDefaultedLoss', + _DummyDefaultedLoss, + raising=False) + + loss_fn, _ = build_loss_fn({'name': '_DummyDefaultedLoss'}) + + assert loss_fn.gamma == 1.0 + + +def test_build_loss_fn_raises_without_name(): + """A loss config without a `name` should raise, not silently no-op.""" + with pytest.raises(ValueError): + build_loss_fn({}) diff --git a/scalr/refinement.py b/scalr/refinement.py new file mode 100644 index 0000000..9bebe77 --- /dev/null +++ b/scalr/refinement.py @@ -0,0 +1,90 @@ +"""This file implements cluster-aware post-processing refinement of predictions. + +Per-cell classifier predictions can be noisy within a transcriptionally +homogeneous cluster. This module optionally relabels cells to their cluster's +confidence-weighted majority class, while always preserving the original, +unrefined predictions rather than silently overwriting them. +""" + +from dataclasses import replace +from typing import Optional + +from anndata import AnnData +import numpy as np +import pandas as pd + +from scalr.result import PredictionResult + +_DEFAULT_CLUSTER_KEYS = ('leiden', 'louvain', 'cluster', 'clusters') + + +def detect_cluster_key(adata: AnnData) -> Optional[str]: + """Return the first column in `adata.obs` that looks like a cluster assignment, if any.""" + for key in _DEFAULT_CLUSTER_KEYS: + if key in adata.obs.columns: + return key + return None + + +def refine_with_clusters( + result: PredictionResult, + adata: AnnData, + cluster_key: str = 'auto', + min_cluster_agreement: float = 0.5, +) -> PredictionResult: + """Refine predictions using per-cluster, confidence-weighted majority vote. + + For each cluster, if a class holds at least `min_cluster_agreement` of + that cluster's total confidence-weighted vote, every cell in the cluster + is relabeled to that class. This never mutates `result`; it returns a new + `PredictionResult` whose `labels` are the refined labels, with + `metadata['raw_labels']` preserving the original, unrefined predictions + so both are always available. + + Args: + result: A `PredictionResult` whose `obs_names` match `adata.obs_names`. + adata: AnnData carrying a cluster assignment column in `.obs`. + cluster_key: Column name in `adata.obs`, or 'auto' to detect one of + 'leiden'/'louvain'/'cluster'/'clusters'. + min_cluster_agreement: Minimum confidence-weighted vote share (0-1) + required for a cluster's majority class to be applied. + + Returns: + A new `PredictionResult` with refined `labels` and + `metadata['raw_labels']` holding the pre-refinement labels. + """ + if list(adata.obs_names) != list(result.obs_names): + raise ValueError( + 'adata.obs_names must match result.obs_names for cluster refinement.' + ) + + if cluster_key == 'auto': + detected = detect_cluster_key(adata) + if detected is None: + raise ValueError( + 'No cluster column found automatically (looked for ' + f'{_DEFAULT_CLUSTER_KEYS}); pass `cluster_key` explicitly.') + cluster_key = detected + elif cluster_key not in adata.obs.columns: + raise ValueError(f'cluster_key="{cluster_key}" not found in adata.obs.') + + clusters = adata.obs[cluster_key].astype(str).values + labels = np.array(result.labels) + confidence = np.asarray(result.confidence) + + refined = labels.copy() + for cluster_id in np.unique(clusters): + mask = clusters == cluster_id + weights = pd.Series(confidence[mask]).groupby(labels[mask]).sum() + total_weight = weights.sum() + if total_weight == 0: + continue + majority_label = weights.idxmax() + if weights[majority_label] / total_weight >= min_cluster_agreement: + refined[mask] = majority_label + + metadata = dict(result.metadata) + metadata['raw_labels'] = labels.tolist() + metadata['cluster_key'] = cluster_key + + return replace(result, labels=refined.tolist(), metadata=metadata) diff --git a/scalr/result.py b/scalr/result.py new file mode 100644 index 0000000..3d60590 --- /dev/null +++ b/scalr/result.py @@ -0,0 +1,87 @@ +"""This file defines the typed prediction result contract returned by the public API.""" + +from dataclasses import dataclass +from dataclasses import field +from typing import Optional + +import numpy as np +import pandas as pd + + +@dataclass +class PredictionResult: + """Typed container for annotation/prediction outputs. + + Row order and identity always match the `obs_names` of the AnnData object + that was passed to `predict`/`annotate`, so predictions cannot silently + become misaligned with the input data. + + Attributes: + obs_names: Cell identifiers, in the same order as all other arrays. + labels: Top-1 predicted label (string) for each cell. + probabilities: [n_cells, n_classes] calibrated class probabilities. + class_names: Ordered class names corresponding to `probabilities` columns. + confidence: Calibrated top-1 confidence for each cell. + entropy: Predictive (normalized) entropy for each cell. + margin: Top-1 minus top-2 probability margin for each cell. + top_k: List of (label, probability) tuples per cell, top-k only. + is_unknown: Whether the cell was flagged as an open-set/unknown prediction. + ood_score: Optional out-of-distribution score (higher = more OOD). + explanations: Optional per-cell explanation payload. + metadata: Model/preprocessing/run metadata used to produce this result. + is_possible_doublet: Optional screening flag for cells whose top-two + class probabilities are both substantial and close together — + a heuristic signal, not a validated doublet call. + """ + + obs_names: list[str] + labels: list[str] + probabilities: np.ndarray + class_names: list[str] + confidence: np.ndarray + entropy: np.ndarray + margin: np.ndarray + is_unknown: np.ndarray + top_k: list[list[tuple]] = field(default_factory=list) + ood_score: Optional[np.ndarray] = None + explanations: Optional[dict] = None + metadata: dict = field(default_factory=dict) + is_possible_doublet: Optional[np.ndarray] = None + + def __len__(self) -> int: + return len(self.labels) + + def to_frame(self) -> pd.DataFrame: + """Return a DataFrame indexed by `obs_names`, suitable for merging into `adata.obs`.""" + data = { + 'scalr_pred': self.labels, + 'scalr_confidence': self.confidence, + 'scalr_entropy': self.entropy, + 'scalr_margin': self.margin, + 'scalr_unknown': self.is_unknown, + } + if self.is_possible_doublet is not None: + data['scalr_possible_doublet'] = self.is_possible_doublet + return pd.DataFrame(data, index=self.obs_names) + + def write_to_adata(self, adata, prefix: str = 'scalr') -> None: + """Write prediction fields onto an AnnData object's `.obs`/`.obsm`/`.uns`. + + Args: + adata: AnnData object whose `obs_names` must match `self.obs_names`. + prefix: Prefix used for all written columns/keys. + """ + if list(adata.obs_names) != list(self.obs_names): + raise ValueError( + 'adata.obs_names does not match the prediction result obs_names; ' + 'refusing to write to avoid misaligned predictions.') + + adata.obs[f'{prefix}_pred'] = self.labels + adata.obs[f'{prefix}_confidence'] = self.confidence + adata.obs[f'{prefix}_entropy'] = self.entropy + adata.obs[f'{prefix}_margin'] = self.margin + adata.obs[f'{prefix}_unknown'] = self.is_unknown + if self.is_possible_doublet is not None: + adata.obs[f'{prefix}_possible_doublet'] = self.is_possible_doublet + adata.obsm[f'{prefix}_probabilities'] = self.probabilities + adata.uns[prefix] = self.metadata diff --git a/scalr/test_api.py b/scalr/test_api.py new file mode 100644 index 0000000..6da317c --- /dev/null +++ b/scalr/test_api.py @@ -0,0 +1,502 @@ +"""This is an end-to-end test file for the simple public API (api.py, artifact.py).""" + +import numpy as np +import pytest + +import scalr +from scalr.utils import generate_dummy_dge_anndata + + +def _toy_adata(): + return generate_dummy_dge_anndata(n_donors=8, + cell_type_list=['B_cell', 'T_cell', 'DC'], + cell_replicate=10, + n_vars=20) + + +def test_train_returns_annotation_model_with_metrics(): + """scalr.train should return a calibrated model with held-out metrics.""" + adata = _toy_adata() + + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + + assert set(model.class_names) == {'B_cell', 'T_cell', 'DC'} + assert model.features == list(adata.var_names) + assert 'macro_f1' in model.metrics + assert 0.0 <= model.metrics['macro_f1'] <= 1.0 + + +def test_annotate_result_is_aligned_and_typed(): + """scalr.annotate should return a PredictionResult aligned to obs_names.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + + result = scalr.annotate(adata, model=model, device='cpu') + + assert isinstance(result, scalr.PredictionResult) + assert result.obs_names == list(adata.obs_names) + assert len(result.labels) == len(adata) + assert result.probabilities.shape == (len(adata), 3) + assert all(l in model.class_names for l in result.labels) + assert result.is_unknown.dtype == bool + + +def test_model_save_and_load_roundtrip(tmp_path): + """A saved model artifact should reload and produce identical predictions.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + + artifact_dir = tmp_path / 'pbmc_v1' + model.save(str(artifact_dir)) + + assert (artifact_dir / 'manifest.json').exists() + assert (artifact_dir / 'model.pt').exists() + assert (artifact_dir / 'features.json').exists() + assert (artifact_dir / 'label_mapping.json').exists() + + reloaded = scalr.load_model(str(artifact_dir)) + assert reloaded.class_names == model.class_names + assert reloaded.features == model.features + + original_result = model.predict(adata, device='cpu') + reloaded_result = reloaded.predict(adata, device='cpu') + + assert np.allclose(original_result.probabilities, + reloaded_result.probabilities, + atol=1e-5) + assert original_result.labels == reloaded_result.labels + + +def test_annotate_from_artifact_path(tmp_path): + """scalr.annotate should accept a model artifact directory path directly.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + artifact_dir = tmp_path / 'pbmc_v1' + model.save(str(artifact_dir)) + + result = scalr.annotate(adata, model=str(artifact_dir), device='cpu') + assert len(result.labels) == len(adata) + + +def test_annotate_with_gene_subset_query(): + """annotate should align a query with fewer/reordered genes automatically.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + + # Query with a shuffled subset of genes (still above min_feature_overlap). + subset_genes = list(adata.var_names[:15])[::-1] + query = adata[:, subset_genes].copy() + + result = scalr.annotate(query, + model=model, + device='cpu', + min_feature_overlap=0.1) + assert len(result.labels) == len(query) + assert result.metadata['gene_coverage'] < 1.0 + + +def test_predict_raises_when_gene_overlap_too_low(): + """predict should refuse to score a query with too little gene overlap.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + + unrelated_query = adata[:, adata.var_names[:1]].copy() + unrelated_query.var_names = ['completely_unrelated_gene'] + + with pytest.raises(ValueError): + model.predict(unrelated_query, device='cpu', min_feature_overlap=0.5) + + +def test_annotate_rejects_unsupported_preprocess_mode(): + """annotate should raise NotImplementedError for unsupported preprocess modes.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + + with pytest.raises(NotImplementedError): + scalr.annotate(adata, model=model, device='cpu', preprocess='force') + + +def test_predict_open_set_false_never_flags_unknown(): + """With open_set=False, no cell should be flagged unknown regardless of confidence.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + + result = model.predict(adata, device='cpu', open_set=False) + assert not result.is_unknown.any() + + +def test_train_without_group_key_still_produces_usable_model(): + """Training without a group_key should fall back to stratified splitting and still work.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key=None, + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + + result = model.predict(adata, device='cpu') + assert len(result.labels) == len(adata) + + +def test_load_model_rejects_newer_format_version(tmp_path): + """Loading an artifact with a newer format_version than supported should raise.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + artifact_dir = tmp_path / 'pbmc_v1' + model.save(str(artifact_dir)) + + from scalr.utils import read_data + from scalr.utils import write_data + manifest_path = artifact_dir / 'manifest.json' + manifest = read_data(str(manifest_path)) + manifest['format_version'] = 999 + write_data(manifest, str(manifest_path)) + + with pytest.raises(ValueError): + scalr.load_model(str(artifact_dir)) + + +def test_train_with_taxonomy_supports_broad_level_prediction(): + """A model trained with a taxonomy should support level='broad' prediction.""" + adata = _toy_adata() + taxonomy = {'B_cell': 'B_lineage', 'T_cell': 'T_lineage', 'DC': 'Myeloid'} + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False, + taxonomy=taxonomy) + + fine_result = model.predict(adata, device='cpu') + broad_result = model.predict(adata, device='cpu', level='broad') + + assert sorted(broad_result.class_names) == sorted(set(taxonomy.values())) + assert len(broad_result.labels) == len(fine_result.labels) + + +def test_predict_broad_level_without_taxonomy_raises(): + """Requesting level='broad' on a model with no taxonomy should raise clearly.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + + with pytest.raises(ValueError): + model.predict(adata, device='cpu', level='broad') + + +def test_predict_flag_doublets_populates_result_field(): + """flag_doublets=True should populate is_possible_doublet as a boolean array.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + + result = model.predict(adata, device='cpu', flag_doublets=True) + assert result.is_possible_doublet is not None + assert result.is_possible_doublet.dtype == bool + assert len(result.is_possible_doublet) == len(adata) + + +def test_predict_without_flag_doublets_leaves_field_none(): + """By default (flag_doublets=False), is_possible_doublet should remain None.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + + result = model.predict(adata, device='cpu') + assert result.is_possible_doublet is None + + +def test_predict_cluster_refinement_via_annotate(): + """scalr.annotate should support cluster_refinement end-to-end.""" + adata = _toy_adata() + adata.obs['leiden'] = (adata.obs['donor_id'].astype('category').cat.codes % + 3).astype(str) + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + + result = scalr.annotate(adata, + model=model, + device='cpu', + cluster_refinement='auto') + assert 'raw_labels' in result.metadata + assert result.metadata['cluster_key'] == 'leiden' + + +def test_model_save_writes_model_card(tmp_path): + """Saving a model should write a human-readable README.md model card.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + artifact_dir = tmp_path / 'pbmc_v1' + model.save(str(artifact_dir)) + + card = (artifact_dir / 'README.md').read_text() + assert 'Model Card' in card + assert 'donor_id' in card + + +def test_predict_streaming_matches_non_streaming(): + """streaming=True with a small chunk_size should reproduce the same predictions.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + + non_streaming = model.predict(adata, device='cpu') + # chunk_size smaller than len(adata) forces multiple chunks. + streaming = model.predict(adata, device='cpu', streaming=True, chunk_size=7) + + assert streaming.obs_names == non_streaming.obs_names + assert streaming.labels == non_streaming.labels + assert np.allclose(streaming.probabilities, + non_streaming.probabilities, + atol=1e-5) + + +def test_predict_from_h5ad_path_streams_automatically(tmp_path): + """Passing a file path should stream chunk-by-chunk without needing streaming=True.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + + from scalr.utils import write_data + h5ad_path = tmp_path / 'query.h5ad' + write_data(adata, str(h5ad_path)) + + result = model.predict(str(h5ad_path), device='cpu', chunk_size=9) + + assert result.obs_names == list(adata.obs_names) + assert len(result.labels) == len(adata) + + +def test_load_model_from_local_registry(tmp_path, monkeypatch): + """scalr.load_model should resolve a bare name via the local model registry.""" + monkeypatch.setenv('SCALR_MODEL_REGISTRY', str(tmp_path / 'registry')) + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + artifact_dir = tmp_path / 'pbmc_v1' + model.save(str(artifact_dir)) + scalr.models.register(str(artifact_dir), 'my_registered_model') + + loaded = scalr.load_model('my_registered_model') + assert loaded.class_names == model.class_names + + +def test_load_model_auto_downloads_from_hub_for_repo_id_names( + tmp_path, monkeypatch): + """A '/'-shaped name not found locally should be pulled from the hub.""" + registry_dir = str(tmp_path / 'registry') + monkeypatch.setenv('SCALR_MODEL_REGISTRY', registry_dir) + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + artifact_dir = tmp_path / 'pbmc_v1' + model.save(str(artifact_dir)) + + def fake_download(repo_id, registry_dir=None, **kwargs): + return scalr.models.register(str(artifact_dir), repo_id, registry_dir) + + monkeypatch.setattr(scalr.models, 'download', fake_download) + + loaded = scalr.load_model('org/pbmc-v1') + + assert loaded.class_names == model.class_names + assert scalr.models.list(registry_dir=registry_dir) == ['org/pbmc-v1'] + + +def test_train_records_a_drift_reference(): + """scalr.train should record training-set stats so drift detection works out of the box.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + + assert model.drift_reference is not None + assert len(model.drift_reference['gene_mean']) == len(model.features) + + +def test_detect_drift_top_level_matches_same_distribution_query(): + """scalr.detect_drift(adata, model) should not flag drift when querying with in-distribution data.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + + report = scalr.detect_drift(adata, model) + + assert report.is_drifted is False + assert report.n_query_cells == adata.shape[0] + + +def test_detect_drift_flags_a_shifted_query_dataset(): + """scalr.detect_drift should flag a query dataset whose expression is shifted far from training.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + + drifted = adata.copy() + X = drifted.X + X = X.toarray() if not isinstance(X, np.ndarray) else X + drifted.X = X + 10.0 + + report = scalr.detect_drift(drifted, model) + + assert report.is_drifted is True + + +def test_detect_drift_accepts_a_model_path(tmp_path): + """scalr.detect_drift(adata, model_path) should resolve the model like scalr.annotate does.""" + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + model_dir = tmp_path / 'model' + model.save(str(model_dir)) + + report = scalr.detect_drift(adata, str(model_dir)) + + assert report.n_query_cells == adata.shape[0] diff --git a/scalr/test_benchmark.py b/scalr/test_benchmark.py new file mode 100644 index 0000000..0a5f4e3 --- /dev/null +++ b/scalr/test_benchmark.py @@ -0,0 +1,53 @@ +"""This is a test file for benchmark.py""" + +from scalr.benchmark import capture_environment_metadata +from scalr.benchmark import peak_rss_mb +from scalr.benchmark import run_benchmark +from scalr.utils import generate_dummy_dge_anndata + + +def _toy_adata(n_donors, cell_replicate): + return generate_dummy_dge_anndata(n_donors=n_donors, + cell_type_list=['B_cell', 'T_cell', 'DC'], + cell_replicate=cell_replicate, + n_vars=15) + + +def test_run_benchmark_reports_quality_and_resource_metrics(): + """run_benchmark should produce one row per dataset with quality + cost metrics.""" + datasets = { + 'small': _toy_adata(n_donors=6, cell_replicate=5), + 'medium': _toy_adata(n_donors=6, cell_replicate=10), + } + + df = run_benchmark(datasets, + labels_key='cell_type', + group_key='donor_id', + train_kwargs={ + 'hidden_layers': (8,), + 'epochs': 1, + 'batch_size': 32, + 'device': 'cpu', + }) + + assert list(df.index) == ['small', 'medium'] + for col in ('n_cells', 'n_genes', 'n_classes', 'train_seconds', + 'predict_seconds', 'cells_per_second_predict', 'peak_rss_mb', + 'macro_f1', 'balanced_accuracy'): + assert col in df.columns + assert (df['train_seconds'] > 0).all() + assert (df['n_cells'] == [90, 180]).all() + + +def test_peak_rss_mb_is_positive(): + """Peak RSS should be a positive number of megabytes for a running process.""" + assert peak_rss_mb() > 0 + + +def test_capture_environment_metadata_has_expected_keys(): + """Environment metadata should record versions and hardware availability.""" + meta = capture_environment_metadata() + assert set(meta) == { + 'scalr_version', 'python_version', 'torch_version', 'cuda_available', + 'platform' + } diff --git a/scalr/test_calibration.py b/scalr/test_calibration.py new file mode 100644 index 0000000..b1bc271 --- /dev/null +++ b/scalr/test_calibration.py @@ -0,0 +1,122 @@ +"""This is a test file for calibration.py""" + +import numpy as np +import torch + +from scalr.calibration import brier_score +from scalr.calibration import expected_calibration_error +from scalr.calibration import OpenSetThresholds +from scalr.calibration import predictive_entropy +from scalr.calibration import TemperatureScaler +from scalr.calibration import top1_top2_margin + + +def test_temperature_scaler_roundtrip(): + """A TemperatureScaler's dict form should reconstruct an equivalent scaler.""" + scaler = TemperatureScaler(temperature=2.5) + d = scaler.to_dict() + restored = TemperatureScaler.from_dict(d) + assert restored.temperature == scaler.temperature + + +def test_temperature_scaler_fit_reduces_overconfidence(): + """Fitting on confident-but-frequently-wrong logits should push temperature above 1.""" + torch.manual_seed(0) + n, n_classes = 300, 3 + labels = torch.randint(0, n_classes, (n,)) + + # Model confidently predicts a class that is wrong 40% of the time: + # classic overconfidence that temperature scaling should soften (T > 1). + predicted_class = labels.clone() + flip_idx = torch.randperm(n)[:int(n * 0.4)] + predicted_class[flip_idx] = (labels[flip_idx] + 1) % n_classes + + logits = torch.randn(n, n_classes) * 0.1 + logits[torch.arange(n), predicted_class] += 8.0 + + scaler = TemperatureScaler() + scaler.fit(logits, labels) + + assert scaler.temperature > 1.0 + + +def test_predictive_entropy_bounds(): + """Entropy should be ~0 for a one-hot distribution and ~1 for uniform.""" + one_hot = np.array([[1.0, 0.0, 0.0]]) + uniform = np.array([[1 / 3, 1 / 3, 1 / 3]]) + + assert predictive_entropy(one_hot)[0] < 1e-3 + assert predictive_entropy(uniform)[0] > 0.99 + + +def test_top1_top2_margin(): + """Margin should equal the gap between the two largest probabilities.""" + probs = np.array([[0.6, 0.3, 0.1]]) + margin = top1_top2_margin(probs) + assert np.isclose(margin[0], 0.3) + + +def test_expected_calibration_error_perfect_calibration(): + """ECE should be ~0 when confidence exactly matches accuracy.""" + probabilities = np.array([[0.9, 0.1]] * 10 + [[0.1, 0.9]] * 10) + labels = np.array([0] * 9 + [1] + [1] * 9 + [0]) + ece = expected_calibration_error(probabilities, labels, n_bins=5) + assert ece < 0.2 + + +def test_brier_score_perfect_predictions(): + """Brier score should be 0 for confident, correct one-hot predictions.""" + probabilities = np.array([[1.0, 0.0], [0.0, 1.0]]) + labels = np.array([0, 1]) + assert brier_score(probabilities, labels) == 0.0 + + +def test_open_set_thresholds_flags_low_confidence(): + """A cell with low confidence, high entropy and low margin should be flagged unknown.""" + thresholds = OpenSetThresholds(min_confidence=0.5, + max_entropy=0.7, + min_margin=0.05) + confidence = np.array([0.9, 0.3]) + entropy = np.array([0.1, 0.9]) + margin = np.array([0.8, 0.02]) + + flags = thresholds.flag_unknown(confidence, entropy, margin) + assert list(flags) == [False, True] + + +def test_open_set_thresholds_roundtrip(): + """An OpenSetThresholds' dict form should reconstruct an equivalent object.""" + thresholds = OpenSetThresholds(min_confidence=0.6, + max_entropy=0.5, + min_margin=0.1) + restored = OpenSetThresholds.from_dict(thresholds.to_dict()) + assert restored == thresholds + + +def test_predictive_entropy_single_class_is_zero(): + """A single-class distribution has no uncertainty to measure; entropy should be ~0.""" + probs = np.array([[1.0]]) + assert abs(predictive_entropy(probs)[0]) < 1e-6 + + +def test_top1_top2_margin_single_class_returns_top1(): + """With only one class, the margin degenerates to the top-1 probability.""" + probs = np.array([[0.8]]) + margin = top1_top2_margin(probs) + assert np.isclose(margin[0], 0.8) + + +def test_expected_calibration_error_worst_case_is_bounded(): + """A maximally overconfident-and-wrong model should have high but <=1 ECE.""" + probabilities = np.array([[0.99, 0.01]] * 20) + labels = np.array([1] * 20) # always wrong despite high confidence. + ece = expected_calibration_error(probabilities, labels, n_bins=10) + assert 0.9 <= ece <= 1.0 + + +def test_temperature_scaler_calibrate_probabilities_sum_to_one(): + """Calibrated probabilities should form a valid distribution per row.""" + scaler = TemperatureScaler(temperature=2.0) + logits = np.array([[2.0, 1.0, 0.1], [0.5, 0.5, 0.5]]) + probs = scaler.calibrate_probabilities(logits) + assert np.allclose(probs.sum(axis=1), 1.0) diff --git a/scalr/test_cli.py b/scalr/test_cli.py new file mode 100644 index 0000000..3ef913a --- /dev/null +++ b/scalr/test_cli.py @@ -0,0 +1,174 @@ +"""This is a test file for cli.py""" + +from unittest.mock import MagicMock + +import scalr +from scalr.cli import main +from scalr.utils import generate_dummy_dge_anndata +from scalr.utils import write_data + + +def _write_toy_h5ad(path_): + adata = generate_dummy_dge_anndata( + n_donors=8, + cell_type_list=['B_cell', 'T_cell', 'DC'], + cell_replicate=10, + n_vars=20) + write_data(adata, str(path_)) + return adata + + +def test_cli_validate_returns_zero_for_usable_data(tmp_path, capsys): + """`scalr validate` should print a report and exit 0 for usable data.""" + input_path = tmp_path / 'data.h5ad' + _write_toy_h5ad(input_path) + + code = main([ + 'validate', '--input', + str(input_path), '--labels-key', 'cell_type', '--group-key', 'donor_id' + ]) + + captured = capsys.readouterr() + assert code == 0 + assert 'scaLR validation' in captured.out + + +def test_cli_validate_returns_nonzero_for_missing_labels_key(tmp_path, capsys): + """`scalr validate` should exit non-zero when labels_key is missing.""" + input_path = tmp_path / 'data.h5ad' + _write_toy_h5ad(input_path) + + code = main([ + 'validate', '--input', + str(input_path), '--labels-key', 'does_not_exist' + ]) + assert code == 1 + + +def test_cli_train_then_annotate_roundtrip(tmp_path, capsys): + """`scalr train` then `scalr annotate` should produce an annotated h5ad file.""" + input_path = tmp_path / 'data.h5ad' + _write_toy_h5ad(input_path) + model_dir = tmp_path / 'model' + output_path = tmp_path / 'annotated.h5ad' + + train_code = main([ + 'train', '--input', + str(input_path), '--labels-key', 'cell_type', '--group-key', 'donor_id', + '--output', + str(model_dir), '--epochs', '2', '--device', 'cpu', '--quiet' + ]) + assert train_code == 0 + assert (model_dir / 'manifest.json').exists() + assert (model_dir / 'README.md').exists() + + annotate_code = main([ + 'annotate', '--input', + str(input_path), '--model', + str(model_dir), '--output', + str(output_path), '--device', 'cpu' + ]) + captured = capsys.readouterr() + assert annotate_code == 0 + assert output_path.exists() + assert 'Annotated' in captured.out + + +def test_cli_models_list_reports_empty_registry(tmp_path, monkeypatch): + """`scalr models list` should not crash against an empty/default registry.""" + monkeypatch.setenv('SCALR_MODEL_REGISTRY', str(tmp_path / 'empty_registry')) + code = main(['models', 'list']) + assert code == 0 + + +def test_cli_models_push_calls_hub_push(tmp_path, monkeypatch, capsys): + """`scalr models push` should forward to scalr.models.push and print the URL.""" + fake_push = MagicMock(return_value='https://huggingface.co/org/my-model') + monkeypatch.setattr(scalr.models, 'push', fake_push) + + code = main([ + 'models', 'push', + str(tmp_path / 'model'), 'org/my-model', '--private' + ]) + + captured = capsys.readouterr() + assert code == 0 + fake_push.assert_called_once_with(str(tmp_path / 'model'), + 'org/my-model', + token=None, + private=True) + assert 'https://huggingface.co/org/my-model' in captured.out + + +def test_cli_models_pull_calls_hub_download(tmp_path, monkeypatch, capsys): + """`scalr models pull` should forward to scalr.models.download and report the local path.""" + fake_download = MagicMock(return_value=str(tmp_path / + 'registry/org/my-model')) + monkeypatch.setattr(scalr.models, 'download', fake_download) + + code = main(['models', 'pull', 'org/my-model']) + + captured = capsys.readouterr() + assert code == 0 + fake_download.assert_called_once_with('org/my-model', + revision=None, + token=None, + name=None) + assert 'org/my-model' in captured.out + + +def test_cli_drift_reports_ok_for_in_distribution_data(tmp_path, capsys): + """`scalr drift` should print a report and exit 0 when the query looks like training data.""" + input_path = tmp_path / 'data.h5ad' + _write_toy_h5ad(input_path) + model_dir = tmp_path / 'model' + + train_code = main([ + 'train', '--input', + str(input_path), '--labels-key', 'cell_type', '--group-key', 'donor_id', + '--output', + str(model_dir), '--epochs', '2', '--device', 'cpu', '--quiet' + ]) + assert train_code == 0 + + code = main( + ['drift', '--input', + str(input_path), '--model', + str(model_dir)]) + + captured = capsys.readouterr() + assert code == 0 + assert 'scaLR drift report' in captured.out + assert 'Status: OK' in captured.out + + +def test_cli_drift_fail_on_drift_exits_nonzero_for_shifted_data( + tmp_path, capsys): + """`--fail-on-drift` should make `scalr drift` exit non-zero once drift is detected.""" + input_path = tmp_path / 'data.h5ad' + adata = _write_toy_h5ad(input_path) + model_dir = tmp_path / 'model' + + train_code = main([ + 'train', '--input', + str(input_path), '--labels-key', 'cell_type', '--group-key', 'donor_id', + '--output', + str(model_dir), '--epochs', '2', '--device', 'cpu', '--quiet' + ]) + assert train_code == 0 + + drifted_path = tmp_path / 'drifted.h5ad' + X = adata.X + X = X.toarray() if hasattr(X, 'toarray') else X + adata.X = X + 10.0 + write_data(adata, str(drifted_path)) + + code = main([ + 'drift', '--input', + str(drifted_path), '--model', + str(model_dir), '--fail-on-drift' + ]) + + captured = capsys.readouterr() + assert code == 1 + assert 'Status: DRIFTED' in captured.out diff --git a/scalr/test_doublet.py b/scalr/test_doublet.py new file mode 100644 index 0000000..113a8bd --- /dev/null +++ b/scalr/test_doublet.py @@ -0,0 +1,45 @@ +"""This is a test file for doublet.py""" + +import numpy as np + +from scalr.calibration import predictive_entropy +from scalr.calibration import top1_top2_margin +from scalr.doublet import flag_possible_doublets + + +def test_flag_possible_doublets_flags_close_competing_classes(): + """A cell with two close, substantial class probabilities should be flagged.""" + probabilities = np.array([ + [0.48, 0.47, 0.05], # ambiguous: two competing classes. + [0.95, 0.03, 0.02], # confident single class. + ]) + entropy = predictive_entropy(probabilities) + margin = top1_top2_margin(probabilities) + + flags = flag_possible_doublets(probabilities, entropy, margin) + + assert list(flags) == [True, False] + + +def test_flag_possible_doublets_low_second_probability_not_flagged(): + """A low-confidence prediction without a real second candidate should not be flagged.""" + probabilities = np.array([[0.4, 0.35, 0.25]]) + entropy = predictive_entropy(probabilities) + margin = top1_top2_margin(probabilities) + + flags = flag_possible_doublets(probabilities, + entropy, + margin, + min_second_probability=0.5) + + assert not flags[0] + + +def test_flag_possible_doublets_single_class_never_flags(): + """With only one class, there is no competing class to flag as a doublet.""" + probabilities = np.array([[1.0]]) + entropy = np.zeros(1) + margin = np.ones(1) + + flags = flag_possible_doublets(probabilities, entropy, margin) + assert list(flags) == [False] diff --git a/scalr/test_drift.py b/scalr/test_drift.py new file mode 100644 index 0000000..7f56e86 --- /dev/null +++ b/scalr/test_drift.py @@ -0,0 +1,95 @@ +"""This is a test file for drift.py (dataset-level drift detection)""" + +from anndata import AnnData +import numpy as np +import pytest +from scipy import sparse + +from scalr.drift import compute_reference_stats +from scalr.drift import detect_drift + + +def _adata(X, n_genes=4): + features = [f'gene_{i}' for i in range(n_genes)] + return AnnData(X=sparse.csr_matrix(X)), features + + +def test_compute_reference_stats_matches_manual_mean_std(): + """Reference gene mean/std should match plain numpy over the same matrix.""" + X = np.array([[1.0, 0.0, 3.0, 0.0], [1.0, 2.0, 3.0, 0.0], + [1.0, 4.0, 3.0, 0.0]]) + adata, _ = _adata(X) + + stats = compute_reference_stats(adata) + + np.testing.assert_allclose(stats['gene_mean'], X.mean(axis=0)) + np.testing.assert_allclose(stats['gene_std'], X.std(axis=0)) + assert stats['n_cells'] == 3 + assert stats['library_size_mean'] == pytest.approx(X.sum(axis=1).mean()) + assert stats['n_genes_detected_mean'] == pytest.approx( + (X > 0).sum(axis=1).mean()) + + +def test_detect_drift_reports_no_drift_for_identical_distribution(): + """Querying with data drawn from the same distribution as the reference should not flag drift.""" + rng = np.random.default_rng(0) + X_ref = rng.normal(loc=2.0, scale=0.5, size=(200, 5)).clip(min=0) + X_query = rng.normal(loc=2.0, scale=0.5, size=(200, 5)).clip(min=0) + features = [f'gene_{i}' for i in range(5)] + + reference = compute_reference_stats(AnnData(X=sparse.csr_matrix(X_ref))) + query = AnnData(X=sparse.csr_matrix(X_query)) + + report = detect_drift(query, reference, features) + + assert report.is_drifted is False + assert report.drifted_gene_fraction < 0.1 + + +def test_detect_drift_flags_a_systematic_mean_shift(): + """A large, uniform shift in every gene's mean expression should be flagged as drifted.""" + rng = np.random.default_rng(0) + X_ref = rng.normal(loc=2.0, scale=0.3, size=(200, 5)).clip(min=0) + X_query = X_ref[:150] + 5.0 # shift every gene far outside its reference std + features = [f'gene_{i}' for i in range(5)] + + reference = compute_reference_stats(AnnData(X=sparse.csr_matrix(X_ref))) + query = AnnData(X=sparse.csr_matrix(X_query)) + + report = detect_drift(query, reference, features) + + assert report.is_drifted is True + assert report.drifted_gene_fraction == 1.0 + assert len(report.top_drifted_genes) == 5 + assert report.n_query_cells == 150 + assert report.n_reference_cells == 200 + assert 'DRIFTED' in str(report) + + +def test_detect_drift_flags_severe_library_size_shift_even_if_genes_look_stable( +): + """A big library-size/detected-gene shift should flag drift even if per-gene z-scores stay mild.""" + rng = np.random.default_rng(0) + X_ref = rng.normal(loc=2.0, scale=0.5, size=(200, 5)).clip(min=0) + # Scale every cell up proportionally: per-gene relative pattern is similar, + # but total library size/detected genes shift hard. + X_query = X_ref[:100] * 20.0 + features = [f'gene_{i}' for i in range(5)] + + reference = compute_reference_stats(AnnData(X=sparse.csr_matrix(X_ref))) + query = AnnData(X=sparse.csr_matrix(X_query)) + + report = detect_drift(query, reference, features, severity_z_threshold=3.0) + + assert report.is_drifted is True + assert abs(report.library_size_z) >= 3.0 + + +def test_detect_drift_rejects_feature_count_mismatch(): + """A query with a different number of columns than the reference should raise, not silently misalign.""" + X_ref = np.ones((10, 4)) + reference = compute_reference_stats(AnnData(X=sparse.csr_matrix(X_ref))) + query = AnnData(X=sparse.csr_matrix(np.ones((5, 3)))) + + with pytest.raises(ValueError): + detect_drift(query, reference, features=['a', 'b', 'c']) diff --git a/scalr/test_explain.py b/scalr/test_explain.py new file mode 100644 index 0000000..b0d9e7d --- /dev/null +++ b/scalr/test_explain.py @@ -0,0 +1,93 @@ +"""This is a test file for the per-cell/per-class explainability API (explain.py).""" + +import pytest + +import scalr +from scalr.utils import generate_dummy_dge_anndata + + +def _toy_adata(): + return generate_dummy_dge_anndata(n_donors=8, + cell_type_list=['B_cell', 'T_cell', 'DC'], + cell_replicate=10, + n_vars=20) + + +def _toy_model(): + adata = _toy_adata() + model = scalr.train(adata, + labels_key='cell_type', + group_key='donor_id', + hidden_layers=(16,), + epochs=2, + batch_size=32, + device='cpu', + verbose=False) + return adata, model + + +def test_explain_returns_one_explanation_per_index(): + """model.explain should return one CellExplanation per requested index, in order.""" + adata, model = _toy_model() + indices = [0, 5, 10] + + explanations = model.explain(adata, indices=indices, top_k=5, device='cpu') + + assert len(explanations) == len(indices) + for i, explanation in zip(indices, explanations): + assert isinstance(explanation, scalr.CellExplanation) + assert explanation.obs_name == str(adata.obs_names[i]) + assert explanation.prediction in model.class_names + assert 0.0 <= explanation.confidence <= 1.0 + + +def test_explain_gene_scores_use_model_features(): + """Supporting/contradictory gene names should come from the model's feature list.""" + adata, model = _toy_model() + + explanations = model.explain(adata, indices=[0, 1], top_k=3, device='cpu') + + for explanation in explanations: + for gene, score in explanation.supporting_genes: + assert gene in model.features + assert score > 0 + for gene, score in explanation.contradictory_genes: + assert gene in model.features + assert score < 0 + + +def test_explain_respects_top_k(): + """No more than top_k supporting/contradictory genes should be returned per cell.""" + adata, model = _toy_model() + + explanations = model.explain(adata, indices=[0], top_k=2, device='cpu') + + assert len(explanations[0].supporting_genes) <= 2 + assert len(explanations[0].contradictory_genes) <= 2 + + +def test_explain_class_returns_supporting_and_contradictory_genes(): + """model.explain_class should explain a class independent of any specific cell.""" + _, model = _toy_model() + class_name = model.class_names[0] + + supporting, contradictory = model.explain_class(class_name, + top_k=5, + device='cpu') + + assert isinstance(supporting, list) + assert isinstance(contradictory, list) + for gene, score in supporting: + assert gene in model.features + assert score > 0 + for gene, score in contradictory: + assert gene in model.features + assert score < 0 + + +def test_explain_class_rejects_unknown_class(): + """explain_class should raise a clear error for a class the model doesn't have.""" + _, model = _toy_model() + + with pytest.raises(ValueError): + model.explain_class('not_a_real_class', device='cpu') diff --git a/scalr/test_feature_extraction_pipeline.py b/scalr/test_feature_extraction_pipeline.py new file mode 100644 index 0000000..946e572 --- /dev/null +++ b/scalr/test_feature_extraction_pipeline.py @@ -0,0 +1,125 @@ +"""This is a test file for feature_extraction_pipeline.py's `feature_scoring` step.""" + +import numpy as np + +from scalr.feature.scoring import ScoringBase +import scalr.feature.scoring as scoring_module +from scalr.feature_extraction_pipeline import FeatureExtractionPipeline +from scalr.utils import generate_dummy_dge_anndata + + +class _TaggedFakeScorer(ScoringBase): + """Returns a score block filled with the model's own tag, so we can check + each chunk's output ends up at the right position after (possibly + out-of-order) parallel completion.""" + + def generate_scores(self, model, train_data, val_data, target, mappings): + n_classes = len(mappings[target]['id2label']) + n_features = train_data.shape[1] + return np.full((n_classes, n_features), fill_value=model.tag) + + @classmethod + def get_default_params(cls): + return dict() + + +class _Tag: + """Minimal stand-in for a trained chunked model: just carries an identity tag.""" + + def __init__(self, tag): + self.tag = tag + + +def _build_pipeline(tmp_path, num_workers, monkeypatch): + monkeypatch.setattr(scoring_module, + '_TaggedFakeScorer', + _TaggedFakeScorer, + raising=False) + + adata = generate_dummy_dge_anndata(n_donors=4, + cell_type_list=['B_cell', 'T_cell'], + cell_replicate=5, + n_vars=12) + mappings = { + 'cell_type': { + 'id2label': { + 0: 'B_cell', + 1: 'T_cell' + }, + 'label2id': { + 'B_cell': 0, + 'T_cell': 1 + }, + } + } + + pipeline = FeatureExtractionPipeline( + feature_selection_config={ + 'feature_subsetsize': 4, + 'scoring_config': { + 'name': '_TaggedFakeScorer' + }, + }, + dirpath=str(tmp_path), + device='cpu', + ) + pipeline.set_data_and_targets(adata, adata, 'cell_type', mappings) + pipeline.feature_subsetsize = 4 + pipeline.num_workers = num_workers + pipeline.set_model([_Tag(0), _Tag(1), _Tag(2)]) + return pipeline + + +def test_feature_scoring_preserves_chunk_order_single_worker( + tmp_path, monkeypatch): + """With num_workers=1, each feature chunk's score block should land at its own columns.""" + pipeline = _build_pipeline(tmp_path, num_workers=1, monkeypatch=monkeypatch) + + score_matrix = pipeline.feature_scoring() + + assert score_matrix.shape == (2, 12) + assert (score_matrix.iloc[:, 0:4] == 0).all().all() + assert (score_matrix.iloc[:, 4:8] == 1).all().all() + assert (score_matrix.iloc[:, 8:12] == 2).all().all() + + +def test_set_data_and_targets_stores_sample_chunksize(tmp_path): + """set_data_and_targets should retain sample_chunksize, not just accept and drop it.""" + adata = generate_dummy_dge_anndata(n_donors=2, + cell_type_list=['B_cell'], + cell_replicate=2, + n_vars=4) + mappings = { + 'cell_type': { + 'id2label': { + 0: 'B_cell' + }, + 'label2id': { + 'B_cell': 0 + } + } + } + + pipeline = FeatureExtractionPipeline(feature_selection_config={}, + dirpath=str(tmp_path), + device='cpu') + pipeline.set_data_and_targets(adata, + adata, + 'cell_type', + mappings, + sample_chunksize=500) + + assert pipeline.sample_chunksize == 500 + + +def test_feature_scoring_preserves_chunk_order_multiple_workers( + tmp_path, monkeypatch): + """Parallel (threaded) scoring across chunks should still preserve column order.""" + pipeline = _build_pipeline(tmp_path, num_workers=3, monkeypatch=monkeypatch) + + score_matrix = pipeline.feature_scoring() + + assert score_matrix.shape == (2, 12) + assert (score_matrix.iloc[:, 0:4] == 0).all().all() + assert (score_matrix.iloc[:, 4:8] == 1).all().all() + assert (score_matrix.iloc[:, 8:12] == 2).all().all() diff --git a/scalr/test_feature_stability.py b/scalr/test_feature_stability.py new file mode 100644 index 0000000..18ecd96 --- /dev/null +++ b/scalr/test_feature_stability.py @@ -0,0 +1,62 @@ +"""This is a test file for feature_stability.py""" + +from scalr.feature_stability import compute_feature_stability +from scalr.utils import generate_dummy_dge_anndata + + +def test_compute_feature_stability_returns_bounded_report(): + """Stability report should have a bounded Jaccard similarity and n_runs selection counts.""" + adata = generate_dummy_dge_anndata( + n_donors=6, + cell_type_list=['B_cell', 'T_cell', 'DC'], + cell_replicate=8, + n_vars=15) + + report = compute_feature_stability(adata, + labels_key='cell_type', + top_k=5, + n_runs=4, + subsample_frac=0.8, + seed=0) + + assert report.n_runs == 4 + assert report.top_k == 5 + assert 0.0 <= report.mean_jaccard <= 1.0 + assert (report.selection_frequency <= report.n_runs).all() + assert (report.selection_frequency >= 1).all() + + +def test_compute_feature_stability_top_stable_features_respects_min_runs(): + """top_stable_features should only include genes meeting the min_runs threshold.""" + adata = generate_dummy_dge_anndata( + n_donors=6, + cell_type_list=['B_cell', 'T_cell', 'DC'], + cell_replicate=8, + n_vars=15) + report = compute_feature_stability(adata, + labels_key='cell_type', + top_k=5, + n_runs=4, + seed=1) + + all_selected = report.top_stable_features(min_runs=1) + strict_selected = report.top_stable_features(min_runs=report.n_runs) + + assert set(strict_selected) <= set(all_selected) + + +def test_feature_stability_report_str_contains_summary(): + """The human-readable report should mention the run count and Jaccard summary.""" + adata = generate_dummy_dge_anndata( + n_donors=6, + cell_type_list=['B_cell', 'T_cell', 'DC'], + cell_replicate=8, + n_vars=15) + report = compute_feature_stability(adata, + labels_key='cell_type', + top_k=5, + n_runs=3, + seed=2) + text = str(report) + assert 'Jaccard' in text + assert 'runs' in text.lower() diff --git a/scalr/test_genes.py b/scalr/test_genes.py new file mode 100644 index 0000000..977c286 --- /dev/null +++ b/scalr/test_genes.py @@ -0,0 +1,98 @@ +"""This is a test file for genes.py""" + +from anndata import AnnData +import numpy as np + +from scalr.genes import align_genes +from scalr.utils import generate_dummy_anndata + + +def test_align_genes_full_overlap(): + """When query genes are a superset of reference features, alignment should + reorder/subset columns exactly to the reference order.""" + adata = generate_dummy_anndata(n_samples=10, n_features=10) + reference_features = list(adata.var_names[::-1]) # reversed order. + + aligned, report = align_genes(adata, reference_features) + + assert list(aligned.var_names) == reference_features + assert report.matched_features == len(reference_features) + assert report.missing_features == [] + assert aligned.shape == (10, len(reference_features)) + + +def test_align_genes_missing_features_are_zero_filled(): + """Reference genes absent from the query should be zero-filled, not dropped.""" + adata = generate_dummy_anndata(n_samples=5, n_features=5) + reference_features = list(adata.var_names) + ['unseen_gene'] + + aligned, report = align_genes(adata, reference_features) + + assert report.missing_features == ['unseen_gene'] + unseen_col = list(aligned.var_names).index('unseen_gene') + X = aligned.X + if not isinstance(X, np.ndarray): + X = X.toarray() + assert np.all(X[:, unseen_col] == 0) + + +def test_align_genes_low_overlap_is_unusable(): + """When gene overlap is below the minimum threshold, the report should say so.""" + adata = generate_dummy_anndata(n_samples=5, n_features=5) + reference_features = ['g1', 'g2', 'g3', 'g4', 'g5', 'g6', 'g7', 'g8'] + + _, report = align_genes(adata, reference_features, min_feature_overlap=0.9) + + assert not report.is_usable + + +def test_align_genes_no_overlap_at_all(): + """When none of the reference genes are present, all columns should be zero-filled.""" + adata = generate_dummy_anndata(n_samples=4, n_features=3) + reference_features = ['unrelated_1', 'unrelated_2'] + + aligned, report = align_genes(adata, + reference_features, + min_feature_overlap=0.1) + + assert report.matched_features == 0 + assert report.coverage == 0.0 + assert not report.is_usable + X = aligned.X + if not isinstance(X, np.ndarray): + X = X.toarray() + assert np.all(X == 0) + + +def test_align_genes_preserves_obs(): + """Alignment should carry over the original `obs` metadata unchanged.""" + adata = generate_dummy_anndata(n_samples=6, n_features=4) + reference_features = list(adata.var_names) + + aligned, _ = align_genes(adata, reference_features) + + assert list(aligned.obs_names) == list(adata.obs_names) + assert list(aligned.obs.columns) == list(adata.obs.columns) + + +def test_align_genes_report_str_contains_counts(): + """The report's string form should surface the key coverage numbers.""" + adata = generate_dummy_anndata(n_samples=4, n_features=4) + reference_features = list(adata.var_names) + ['missing_gene'] + + _, report = align_genes(adata, reference_features, min_feature_overlap=0.1) + text = str(report) + + assert 'Reference features: 5' in text + assert 'Missing features: 1' in text + + +def test_align_genes_duplicate_var_names_uses_first_occurrence(): + """A duplicated gene name should resolve to its first column, matching the + "only the first occurrence will be used" contract from scalr.validation.""" + adata = AnnData(X=np.array([[1.0, 999.0, 3.0]])) + adata.var_names = ['A', 'B', 'A'] + + aligned, _ = align_genes(adata, ['A', 'B']) + + np.testing.assert_allclose(aligned.X.toarray(), [[1.0, 999.0]]) diff --git a/scalr/test_hierarchy.py b/scalr/test_hierarchy.py new file mode 100644 index 0000000..bcf8cdf --- /dev/null +++ b/scalr/test_hierarchy.py @@ -0,0 +1,71 @@ +"""This is a test file for hierarchy.py""" + +import numpy as np +import pytest + +from scalr.hierarchy import aggregate_to_broad +from scalr.hierarchy import validate_taxonomy +from scalr.result import PredictionResult + + +def _fine_result(): + class_names = ['CD4_T', 'CD8_T', 'B_cell'] + probabilities = np.array([ + [0.6, 0.3, 0.1], + [0.1, 0.1, 0.8], + [0.2, 0.7, 0.1], + ]) + return PredictionResult( + obs_names=['c0', 'c1', 'c2'], + labels=['CD4_T', 'B_cell', 'CD8_T'], + probabilities=probabilities, + class_names=class_names, + confidence=probabilities.max(axis=1), + entropy=np.zeros(3), + margin=np.full(3, 0.3), + is_unknown=np.array([False, False, True]), + ) + + +def test_validate_taxonomy_passes_with_full_coverage(): + """A taxonomy covering every class should validate without error.""" + validate_taxonomy(['A', 'B'], {'A': 'x', 'B': 'y'}) + + +def test_validate_taxonomy_raises_on_missing_class(): + """A taxonomy missing an entry for a known class should raise.""" + with pytest.raises(ValueError): + validate_taxonomy(['A', 'B'], {'A': 'x'}) + + +def test_aggregate_to_broad_sums_fine_probabilities(): + """Broad-class probability should equal the sum of its fine-class probabilities.""" + result = _fine_result() + taxonomy = {'CD4_T': 'T_cell', 'CD8_T': 'T_cell', 'B_cell': 'B_cell'} + + broad = aggregate_to_broad(result, taxonomy) + + assert sorted(broad.class_names) == ['B_cell', 'T_cell'] + t_cell_idx = broad.class_names.index('T_cell') + # Cell 0: CD4_T (0.6) + CD8_T (0.3) = 0.9 T_cell probability. + assert np.isclose(broad.probabilities[0, t_cell_idx], 0.9) + + +def test_aggregate_to_broad_preserves_obs_names_and_unknown(): + """Aggregation should not change obs_names or the fine-level unknown flag.""" + result = _fine_result() + taxonomy = {'CD4_T': 'T_cell', 'CD8_T': 'T_cell', 'B_cell': 'B_cell'} + + broad = aggregate_to_broad(result, taxonomy) + + assert broad.obs_names == result.obs_names + assert list(broad.is_unknown) == list(result.is_unknown) + assert broad.metadata['level'] == 'broad' + assert broad.metadata['fine_class_names'] == result.class_names + + +def test_aggregate_to_broad_raises_on_incomplete_taxonomy(): + """Aggregating with a taxonomy missing a fine class should raise.""" + result = _fine_result() + with pytest.raises(ValueError): + aggregate_to_broad(result, {'CD4_T': 'T_cell'}) diff --git a/scalr/test_lazy_imports.py b/scalr/test_lazy_imports.py new file mode 100644 index 0000000..31f2e0d --- /dev/null +++ b/scalr/test_lazy_imports.py @@ -0,0 +1,57 @@ +"""This is a test file verifying scalr.analysis/scalr.feature.scoring's heavy +(scanpy/shap-backed) submodules are lazily imported, not paid for by every +`import scalr` (e.g. the `scalr` CLI's `models`/`train`/`annotate` commands). + +Runs in a subprocess: other tests in the same pytest session may already have +imported scanpy/shap, which would make an in-process check order-dependent. +""" + +import subprocess +import sys + + +def _run(code: str) -> str: + result = subprocess.run([sys.executable, '-c', code], + capture_output=True, + text=True, + check=True) + return result.stdout.strip() + + +def test_import_scalr_does_not_eagerly_load_scanpy_or_shap(): + """`import scalr` alone should not pull in scanpy/shap.""" + output = _run("import sys\n" + "import scalr\n" + "print('scanpy' in sys.modules)\n" + "print('shap' in sys.modules)\n") + + assert output == 'False\nFalse' + + +def test_lazy_analysis_attribute_resolves_and_loads_scanpy_on_demand(): + """Accessing a lazy scalr.analysis attribute should still work, loading scanpy on demand.""" + output = _run("import sys\n" + "import scalr\n" + "scalr.analysis.DgeLMEM\n" + "print('scanpy' in sys.modules)\n") + + assert output == 'True' + + +def test_lazy_scoring_attribute_resolves_and_loads_shap_on_demand(): + """Accessing scalr.feature.scoring.ShapScorer should still work, loading shap on demand.""" + output = _run("import sys\n" + "from scalr.feature.scoring import ShapScorer\n" + "print('shap' in sys.modules)\n") + + assert output == 'True' + + +def test_build_analyser_resolves_lazy_class_by_name(): + """The name/params builder pattern (getattr-based) should still find lazy analysis classes.""" + output = _run( + "from scalr.analysis import build_analyser\n" + "analyser, _ = build_analyser({'name': 'Heatmap', 'params': {}})\n" + "print(type(analyser).__name__)\n") + + assert output == 'Heatmap' diff --git a/scalr/test_leakage.py b/scalr/test_leakage.py new file mode 100644 index 0000000..625ba73 --- /dev/null +++ b/scalr/test_leakage.py @@ -0,0 +1,87 @@ +"""This is a test file for leakage.py""" + +import numpy as np +import pandas as pd + +from scalr.leakage import check_group_leakage +from scalr.leakage import detect_likely_grouping_column +from scalr.leakage import group_safe_split +from scalr.utils import generate_dummy_dge_anndata + + +def test_group_safe_split_no_leakage(): + """No donor should appear in more than one split when group_key is used.""" + adata = generate_dummy_dge_anndata(n_donors=10, cell_replicate=5) + obs = adata.obs + + split = group_safe_split(obs, + labels_key='cell_type', + group_key='donor_id', + split_ratio=(0.6, 0.2, 0.2)) + + splits_obs = {name: obs.iloc[idx] for name, idx in split.items()} + report = check_group_leakage(splits_obs, 'donor_id') + + assert not report.has_leakage + # every sample should be assigned to exactly one split. + total = sum(len(v) for v in split.values()) + assert total == len(obs) + + +def test_check_group_leakage_detects_overlap(): + """A donor present in two splits should be flagged as leakage.""" + train = pd.DataFrame({'donor_id': ['d1', 'd1', 'd2']}) + test = pd.DataFrame({'donor_id': ['d2', 'd3']}) + + report = check_group_leakage({'train': train, 'test': test}, 'donor_id') + + assert report.has_leakage + assert 'd2' in report.leaked_groups['train <-> test'] + + +def test_detect_likely_grouping_column(): + """Columns with donor/patient/sample-like names should be detected.""" + obs = pd.DataFrame({'donor_id': ['d1'], 'random_col': [1]}) + detected = detect_likely_grouping_column(obs) + assert detected == ['donor_id'] + + +def test_detect_likely_grouping_column_none_found(): + """No warning-worthy columns should be detected when none look like an identifier.""" + obs = pd.DataFrame({'cell_type': ['T'], 'random_col': [1]}) + assert detect_likely_grouping_column(obs) == [] + + +def test_group_safe_split_without_group_key_falls_back_to_stratified(): + """Without a group_key, splitting should still stratify by label and cover all rows.""" + adata = generate_dummy_dge_anndata(n_donors=6, cell_replicate=5) + obs = adata.obs + + split = group_safe_split(obs, + labels_key='cell_type', + group_key=None, + split_ratio=(0.6, 0.2, 0.2)) + + total = sum(len(v) for v in split.values()) + assert total == len(obs) + # No overlap between splits. + all_idx = np.concatenate(list(split.values())) + assert len(all_idx) == len(set(all_idx.tolist())) + + +def test_check_group_leakage_no_group_column_in_split(): + """A split missing the group column entirely should not be compared or crash.""" + train = pd.DataFrame({'donor_id': ['d1', 'd2']}) + other = pd.DataFrame({'some_other_col': [1, 2]}) + + report = check_group_leakage({'train': train, 'other': other}, 'donor_id') + + assert not report.has_leakage + + +def test_leakage_report_str_no_leakage(): + """The report string should clearly state when no leakage is detected.""" + train = pd.DataFrame({'donor_id': ['d1']}) + test = pd.DataFrame({'donor_id': ['d2']}) + report = check_group_leakage({'train': train, 'test': test}, 'donor_id') + assert 'No leakage detected' in str(report) diff --git a/scalr/test_metrics.py b/scalr/test_metrics.py new file mode 100644 index 0000000..cbf8234 --- /dev/null +++ b/scalr/test_metrics.py @@ -0,0 +1,93 @@ +"""This is a test file for metrics.py""" + +import numpy as np +import pandas as pd + +from scalr.metrics import check_class_imbalance +from scalr.metrics import compute_class_weights +from scalr.metrics import compute_classification_metrics + + +def test_compute_classification_metrics_perfect_predictions(): + """Perfect predictions should give macro-F1 and balanced accuracy of 1.""" + labels = np.array([0, 0, 1, 1, 2, 2]) + predictions = np.array([0, 0, 1, 1, 2, 2]) + class_names = ['A', 'B', 'C'] + + metrics = compute_classification_metrics(labels, predictions, class_names) + + assert metrics.macro_f1 == 1.0 + assert metrics.weighted_f1 == 1.0 + assert metrics.balanced_accuracy == 1.0 + assert metrics.per_class.loc['A', 'support'] == 2 + + +def test_compute_classification_metrics_rare_class_recall(): + """A rare, always-misclassified class should show up with low recall.""" + # Class 'C' (rare) is always predicted as 'A'. + labels = np.array([0] * 20 + [1] * 20 + [2] * 2) + predictions = np.array([0] * 20 + [1] * 20 + [0] * 2) + class_names = ['A', 'B', 'C'] + + metrics = compute_classification_metrics(labels, predictions, class_names) + + assert metrics.per_class.loc['C', 'recall'] == 0.0 + # Macro-F1 penalizes the rare-class failure much more than accuracy would. + assert metrics.macro_f1 < 0.9 + + +def test_check_class_imbalance_flags_small_classes(): + """A severely imbalanced label distribution should flag the small class.""" + labels = pd.Series(['T_cell'] * 1000 + ['DC'] * 5) + report = check_class_imbalance(labels, min_class_size=50) + + assert 'DC' in report.small_classes + assert report.imbalance_ratio > 1 + + +def test_compute_class_weights_inverse_frequency(): + """Rarer classes should get larger weights than common classes.""" + labels = np.array([0] * 90 + [1] * 10) + weights = compute_class_weights(labels, n_classes=2) + + assert weights[1] > weights[0] + + +def test_compute_classification_metrics_class_with_zero_support(): + """A class absent from the evaluation set should show support=0, not crash.""" + labels = np.array([0, 0, 1, 1]) + predictions = np.array([0, 0, 1, 1]) + class_names = ['A', 'B', 'Unseen'] + + metrics = compute_classification_metrics(labels, predictions, class_names) + + assert metrics.per_class.loc['Unseen', 'support'] == 0 + assert metrics.confusion.shape == (3, 3) + + +def test_rare_class_recall_filters_by_support(): + """rare_class_recall should return only classes below the given support threshold.""" + labels = np.array([0] * 50 + [1] * 3) + predictions = np.array([0] * 50 + [1] * 3) + class_names = ['common', 'rare'] + + metrics = compute_classification_metrics(labels, predictions, class_names) + rare = metrics.rare_class_recall(min_support=10) + + assert list(rare.index) == ['rare'] + + +def test_check_class_imbalance_balanced_data_flags_nothing(): + """A roughly balanced label distribution should not flag any small classes.""" + labels = pd.Series(['A'] * 100 + ['B'] * 100) + report = check_class_imbalance(labels, min_class_size=10) + + assert report.small_classes == [] + assert report.imbalance_ratio == 1.0 + + +def test_compute_class_weights_balanced_labels_are_equal(): + """Class weights should be equal (1.0) when classes are perfectly balanced.""" + labels = np.array([0] * 50 + [1] * 50) + weights = compute_class_weights(labels, n_classes=2) + assert np.allclose(weights, [1.0, 1.0]) diff --git a/scalr/test_model_card.py b/scalr/test_model_card.py new file mode 100644 index 0000000..49c3ce3 --- /dev/null +++ b/scalr/test_model_card.py @@ -0,0 +1,55 @@ +"""This is a test file for model_card.py""" + +from scalr.model_card import render_model_card + + +def test_render_model_card_includes_key_fields(): + """The rendered card should surface labels, version, and metrics.""" + metadata = { + 'model_name': 'human_pbmc', + 'model_version': '1.0.0', + 'scalr_version': '2.0.0.dev0', + 'species': 'human', + 'random_seed': 42, + 'labels_key': 'cell_type', + 'group_key': 'donor_id', + 'split_ratio': [0.7, 0.1, 0.2], + 'feature_count': 5000, + 'labels': ['B_cell', 'T_cell', 'DC'], + 'temperature': 1.3, + 'open_set_thresholds': { + 'min_confidence': 0.5 + }, + } + metrics = { + 'macro_f1': 0.812345, + 'weighted_f1': 0.9, + 'balanced_accuracy': 0.85, + 'n_test_cells': 1200, + } + + card = render_model_card(metadata, metrics) + + assert 'human_pbmc' in card + assert 'donor_id' in card + assert 'B_cell, T_cell, DC' in card + assert '0.8123' in card + assert 'T=1.300' in card + + +def test_render_model_card_handles_missing_metrics(): + """A model card with no recorded metrics should still render without error.""" + metadata = {'labels_key': 'cell_type', 'labels': ['A', 'B']} + card = render_model_card(metadata, {}) + assert 'No held-out metrics recorded.' in card + + +def test_render_model_card_handles_missing_group_key(): + """A model trained without group_key should say so explicitly.""" + metadata = { + 'labels_key': 'cell_type', + 'labels': ['A', 'B'], + 'group_key': None + } + card = render_model_card(metadata, {}) + assert 'none (no leakage-safe grouping used)' in card diff --git a/scalr/test_models.py b/scalr/test_models.py new file mode 100644 index 0000000..44384ce --- /dev/null +++ b/scalr/test_models.py @@ -0,0 +1,143 @@ +"""This is a test file for models.py (local model registry)""" + +import sys +from unittest.mock import MagicMock + +import pytest + +from scalr import models + + +def _fake_artifact(tmp_path, name='src_model'): + artifact_dir = tmp_path / name + artifact_dir.mkdir() + (artifact_dir / 'manifest.json').write_text('{"labels": ["A", "B"]}') + return str(artifact_dir) + + +def test_register_and_list_and_info_copy(tmp_path): + """A registered model (copied) should show up in list() and info().""" + artifact_dir = _fake_artifact(tmp_path) + registry_dir = str(tmp_path / 'registry') + + target = models.register(artifact_dir, + 'my_model', + registry_dir=registry_dir) + + assert models.list(registry_dir=registry_dir) == ['my_model'] + manifest = models.info('my_model', registry_dir=registry_dir) + assert manifest['labels'] == ['A', 'B'] + # Copy mode: the registered artifact is independent of the source dir. + assert target != artifact_dir + + +def test_register_pointer_mode_does_not_copy(tmp_path): + """copy=False should register a pointer to the original directory, not duplicate files.""" + artifact_dir = _fake_artifact(tmp_path) + registry_dir = str(tmp_path / 'registry') + + models.register(artifact_dir, + 'ptr_model', + registry_dir=registry_dir, + copy=False) + + resolved = models.resolve_path('ptr_model', registry_dir=registry_dir) + assert resolved == artifact_dir + + +def test_list_empty_registry_returns_empty_list(tmp_path): + """An empty/nonexistent registry directory should list as empty, not raise.""" + registry_dir = str(tmp_path / 'does_not_exist') + assert models.list(registry_dir=registry_dir) == [] + + +def test_info_unknown_model_raises(tmp_path): + """Requesting info for an unregistered model name should raise a clear error.""" + registry_dir = str(tmp_path / 'registry') + with pytest.raises(FileNotFoundError): + models.info('nonexistent', registry_dir=registry_dir) + + +def test_register_rejects_non_artifact_directory(tmp_path): + """Registering a directory without manifest.json should raise.""" + not_an_artifact = tmp_path / 'random_dir' + not_an_artifact.mkdir() + with pytest.raises(ValueError): + models.register(str(not_an_artifact), + 'bad_model', + registry_dir=str(tmp_path / 'registry')) + + +def test_require_hub_missing_dependency_raises_helpful_error(monkeypatch): + """Without the "hub" extra installed, hub access should fail with an actionable message.""" + monkeypatch.setitem(sys.modules, 'huggingface_hub', None) + with pytest.raises(ImportError, match=r'pyscaLR\[hub\]'): + models._require_hub() + + +def test_download_pulls_from_hub_and_registers_locally(monkeypatch, tmp_path): + """download() should snapshot the hub repo and register it under repo_id by default.""" + snapshot_dir = _fake_artifact(tmp_path, name='snapshot') + fake_hub = MagicMock() + fake_hub.snapshot_download.return_value = snapshot_dir + monkeypatch.setattr(models, '_require_hub', lambda: fake_hub) + + registry_dir = str(tmp_path / 'registry') + target = models.download('org/some-model', registry_dir=registry_dir) + + fake_hub.snapshot_download.assert_called_once_with( + repo_id='org/some-model', + revision=None, + token=None, + repo_type='model', + ) + assert models.list(registry_dir=registry_dir) == ['org/some-model'] + assert models.info('org/some-model', registry_dir=registry_dir) == { + 'labels': ['A', 'B'] + } + assert target != snapshot_dir # copied, not pointed at the cache dir. + + +def test_download_rejects_non_artifact_repo(monkeypatch, tmp_path): + """A hub repo without a manifest.json should not be silently registered.""" + not_an_artifact = tmp_path / 'snapshot' + not_an_artifact.mkdir() + fake_hub = MagicMock() + fake_hub.snapshot_download.return_value = str(not_an_artifact) + monkeypatch.setattr(models, '_require_hub', lambda: fake_hub) + + with pytest.raises(ValueError): + models.download('org/bad-model', + registry_dir=str(tmp_path / 'registry')) + + +def test_push_uploads_artifact_and_returns_url(monkeypatch, tmp_path): + """push() should create the repo (if needed) and upload the artifact folder.""" + artifact_dir = _fake_artifact(tmp_path) + fake_api = MagicMock() + fake_hub = MagicMock() + fake_hub.HfApi.return_value = fake_api + monkeypatch.setattr(models, '_require_hub', lambda: fake_hub) + + url = models.push(artifact_dir, 'org/my-model', private=True) + + fake_hub.HfApi.assert_called_once_with(token=None) + fake_api.create_repo.assert_called_once_with('org/my-model', + repo_type='model', + private=True, + exist_ok=True) + fake_api.upload_folder.assert_called_once_with( + repo_id='org/my-model', + folder_path=artifact_dir, + token=None, + commit_message='Upload scaLR model artifact', + ) + assert url == 'https://huggingface.co/org/my-model' + + +def test_push_rejects_non_artifact_directory(tmp_path): + """Pushing a directory without manifest.json should raise before touching the hub.""" + not_an_artifact = tmp_path / 'random_dir' + not_an_artifact.mkdir() + with pytest.raises(ValueError): + models.push(str(not_an_artifact), 'org/bad-model') diff --git a/scalr/test_refinement.py b/scalr/test_refinement.py new file mode 100644 index 0000000..f5cd91e --- /dev/null +++ b/scalr/test_refinement.py @@ -0,0 +1,97 @@ +"""This is a test file for refinement.py""" + +import numpy as np +import pytest + +from scalr.refinement import detect_cluster_key +from scalr.refinement import refine_with_clusters +from scalr.result import PredictionResult +from scalr.utils import generate_dummy_anndata + + +def _result_for(adata, labels, confidence): + n = len(adata) + probabilities = np.zeros((n, 2)) + return PredictionResult( + obs_names=list(adata.obs_names), + labels=labels, + probabilities=probabilities, + class_names=['A', 'B'], + confidence=np.array(confidence), + entropy=np.zeros(n), + margin=np.zeros(n), + is_unknown=np.zeros(n, dtype=bool), + ) + + +def test_detect_cluster_key_finds_known_column(): + """A leiden column should be auto-detected as the cluster key.""" + adata = generate_dummy_anndata(n_samples=4, n_features=3) + adata.obs['leiden'] = ['0', '0', '1', '1'] + assert detect_cluster_key(adata) == 'leiden' + + +def test_detect_cluster_key_returns_none_when_absent(): + """No cluster-like column should return None, not raise.""" + adata = generate_dummy_anndata(n_samples=4, n_features=3) + assert detect_cluster_key(adata) is None + + +def test_refine_with_clusters_applies_majority_vote(): + """A cluster where one label dominates by confidence should be relabeled uniformly.""" + adata = generate_dummy_anndata(n_samples=4, n_features=3) + adata.obs['cluster'] = ['0', '0', '0', '1'] + # cluster 0: 2xA (high conf) + 1xB (low conf) -> majority A. + result = _result_for(adata, ['A', 'A', 'B', 'B'], + confidence=[0.9, 0.9, 0.1, 0.5]) + + refined = refine_with_clusters(result, + adata, + cluster_key='cluster', + min_cluster_agreement=0.5) + + assert refined.labels[:3] == ['A', 'A', 'A'] + assert refined.labels[3] == 'B' # cluster 1 has only one cell, unchanged. + + +def test_refine_with_clusters_preserves_raw_labels(): + """Refinement must keep the pre-refinement labels in metadata['raw_labels'].""" + adata = generate_dummy_anndata(n_samples=3, n_features=3) + adata.obs['cluster'] = ['0', '0', '0'] + result = _result_for(adata, ['A', 'A', 'B'], confidence=[0.9, 0.9, 0.1]) + + refined = refine_with_clusters(result, adata, cluster_key='cluster') + + assert refined.metadata['raw_labels'] == ['A', 'A', 'B'] + # Original result object must not be mutated. + assert result.labels == ['A', 'A', 'B'] + + +def test_refine_with_clusters_auto_detection(): + """cluster_key='auto' should find a 'louvain' column automatically.""" + adata = generate_dummy_anndata(n_samples=3, n_features=3) + adata.obs['louvain'] = ['0', '0', '0'] + result = _result_for(adata, ['A', 'A', 'B'], confidence=[0.9, 0.9, 0.1]) + + refined = refine_with_clusters(result, adata, cluster_key='auto') + assert refined.metadata['cluster_key'] == 'louvain' + + +def test_refine_with_clusters_raises_without_cluster_column(): + """auto detection with no cluster-like column should raise a clear error.""" + adata = generate_dummy_anndata(n_samples=3, n_features=3) + result = _result_for(adata, ['A', 'A', 'B'], confidence=[0.9, 0.9, 0.1]) + + with pytest.raises(ValueError): + refine_with_clusters(result, adata, cluster_key='auto') + + +def test_refine_with_clusters_raises_on_mismatched_obs_names(): + """Mismatched obs_names between result and adata should raise, not silently misalign.""" + adata = generate_dummy_anndata(n_samples=3, n_features=3) + adata.obs['cluster'] = ['0', '0', '0'] + result = _result_for(adata, ['A', 'A', 'B'], confidence=[0.9, 0.9, 0.1]) + result.obs_names = ['x', 'y', 'z'] + + with pytest.raises(ValueError): + refine_with_clusters(result, adata, cluster_key='cluster') diff --git a/scalr/test_result.py b/scalr/test_result.py new file mode 100644 index 0000000..b9671de --- /dev/null +++ b/scalr/test_result.py @@ -0,0 +1,75 @@ +"""This is a test file for result.py""" + +import numpy as np +import pytest + +from scalr.result import PredictionResult +from scalr.utils import generate_dummy_anndata + + +def _toy_result(obs_names): + n = len(obs_names) + probabilities = np.tile([0.7, 0.2, 0.1], (n, 1)) + return PredictionResult( + obs_names=obs_names, + labels=['A'] * n, + probabilities=probabilities, + class_names=['A', 'B', 'C'], + confidence=probabilities.max(axis=1), + entropy=np.zeros(n), + margin=np.full(n, 0.5), + is_unknown=np.zeros(n, dtype=bool), + ) + + +def test_len_matches_labels(): + """__len__ should equal the number of predicted cells.""" + result = _toy_result(['c0', 'c1', 'c2']) + assert len(result) == 3 + + +def test_to_frame_shape_and_index(): + """to_frame should return one row per cell, indexed by obs_names.""" + result = _toy_result(['c0', 'c1']) + frame = result.to_frame() + + assert list(frame.index) == ['c0', 'c1'] + assert set(frame.columns) == { + 'scalr_pred', 'scalr_confidence', 'scalr_entropy', 'scalr_margin', + 'scalr_unknown' + } + assert len(frame) == 2 + + +def test_write_to_adata_success(): + """write_to_adata should populate obs/obsm/uns when obs_names match.""" + adata = generate_dummy_anndata(n_samples=5, n_features=4) + result = _toy_result(list(adata.obs_names)) + + result.write_to_adata(adata) + + assert list(adata.obs['scalr_pred']) == ['A'] * 5 + assert 'scalr_probabilities' in adata.obsm + assert adata.obsm['scalr_probabilities'].shape == (5, 3) + assert adata.uns['scalr'] == result.metadata + + +def test_write_to_adata_mismatched_obs_names_raises(): + """write_to_adata must refuse to write when obs_names are misaligned.""" + adata = generate_dummy_anndata(n_samples=5, n_features=4) + result = _toy_result(['wrong_' + n for n in adata.obs_names]) + + with pytest.raises(ValueError): + result.write_to_adata(adata) + + +def test_write_to_adata_custom_prefix(): + """write_to_adata should honor a custom column/key prefix.""" + adata = generate_dummy_anndata(n_samples=3, n_features=4) + result = _toy_result(list(adata.obs_names)) + + result.write_to_adata(adata, prefix='my_model') + + assert 'my_model_pred' in adata.obs.columns + assert 'my_model_probabilities' in adata.obsm + assert 'my_model' in adata.uns diff --git a/scalr/test_validation.py b/scalr/test_validation.py new file mode 100644 index 0000000..5af8667 --- /dev/null +++ b/scalr/test_validation.py @@ -0,0 +1,119 @@ +"""This is a test file for validation.py""" + +from anndata import AnnData +import numpy as np + +from scalr.utils import generate_dummy_anndata +from scalr.validation import detect_normalization_state +from scalr.validation import validate + + +def test_validate_usable_data(): + """A valid AnnData with a labels_key should validate as usable with no errors.""" + adata = generate_dummy_anndata(n_samples=50, n_features=20) + report = validate(adata, labels_key='celltype') + assert report.is_usable + assert report.errors == [] + + +def test_validate_missing_labels_key(): + """A missing labels_key should be reported as an error.""" + adata = generate_dummy_anndata(n_samples=20, n_features=10) + report = validate(adata, labels_key='does_not_exist') + assert not report.is_usable + assert any('does_not_exist' in e for e in report.errors) + + +def test_validate_duplicate_obs_names(): + """Duplicate cell IDs should be reported as an error.""" + adata = generate_dummy_anndata(n_samples=10, n_features=5) + adata.obs_names = ['cell_0'] * len(adata) + report = validate(adata) + assert not report.is_usable + assert any( + 'duplicate cell IDs' in e.lower() or 'duplicate cell' in e.lower() + for e in report.errors) + + +def test_validate_nan_values(): + """NaN values in X should be reported as an error.""" + adata = generate_dummy_anndata(n_samples=10, n_features=5) + adata.X[0, 0] = np.nan + report = validate(adata) + assert not report.is_usable + assert any('nan' in e.lower() for e in report.errors) + + +def test_validate_gene_overlap_report(): + """Gene coverage against model_features should be reported.""" + adata = generate_dummy_anndata(n_samples=10, n_features=10) + model_features = list( + adata.var_names[:5]) + ['missing_gene_1', 'missing_gene_2'] + report = validate(adata, + model_features=model_features, + min_feature_overlap=0.1) + assert report.is_usable + assert any('missing' in w.lower() for w in report.warnings) + + +def test_validate_low_gene_overlap_is_error(): + """Gene coverage below `min_feature_overlap` should make data unusable.""" + adata = generate_dummy_anndata(n_samples=10, n_features=10) + model_features = ['missing_gene_1', 'missing_gene_2', 'missing_gene_3'] + report = validate(adata, + model_features=model_features, + min_feature_overlap=0.5) + assert not report.is_usable + + +def test_detect_normalization_state_raw_counts(): + """Integer-valued, large-magnitude data should be detected as raw counts.""" + adata = generate_dummy_anndata(n_samples=20, n_features=10) + adata.X = np.random.randint(0, 500, size=adata.X.shape).astype(float) + assert detect_normalization_state(adata) == 'raw_counts' + + +def test_detect_normalization_state_scaled(): + """Data containing negative values should be detected as scaled/standardized.""" + adata = generate_dummy_anndata(n_samples=20, n_features=10) + adata.X = adata.X - 5 + assert detect_normalization_state(adata) == 'scaled' + + +def test_validate_rejects_non_anndata_input(): + """A non-AnnData input should produce a single, clear error.""" + report = validate({'not': 'an anndata'}) + assert not report.is_usable + assert len(report.errors) == 1 + + +def test_validate_empty_adata_is_error(): + """An AnnData with 0 observations should be reported as an error.""" + adata = AnnData(X=np.zeros((0, 5))) + report = validate(adata) + assert not report.is_usable + assert any('0 observations' in e for e in report.errors) + + +def test_validate_missing_group_key_is_error(): + """A missing group_key should be reported as an error.""" + adata = generate_dummy_anndata(n_samples=10, n_features=5) + report = validate(adata, group_key='does_not_exist') + assert not report.is_usable + assert any('does_not_exist' in e for e in report.errors) + + +def test_validate_duplicate_var_names_is_warning_not_error(): + """Duplicate gene IDs should warn but not block usage.""" + adata = generate_dummy_anndata(n_samples=10, n_features=4) + adata.var_names = ['g0', 'g0', 'g1', 'g2'] + report = validate(adata) + assert report.is_usable + assert any('duplicate gene' in w.lower() for w in report.warnings) + + +def test_validation_report_str_reflects_status(): + """The human-readable report should mention 'unusable' when there are errors.""" + adata = generate_dummy_anndata(n_samples=10, n_features=5) + report = validate(adata, labels_key='does_not_exist') + assert 'Status: unusable' in str(report) diff --git a/scalr/utils/file_utils.py b/scalr/utils/file_utils.py index b7aeab9..c9d41a5 100644 --- a/scalr/utils/file_utils.py +++ b/scalr/utils/file_utils.py @@ -236,7 +236,13 @@ def load_full_data_from_config(data_config): full_datas += read_data(path.join(train_val_test_paths[datapath], 'test'), return_anncollection=False) - full_data = AnnCollection(full_datas) + # harmonize_dtypes=False: skip AnnCollection's per-chunk dtype probing, + # which crashes on backed sparse chunks under some anndata versions + # (AttributeError: 'backed_csr_matrix' object has no attribute + # '_validate_indices'). Chunks are written by scalr's own + # write_chunkwise_data from a single source, so dtypes are already + # consistent by construction. + full_data = AnnCollection(full_datas, harmonize_dtypes=False) return full_data @@ -279,7 +285,10 @@ def read_chunked_anndatas(dirpath: str, read_anndata(path.join(dirpath, f'{i}.h5ad'), backed=backed)) else: break - data = AnnCollection(datas) + # harmonize_dtypes=False: see the comment in load_full_data_from_config; + # avoids a crash probing dtypes on backed sparse chunks, and these + # chunks are already dtype-consistent by construction. + data = AnnCollection(datas, harmonize_dtypes=False) if return_anncollection: return data diff --git a/scalr/validation.py b/scalr/validation.py new file mode 100644 index 0000000..54f0121 --- /dev/null +++ b/scalr/validation.py @@ -0,0 +1,230 @@ +"""This file implements the AnnData input-validation subsystem. + +It replaces silent downstream failures (obscure numpy/torch/sklearn tracebacks) +with an explicit, user-facing validation report. +""" + +from dataclasses import dataclass +from dataclasses import field +from typing import Optional, Union + +from anndata import AnnData +import numpy as np + + +@dataclass +class ValidationReport: + """Structured validation result with error/warning/info messages. + + `errors` block usage of the data (`is_usable` is False when non-empty). + `warnings` describe issues that should be surfaced but do not block usage. + `info` records descriptive facts about the input. + """ + + errors: list[str] = field(default_factory=list) + warnings: list[str] = field(default_factory=list) + info: list[str] = field(default_factory=list) + + def add_error(self, message: str) -> None: + self.errors.append(message) + + def add_warning(self, message: str) -> None: + self.warnings.append(message) + + def add_info(self, message: str) -> None: + self.info.append(message) + + @property + def is_usable(self) -> bool: + return len(self.errors) == 0 + + def raise_if_unusable(self) -> None: + if not self.is_usable: + raise ValueError('scaLR validation failed:\n' + + '\n'.join(f' - {e}' for e in self.errors)) + + def __str__(self) -> str: + lines = ['scaLR validation', '-' * 17] + lines.append('ERROR') + lines += [f' {e}' for e in self.errors] if self.errors else [' None'] + lines.append('') + lines.append('WARNING') + lines += [f' {w}' for w in self.warnings + ] if self.warnings else [' None'] + lines.append('') + lines.append('INFO') + lines += [f' {i}' for i in self.info] if self.info else [' None'] + lines.append('') + status = 'usable' if self.is_usable else 'unusable' + if self.is_usable and self.warnings: + status += ' with warning' + lines.append(f'Status: {status}') + return '\n'.join(lines) + + +def detect_normalization_state(adata: AnnData, + layer: Optional[str] = None, + sample_size: int = 500) -> str: + """Heuristically detect whether expression values look like raw counts, + log1p-normalized values, or standardized/scaled values. + + This is a heuristic, not a guarantee. It is used to inform the user (and + `preprocess="auto"`); it never silently transforms the data. + + Returns: + One of 'raw_counts', 'log1p_normalized', 'scaled', 'unknown'. + """ + X = adata.layers[layer] if layer else adata.X + n = min(sample_size, X.shape[0]) + if n == 0: + return 'unknown' + sample = X[:n] + if not isinstance(sample, np.ndarray): + sample = sample.toarray() + sample = np.asarray(sample, dtype=np.float64) + + if sample.size == 0: + return 'unknown' + + has_negative = (sample < 0).any() + if has_negative: + return 'scaled' + + is_integer_like = np.allclose(sample, np.round(sample), atol=1e-6) + if is_integer_like and sample.max() > 30: + return 'raw_counts' + + if sample.max() < 30 and not is_integer_like: + return 'log1p_normalized' + + return 'unknown' + + +def validate( + adata: AnnData, + labels_key: Optional[str] = None, + group_key: Optional[str] = None, + model_features: Optional[list[str]] = None, + min_feature_overlap: float = 0.1, +) -> ValidationReport: + """Run the scaLR input-validation subsystem on an AnnData object. + + Args: + adata: Input AnnData object. + labels_key: Column in `adata.obs` expected to hold classification labels + (checked during training-style validation). + group_key: Column in `adata.obs` expected to hold a grouping identifier + such as donor/patient/sample (checked when leakage-safe splitting + is required). + model_features: Ordered gene/feature list of a reference model, used to + report gene-overlap coverage against the query `adata`. + min_feature_overlap: Minimum fraction of `model_features` that must be + present in `adata.var_names` for the data to remain usable. + + Returns: + A `ValidationReport` describing structural, numerical and metadata + issues found in `adata`. + """ + report = ValidationReport() + + if not isinstance(adata, AnnData): + report.add_error( + f'Expected an AnnData object, got {type(adata).__name__}.') + return report + + # --- Data structure checks ------------------------------------------------- + n_obs, n_vars = adata.shape + if n_obs == 0: + report.add_error('AnnData has 0 observations (cells).') + if n_vars == 0: + report.add_error('AnnData has 0 variables (genes).') + + if adata.obs_names.duplicated().any(): + n_dup = int(adata.obs_names.duplicated().sum()) + report.add_error(f'{n_dup} duplicate cell IDs (obs_names) found.') + + if adata.var_names.duplicated().any(): + n_dup = int(adata.var_names.duplicated().sum()) + report.add_warning( + f'{n_dup} duplicate gene IDs (var_names) found; only the first ' + 'occurrence will be used during gene alignment.') + + if not report.is_usable: + return report + + # --- Numerical checks -------------------------------------------------- + X = adata.X + sample_n = min(2000, n_obs) + X_sample = X[:sample_n] + if not isinstance(X_sample, np.ndarray): + X_sample = X_sample.toarray() + X_sample = np.asarray(X_sample) + + if np.isnan(X_sample).any(): + report.add_error('NaN values detected in `adata.X`.') + if np.isinf(X_sample).any(): + report.add_error('Inf values detected in `adata.X`.') + + if not report.is_usable: + return report + + library_sizes = np.asarray(X_sample.sum(axis=1)).flatten() + n_zero_count = int((library_sizes == 0).sum()) + if n_zero_count: + report.add_warning( + f'{n_zero_count}/{sample_n} sampled cells have zero total counts.') + + positive_library_sizes = library_sizes[library_sizes > 0] + if positive_library_sizes.size: + huge = library_sizes > (np.median(positive_library_sizes) * 100) + if huge.any(): + report.add_warning( + f'{int(huge.sum())}/{sample_n} sampled cells have unusually ' + 'large library sizes (>100x the sample median).') + + norm_state = detect_normalization_state(adata) + report.add_info(f'Detected expression state: {norm_state}') + + # --- Metadata checks ----------------------------------------------------- + if labels_key is not None: + if labels_key not in adata.obs.columns: + report.add_error( + f'labels_key="{labels_key}" not found in adata.obs columns.') + else: + n_missing = int(adata.obs[labels_key].isna().sum()) + if n_missing: + report.add_warning( + f'{n_missing} cells have a missing "{labels_key}" label.') + class_counts = adata.obs[labels_key].value_counts() + report.add_info( + f'Label "{labels_key}": {class_counts.shape[0]} classes, ' + f'smallest class size={int(class_counts.min())}.') + + if group_key is not None and group_key not in adata.obs.columns: + report.add_error( + f'group_key="{group_key}" not found in adata.obs columns.') + + # --- Model / gene-alignment compatibility --------------------------------- + if model_features is not None: + query_genes = set(adata.var_names) + model_genes = set(model_features) + matched = model_genes & query_genes + coverage = len(matched) / len(model_genes) if model_genes else 0.0 + report.add_info(f'Input: {n_obs:,} cells x {n_vars:,} genes; ' + f'model expects {len(model_genes):,} genes; ' + f'gene coverage: {coverage * 100:.1f}%.') + if len(model_genes) - len(matched) > 0: + report.add_warning( + f'{len(model_genes) - len(matched)}/{len(model_genes)} model ' + 'genes are missing from the query.') + if coverage < min_feature_overlap: + report.add_error( + f'Gene coverage ({coverage * 100:.1f}%) is below the minimum ' + f'required overlap ({min_feature_overlap * 100:.1f}%).') + else: + report.add_info(f'Input: {n_obs:,} cells x {n_vars:,} genes.') + + is_sparse = not isinstance(adata.X, np.ndarray) + report.add_info(f'Sparse: {"yes" if is_sparse else "no"}') + + return report diff --git a/tutorials/analysis/differential_gene_expression/dge.ipynb b/tutorials/analysis/differential_gene_expression/dge.ipynb index 65c1d47..1af519c 100644 --- a/tutorials/analysis/differential_gene_expression/dge.ipynb +++ b/tutorials/analysis/differential_gene_expression/dge.ipynb @@ -123,7 +123,7 @@ }, "outputs": [], "source": [ - "!pip install anndata==0.10.9 scanpy==1.10.3 shap==0.46.0 pydeseq2==0.4.11" + "!pip install \"anndata>=0.12,<0.13\" \"scanpy>=1.10.3,<2\" \"shap>=0.46,<1\" \"pydeseq2>=0.4.11,<1\"" ] }, { diff --git a/tutorials/analysis/gene_recall_curve/gene_recall_curve.ipynb b/tutorials/analysis/gene_recall_curve/gene_recall_curve.ipynb index 74be3a1..23fbc5a 100644 --- a/tutorials/analysis/gene_recall_curve/gene_recall_curve.ipynb +++ b/tutorials/analysis/gene_recall_curve/gene_recall_curve.ipynb @@ -186,7 +186,7 @@ }, "outputs": [], "source": [ - "!pip install anndata==0.10.9 scanpy==1.10.3 shap==0.46.0 pydeseq2==0.4.11" + "!pip install \"anndata>=0.12,<0.13\" \"scanpy>=1.10.3,<2\" \"shap>=0.46,<1\" \"pydeseq2>=0.4.11,<1\"" ] }, { diff --git a/tutorials/pipeline/scalr_pipeline.ipynb b/tutorials/pipeline/scalr_pipeline.ipynb index 845e131..0fb0c99 100644 --- a/tutorials/pipeline/scalr_pipeline.ipynb +++ b/tutorials/pipeline/scalr_pipeline.ipynb @@ -6,7 +6,7 @@ "id": "dfGECxsGN9bo" }, "source": [ - "\n", + "\n", "\n", "# Single-cell analysis using Low Resource (scaLR)" ] diff --git a/tutorials/pipeline/scalr_pipeline_local_run.ipynb b/tutorials/pipeline/scalr_pipeline_local_run.ipynb index b98bef8..cbd36c3 100644 --- a/tutorials/pipeline/scalr_pipeline_local_run.ipynb +++ b/tutorials/pipeline/scalr_pipeline_local_run.ipynb @@ -6,7 +6,7 @@ "id": "dfGECxsGN9bo" }, "source": [ - "\n", + "\n", "\n", "# Single-cell analysis using Low Resource (scaLR)" ] diff --git a/tutorials/preprocessing/batch_correction.ipynb b/tutorials/preprocessing/batch_correction.ipynb index eae5dd4..7a4e24b 100644 --- a/tutorials/preprocessing/batch_correction.ipynb +++ b/tutorials/preprocessing/batch_correction.ipynb @@ -1 +1 @@ -{"cells":[{"cell_type":"markdown","metadata":{"id":"WZtTrD0KNRJy"},"source":["# Batch Correction using scaLR"]},{"cell_type":"markdown","metadata":{"id":"dEBM-AOGNRJ1"},"source":["Keypoints\n","\n","1. This notebook is designed as a tutorial for using batch correction using **`SimpleMetaDataLoader`** from **`scaLR`** library.\n","2. The dataloader is extensible to add any column from metadata as one hot-encoded vectors which can be useful for model training."]},{"cell_type":"markdown","metadata":{"id":"z90HlqMeNRJ2"},"source":["## What is the batch correction?\n","\n","- Single-cell genomic datasets related to a specific disease or trait are often compiled from multiple experiments, each sequenced from different single cells under varying conditions, such as capturing times, handling personnel, reagent lots, equipment, and even technology platforms.\n","- These differences lead to large variations also known as `batch effects` in the data. When combining these datasets for analysis and modeling, it's crucial to ensure that the model isn't biased towards data from certain batches due to their higher or lower value ranges. Therefore, it's necessary to eliminate `batch effects` from these datasets.\n","\n","\n","## How to perform batch correction?\n","\n","- Many statistical tools such as `Scanpy`, `Seurat`, `Harmony`, `Combat`, etc. performed `batch correction` by performing `normalization` and `dimensional reduction` and then removing the `batch effect` by fitting `linear or mixed models`, calculating the `k-nearest neighbor` or `mutual nearest neighbors distance`, `canonical correlation analysis`, etc.\n","\n","- While traditional `batch corrections` are robust and widely used, they have limitations in handling non-linear relationships, scalability, flexibility, and the preservation of biological signals, AI/ML-based batch correction methods offer significant advantages in these areas, making them a powerful alternative in complex and large-scale single-cell genomic datasets.\n","\n","- The `batch correction` approach in the **`scaLR`** platform is inspired by the [scGPT](https://www.nature.com/articles/s41592-024-02201-0)(Cui et al.) tool, where batch information is integrated into the feature data to inform the model about the origin of each sample. Since batch is a `categorical` variable, directly including it as a `label-encoded` feature is not appropriate, as no batch is inherently superior to another. Instead, the solution is to use a `one-hot encoded vector` to represent batch information.\n","\n","- For example, if we have four batches in the dataset, the one-hot encoding would work as follows:\n"," - Batch 1 -> 0 0 0 1\n"," - Batch 2 -> 0 0 1 0\n"," - Batch 3 -> 0 1 0 0\n"," - Batch 4 -> 1 0 0 0\n","\n","- These encoded vectors represent the batches and are added to the feature data. In this case, four additional columns will be included in the feature data, ensuring that the model is aware of the batch information while training on samples from different batches.\n","\n","\n","## How is it implemented in the scaLR platform?\n","\n","- In the **`scaLR`** platform, we've implemented **`SimpleMetaDataLoader`** data loader that handles this process automatically.\n","- You can specify the `metadata` column you want to one-hot encode and add it to the feature data, and the data loader will take care of the encoding.\n","- We've also extended the functionality to allow users to include multiple columns from the `metadata` as `one-hot encoded vectors` in the feature data. Bypassing a list of columns, you can easily incorporate additional information.\n","- This approach is particularly useful in scenarios like predicting `disease vs. non-disease` outcomes, where certain `metadata`, such as `cell type`, might enhance the model's predictive power.\n","- By adding the `cell type` information to the feature data using this method, you can improve the model's performance.\n","- Generally, this won't be used as a library utility - it will be mostly used as a part of pipeline. Please find below a explaination a code snippet to understand its basics.\n"]},{"cell_type":"markdown","source":["## Cloning scaLR"],"metadata":{"id":"Lyk1yNRW9iMI"}},{"cell_type":"code","source":["!git clone https://github.com/infocusp/scaLR.git"],"metadata":{"id":"ufzA5OylNWsA","collapsed":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Library Installation and Imports\n","\n"],"metadata":{"id":"__FIQG1T9r-V"}},{"cell_type":"code","source":["!pip install anndata==0.10.9 pydeseq2==0.4.11 scanpy==1.10.3"],"metadata":{"collapsed":true,"id":"59UdOT0nRro6"},"execution_count":null,"outputs":[]},{"cell_type":"code","execution_count":null,"metadata":{"id":"nZwdNgQ7NRJ4"},"outputs":[],"source":["# Required imports.\n","import sys\n","sys.path.append('./scaLR')\n","\n","import anndata\n","import numpy as np\n","import pandas as pd\n","\n","from scalr.nn.dataloader import build_dataloader, simple_metadataloader\n","\n","# Setting seed for reproducibility\n","np.random.seed(0)"]},{"cell_type":"markdown","source":["## Downloading data\n","\n","For this tutorial, we will use two datasets from `cellxgene`([Jin et al. (2021) iScience](https://doi.org/10.1016/j.isci.2021.103115)). The first dataset will serve as batch 1, and the second as batch 2."],"metadata":{"id":"RgQ7PSXy_dtY"}},{"cell_type":"code","source":["!wget -P ./data https://datasets.cellxgene.cziscience.com/16acb1d0-4108-4767-9615-0b42abe09992.h5ad\n","!wget -P ./data https://datasets.cellxgene.cziscience.com/8651e63c-0f98-4a87-bdbd-2da41bdf6de5.h5ad"],"metadata":{"id":"xYf5O0MbdcJ7"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Loading datasets and merging"],"metadata":{"id":"6wZhhGrwClT7"}},{"cell_type":"code","source":["data_1 = anndata.read_h5ad('/content/data/16acb1d0-4108-4767-9615-0b42abe09992.h5ad')\n","data_2 = anndata.read_h5ad('/content/data/8651e63c-0f98-4a87-bdbd-2da41bdf6de5.h5ad')\n","print(f'\\nDataset-1 has {data_1.n_obs} cells and {data_1.n_vars} genes\\nDataset-2 has {data_2.n_obs} cells and {data_2.n_vars} genes')"],"metadata":{"id":"u8fBz_7Zdc7g"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["print(f'Shape of \"obs\" before adding batch\\n\\nDataset-1 : {data_1.obs.shape}\\nDataset-2 : {data_2.obs.shape}')"],"metadata":{"id":"R1_oqEnRDBEu"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Adding batch information.\n","data_1.obs['batch'] = 'batch1'\n","data_2.obs['batch'] = 'batch2'\n","print(f'Shape of \"obs\" after adding batch\\n\\nDataset-1 : {data_1.obs.shape}\\nDataset-2 : {data_2.obs.shape}')"],"metadata":{"id":"5OZ99T9jfOMR"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Combining two datasets.\n","adata = anndata.concat([data_1, data_2])\n","print(f'Combined datset has shape : {adata.shape}')"],"metadata":{"id":"g1urx3oLfcLY"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Batch correction"],"metadata":{"id":"ZJTpOrIXGcAK"}},{"cell_type":"code","execution_count":null,"metadata":{"id":"cUMPGYJ-NRJ5"},"outputs":[],"source":["# Below are the required params for SimpleMetaDataLoader.\n","simple_metadataloader.SimpleMetaDataLoader.__init__.__annotations__"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"kmCGPXeLNRJ6"},"outputs":[],"source":["# Defining required parameters for SimpleMetaDataLoader.\n","\n","# For batch correction you can pass the `batch` column inside `metadata_col`.\n","metadata_col = ['batch', ]\n","\n","# Generating mappings for anndata obs columns.\n","mappings = {}\n","for column_name in adata.obs.columns:\n"," mappings[column_name] = {}\n","\n"," id2label = []\n"," id2label += adata.obs[column_name].astype(\n"," 'category').cat.categories.tolist()\n","\n"," label2id = {id2label[i]: i for i in range(len(id2label))}\n"," mappings[column_name]['id2label'] = id2label\n"," mappings[column_name]['label2id'] = label2id"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"BtTzBQ8oNRJ6"},"outputs":[],"source":["# Creating dataloader object.\n","dataloader = simple_metadataloader.SimpleMetaDataLoader(batch_size=3,\n"," target='cell_type',\n"," mappings=mappings,\n"," metadata_col=metadata_col)"]},{"cell_type":"markdown","source":["We can check if `batch` is added as a one hot-encoded vectors in features data`(i.e. 23586 genes + 2 batches)`.\n","Initially features shape is `(batch_size, 23586)` & there are 2 batches in data.\n","So number of features after adding this column to features data should be `23586+2=23588`.\n","Hence features shape has to be `(batch_size, 23588)` post doing batch correction."],"metadata":{"id":"gzv3mvFYOl5S"}},{"cell_type":"code","execution_count":null,"metadata":{"id":"1O2AyKFxNRJ6"},"outputs":[],"source":["# Verifying features in the dataloader\n","for feature, _ in dataloader(adata):\n"," print('Features shape :', feature.shape)\n"," print('Features :', feature)\n"," break"]},{"cell_type":"markdown","source":["We observe that the feature tensor has a shape of `[3, 23588]`, representing `3` samples with `23,588` features each. The batch information is appended to the gene expression values using one-hot encoding: in this case, each feature vector ends with values `[1, 0]`, representing the samples from `batch 1`."],"metadata":{"id":"OVeWJS8h_Ctu"}},{"cell_type":"code","source":["# Checking what is one hot encoding vector for batches.\n","onhotencode_batches = dataloader.metadata_onehotencoder['batch'].transform(np.array(['batch1', 'batch2']).reshape(-1, 1))\n","onhotencode_batches.A"],"metadata":{"id":"WiQYdHjhzoIW"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["It has been verified that the one-hot encoding vector for batch 1 is [1, 0] and for batch 2 is [0, 1]"],"metadata":{"id":"7t9oD5BoBl1t"}},{"cell_type":"code","source":[],"metadata":{"id":"1LHHx8A1CJf7"},"execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"scalr_minerva","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.9.19"},"colab":{"provenance":[]}},"nbformat":4,"nbformat_minor":0} \ No newline at end of file +{"cells":[{"cell_type":"markdown","metadata":{"id":"WZtTrD0KNRJy"},"source":["# Batch Correction using scaLR"]},{"cell_type":"markdown","metadata":{"id":"dEBM-AOGNRJ1"},"source":["Keypoints\n","\n","1. This notebook is designed as a tutorial for using batch correction using **`SimpleMetaDataLoader`** from **`scaLR`** library.\n","2. The dataloader is extensible to add any column from metadata as one hot-encoded vectors which can be useful for model training."]},{"cell_type":"markdown","metadata":{"id":"z90HlqMeNRJ2"},"source":["## What is the batch correction?\n","\n","- Single-cell genomic datasets related to a specific disease or trait are often compiled from multiple experiments, each sequenced from different single cells under varying conditions, such as capturing times, handling personnel, reagent lots, equipment, and even technology platforms.\n","- These differences lead to large variations also known as `batch effects` in the data. When combining these datasets for analysis and modeling, it's crucial to ensure that the model isn't biased towards data from certain batches due to their higher or lower value ranges. Therefore, it's necessary to eliminate `batch effects` from these datasets.\n","\n","\n","## How to perform batch correction?\n","\n","- Many statistical tools such as `Scanpy`, `Seurat`, `Harmony`, `Combat`, etc. performed `batch correction` by performing `normalization` and `dimensional reduction` and then removing the `batch effect` by fitting `linear or mixed models`, calculating the `k-nearest neighbor` or `mutual nearest neighbors distance`, `canonical correlation analysis`, etc.\n","\n","- While traditional `batch corrections` are robust and widely used, they have limitations in handling non-linear relationships, scalability, flexibility, and the preservation of biological signals, AI/ML-based batch correction methods offer significant advantages in these areas, making them a powerful alternative in complex and large-scale single-cell genomic datasets.\n","\n","- The `batch correction` approach in the **`scaLR`** platform is inspired by the [scGPT](https://www.nature.com/articles/s41592-024-02201-0)(Cui et al.) tool, where batch information is integrated into the feature data to inform the model about the origin of each sample. Since batch is a `categorical` variable, directly including it as a `label-encoded` feature is not appropriate, as no batch is inherently superior to another. Instead, the solution is to use a `one-hot encoded vector` to represent batch information.\n","\n","- For example, if we have four batches in the dataset, the one-hot encoding would work as follows:\n"," - Batch 1 -> 0 0 0 1\n"," - Batch 2 -> 0 0 1 0\n"," - Batch 3 -> 0 1 0 0\n"," - Batch 4 -> 1 0 0 0\n","\n","- These encoded vectors represent the batches and are added to the feature data. In this case, four additional columns will be included in the feature data, ensuring that the model is aware of the batch information while training on samples from different batches.\n","\n","\n","## How is it implemented in the scaLR platform?\n","\n","- In the **`scaLR`** platform, we've implemented **`SimpleMetaDataLoader`** data loader that handles this process automatically.\n","- You can specify the `metadata` column you want to one-hot encode and add it to the feature data, and the data loader will take care of the encoding.\n","- We've also extended the functionality to allow users to include multiple columns from the `metadata` as `one-hot encoded vectors` in the feature data. Bypassing a list of columns, you can easily incorporate additional information.\n","- This approach is particularly useful in scenarios like predicting `disease vs. non-disease` outcomes, where certain `metadata`, such as `cell type`, might enhance the model's predictive power.\n","- By adding the `cell type` information to the feature data using this method, you can improve the model's performance.\n","- Generally, this won't be used as a library utility - it will be mostly used as a part of pipeline. Please find below a explaination a code snippet to understand its basics.\n"]},{"cell_type":"markdown","source":["## Cloning scaLR"],"metadata":{"id":"Lyk1yNRW9iMI"}},{"cell_type":"code","source":["!git clone https://github.com/infocusp/scaLR.git"],"metadata":{"id":"ufzA5OylNWsA","collapsed":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Library Installation and Imports\n","\n"],"metadata":{"id":"__FIQG1T9r-V"}},{"cell_type":"code","source":["!pip install \"anndata>=0.12,<0.13\" \"pydeseq2>=0.4.11,<1\" \"scanpy>=1.10.3,<2\""],"metadata":{"collapsed":true,"id":"59UdOT0nRro6"},"execution_count":null,"outputs":[]},{"cell_type":"code","execution_count":null,"metadata":{"id":"nZwdNgQ7NRJ4"},"outputs":[],"source":["# Required imports.\n","import sys\n","sys.path.append('./scaLR')\n","\n","import anndata\n","import numpy as np\n","import pandas as pd\n","\n","from scalr.nn.dataloader import build_dataloader, simple_metadataloader\n","\n","# Setting seed for reproducibility\n","np.random.seed(0)"]},{"cell_type":"markdown","source":["## Downloading data\n","\n","For this tutorial, we will use two datasets from `cellxgene`([Jin et al. (2021) iScience](https://doi.org/10.1016/j.isci.2021.103115)). The first dataset will serve as batch 1, and the second as batch 2."],"metadata":{"id":"RgQ7PSXy_dtY"}},{"cell_type":"code","source":["!wget -P ./data https://datasets.cellxgene.cziscience.com/16acb1d0-4108-4767-9615-0b42abe09992.h5ad\n","!wget -P ./data https://datasets.cellxgene.cziscience.com/8651e63c-0f98-4a87-bdbd-2da41bdf6de5.h5ad"],"metadata":{"id":"xYf5O0MbdcJ7"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Loading datasets and merging"],"metadata":{"id":"6wZhhGrwClT7"}},{"cell_type":"code","source":["data_1 = anndata.read_h5ad('/content/data/16acb1d0-4108-4767-9615-0b42abe09992.h5ad')\n","data_2 = anndata.read_h5ad('/content/data/8651e63c-0f98-4a87-bdbd-2da41bdf6de5.h5ad')\n","print(f'\\nDataset-1 has {data_1.n_obs} cells and {data_1.n_vars} genes\\nDataset-2 has {data_2.n_obs} cells and {data_2.n_vars} genes')"],"metadata":{"id":"u8fBz_7Zdc7g"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["print(f'Shape of \"obs\" before adding batch\\n\\nDataset-1 : {data_1.obs.shape}\\nDataset-2 : {data_2.obs.shape}')"],"metadata":{"id":"R1_oqEnRDBEu"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Adding batch information.\n","data_1.obs['batch'] = 'batch1'\n","data_2.obs['batch'] = 'batch2'\n","print(f'Shape of \"obs\" after adding batch\\n\\nDataset-1 : {data_1.obs.shape}\\nDataset-2 : {data_2.obs.shape}')"],"metadata":{"id":"5OZ99T9jfOMR"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Combining two datasets.\n","adata = anndata.concat([data_1, data_2])\n","print(f'Combined datset has shape : {adata.shape}')"],"metadata":{"id":"g1urx3oLfcLY"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Batch correction"],"metadata":{"id":"ZJTpOrIXGcAK"}},{"cell_type":"code","execution_count":null,"metadata":{"id":"cUMPGYJ-NRJ5"},"outputs":[],"source":["# Below are the required params for SimpleMetaDataLoader.\n","simple_metadataloader.SimpleMetaDataLoader.__init__.__annotations__"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"kmCGPXeLNRJ6"},"outputs":[],"source":["# Defining required parameters for SimpleMetaDataLoader.\n","\n","# For batch correction you can pass the `batch` column inside `metadata_col`.\n","metadata_col = ['batch', ]\n","\n","# Generating mappings for anndata obs columns.\n","mappings = {}\n","for column_name in adata.obs.columns:\n"," mappings[column_name] = {}\n","\n"," id2label = []\n"," id2label += adata.obs[column_name].astype(\n"," 'category').cat.categories.tolist()\n","\n"," label2id = {id2label[i]: i for i in range(len(id2label))}\n"," mappings[column_name]['id2label'] = id2label\n"," mappings[column_name]['label2id'] = label2id"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"BtTzBQ8oNRJ6"},"outputs":[],"source":["# Creating dataloader object.\n","dataloader = simple_metadataloader.SimpleMetaDataLoader(batch_size=3,\n"," target='cell_type',\n"," mappings=mappings,\n"," metadata_col=metadata_col)"]},{"cell_type":"markdown","source":["We can check if `batch` is added as a one hot-encoded vectors in features data`(i.e. 23586 genes + 2 batches)`.\n","Initially features shape is `(batch_size, 23586)` & there are 2 batches in data.\n","So number of features after adding this column to features data should be `23586+2=23588`.\n","Hence features shape has to be `(batch_size, 23588)` post doing batch correction."],"metadata":{"id":"gzv3mvFYOl5S"}},{"cell_type":"code","execution_count":null,"metadata":{"id":"1O2AyKFxNRJ6"},"outputs":[],"source":["# Verifying features in the dataloader\n","for feature, _ in dataloader(adata):\n"," print('Features shape :', feature.shape)\n"," print('Features :', feature)\n"," break"]},{"cell_type":"markdown","source":["We observe that the feature tensor has a shape of `[3, 23588]`, representing `3` samples with `23,588` features each. The batch information is appended to the gene expression values using one-hot encoding: in this case, each feature vector ends with values `[1, 0]`, representing the samples from `batch 1`."],"metadata":{"id":"OVeWJS8h_Ctu"}},{"cell_type":"code","source":["# Checking what is one hot encoding vector for batches.\n","onhotencode_batches = dataloader.metadata_onehotencoder['batch'].transform(np.array(['batch1', 'batch2']).reshape(-1, 1))\n","onhotencode_batches.A"],"metadata":{"id":"WiQYdHjhzoIW"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["It has been verified that the one-hot encoding vector for batch 1 is [1, 0] and for batch 2 is [0, 1]"],"metadata":{"id":"7t9oD5BoBl1t"}},{"cell_type":"code","source":[],"metadata":{"id":"1LHHx8A1CJf7"},"execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"scalr_minerva","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.9.19"},"colab":{"provenance":[]}},"nbformat":4,"nbformat_minor":0} \ No newline at end of file diff --git a/tutorials/preprocessing/normalization.ipynb b/tutorials/preprocessing/normalization.ipynb index d511784..e416599 100644 --- a/tutorials/preprocessing/normalization.ipynb +++ b/tutorials/preprocessing/normalization.ipynb @@ -66,7 +66,7 @@ }, "outputs": [], "source": [ - "!pip install anndata==0.10.9 scanpy==1.10.3 pydeseq2==0.4.11" + "!pip install \"anndata>=0.12,<0.13\" \"scanpy>=1.10.3,<2\" \"pydeseq2>=0.4.11,<1\"" ] }, {