--- license: mit tags: - biology - single-cell - molecular-design - perturbation - dose-response language: - en --- # ReMEDi ![ReMEDi](assets/remedi_overview.png) **Replicate-Enabled Minimax Endpoint-Directed Inference** Generate catalog molecules and concentrations for a desired cellular endpoint. Fit a response ensemble to treated-minus-control features, estimate joint candidate-loss uncertainty from shared draws and replicates, and sample molecule–dose pairs by entropy-regularized minimax regret. ## Installation Python 3.11 or 3.12 on Linux is recommended. The release was tested with Python 3.12. From the repository root: ```bash python -m venv .venv source .venv/bin/activate python -m pip install --upgrade pip pip install -e '.[data,dev]' pytest remedi demo --output runs/demo --max-queries 4 ``` To reproduce the bundled real-count integration checks for both studies and frozen external transfer: ```bash python scripts/validate_release.py --output runs/validation ``` The measured fixtures are included in `examples/`. This command works without downloading the full studies. The demo uses synthetic responses and checks the complete computation. CPU execution is sufficient for the included ridge head, calibration, and convex solver. Frozen transformer embedding extraction requires the optional `encoders` dependencies. ## Data Preparation uses RNA counts, molecular structures, concentrations, cell-line identifiers, and matched control pools from each study. ### Tahoe-100M Download selected official expression shards and metadata: ```bash python scripts/download_data.py tahoe --output data/tahoe --indices 600 601 2600 python scripts/prepare_tahoe.py \ --shards data/tahoe/train-*.parquet --metadata data/tahoe \ --output data/tahoe/cohort --cap 128 ``` The selected shards cover plates 6 and 14. They support a small replicate-based check at 5 µM. Download additional plates for a broader dose and compound study. `prepare_tahoe.py` accepts any collection of official local expression shards. It joins concentration by sample, removes the CLS pair, converts gene tokens to Ensembl identifiers, and retains a reproducible reservoir per compound, dose, cell line, and plate. ### sci-Plex3 The corrected scPerturb v1.4 release is public at [Zenodo](https://doi.org/10.5281/zenodo.13350497). The complete file is approximately 2.53 GB. ```bash python scripts/download_data.py sciplex --output data/sciplex python scripts/fetch_sciplex_subset.py \ --source data/sciplex/sciplex3.h5ad --metadata-cache data/sciplex \ --output data/sciplex/cohort --cells-per-condition 32 --max-molecules 40 ``` This bounded cohort uses A549 at 24 hours, 0.1 and 1 µM, with both replicates and plate-matched controls. Compound structures are joined by exact case-folded names to Tahoe metadata. The nine exact name matches define a small integration test. A larger experiment requires an independently checked `drug,smiles` mapping for the remaining compounds. The `prepare` command accepts the complete corrected file and a corresponding mapping. The subset script also accepts the source HTTPS URL and caches byte ranges. The full-file route is preferable when the server limits range requests. ## Train and evaluate For an end-to-end run on either prepared cohort: ```bash python scripts/run_pipeline.py \ --h5ad data/sciplex/cohort/sciplex_prototype.h5ad \ --mapping configs/sciplex.yaml \ --structures data/sciplex/cohort/structures.csv \ --output runs/sciplex --min-cells 8 --max-cells 32 \ --tolerance 1 --max-queries 6 ``` For Tahoe, replace the three input paths with `data/tahoe/cohort/tahoe.h5ad`, `configs/tahoe.yaml`, and its `structures.csv`. The tolerance is a squared endpoint-distance threshold. Select its value on validation data for a scientific comparison. The command above uses 1 for an integration check. The output includes molecule assignments, condition features, response checkpoints, calibration, all eight generation rules, catalog-pair samples, per-query measurements, molecule-bootstrap summaries, and Ubuntu PDF/PNG plots. The example generation catalog contains measured cohort compounds. Supply your own REAL catalog for REAL generation. For larger studies, use the individual commands to retain fixed splits and hyperparameters: ```bash remedi split --molecules structures.csv --output splits.csv --seed 0 remedi prepare --h5ad cohort.h5ad --mapping configs/tahoe.yaml \ --structures structures.csv --splits splits.csv --output runs/tahoe/data remedi train --data runs/tahoe/data --output runs/tahoe/model --ensemble 5 remedi calibrate --data runs/tahoe/data --model runs/tahoe/model \ --output runs/tahoe/calibration remedi evaluate --data runs/tahoe/data --model runs/tahoe/model \ --calibration runs/tahoe/calibration --output runs/tahoe/evaluation \ --tolerance 1 --max-queries 100 remedi plot --summary runs/tahoe/evaluation/summary.csv --output runs/tahoe/plots ``` Splits are assigned by canonical molecular structure, with every dose and replicate of that structure in one split. `--scaffold` creates grouped Murcko-scaffold splits. Calibration needs at least three held-out compounds with repeated measurements. Retrospective queries use separate replicate/control pools for target construction and scoring. ## Representations The executable CPU configuration uses Morgan fingerprints and training-only log-count PCA. Both are explicit baselines for the frozen-encoder configuration in the manuscript. For frozen UCE features, store per-cell embeddings in `AnnData.obsm['X_uce']` and record the exact checkpoint identifier as `cell_checkpoint` in the mapping YAML. Use `--representation obsm:X_uce` during preparation. UCE extraction is performed with the [released UCE implementation](https://github.com/snap-stanford/UCE). Per-cell embeddings are aggregated after extraction. Cache molecular embeddings with a pinned MolFormer checkpoint: ```bash pip install -e '.[encoders]' remedi encode-molecules --molecules structures.csv --output molecular_features.npz \ --revision CHECKPOINT_COMMIT --device cpu remedi train --data runs/tahoe/data --output runs/tahoe/molformer \ --encoder cached --embeddings molecular_features.npz ``` The cache must cover training, calibration, test, and generation structures. Pass `--embeddings` after moving a cache to another machine. A small MLP response head is available through `train --head mlp`. Encoder extraction and MLP training have separate computational requirements from the CPU reference configuration. ## Generate molecule–dose pairs Provide a CSV with `catalog_id,smiles`, using actual catalog records. ReMEDi retains those IDs in every output. ```bash remedi generate --data runs/tahoe/data --model runs/tahoe/model \ --calibration runs/tahoe/calibration --catalog real_subset.csv \ --target-rows 10 11 --doses 0.05 0.5 5 --shortlist 512 \ --draws 20 --output runs/design ``` Target row numbers refer to `conditions.csv` in the prepared data directory. All selected rows must describe the same compound, dose, and cellular context. Their observed endpoint defines the desired state. `distribution.csv` contains probabilities for the retained pairs. `generated_pairs.csv` contains samples with their catalog IDs and concentrations in µM. Catalog membership is determined from the supplied records. Current stock, synthesis availability, and biological activity require separate confirmation. Catalog preparation is bounded by available memory. Large REAL collections should be partitioned into explicit candidate pools before optimization. ## Formulation For action $a=(m,d)$ and context $c$, a response head predicts $f_\theta(m,\log d,c)$. Shared head and target/control draws give a loss mean $\bar L$ and covariance $C=BB^\top+\operatorname{diag}(v)$. Generation solves $$ \min_{p\in\Delta_N}\max_b\left[(p-e_b)^\top\bar L+ \beta\sqrt{\|B^\top(p-e_b)\|_2^2+\sum_a v_a(p_a-\mathbf 1[a=b])^2}\right] +\tau\,\mathrm{KL}(p\|\pi). $$ Optimization uses CLARABEL with active comparator constraints. Output records include solver status and constraint residuals. The radius is estimated from held-out replicate panels. The selection bound in the manuscript assumes joint coverage of the candidate losses. ## Experiments and plots See [experiment commands](docs/EXPERIMENTS.md) for generation baselines, ablations, frozen external transfer, and figure exports. See [executed checks](docs/VALIDATION.md) for the actual validation scope and dataset sizes. Seven comparison rules are included alongside ReMEDi. These are uniform sampling, minimum predicted loss, Gibbs sampling, marginal uncertainty, absolute robust loss, Thompson sampling, and finite-scenario minimax. All use the same fitted response ensemble and query information. ## Repository layout - `src/remedi/` contains ingestion, response heads, uncertainty estimation, optimization, metrics, and plotting. - `scripts/` contains source download, Tahoe conversion, sci-Plex subsetting, gene alignment, and complete-run commands. - `configs/` contains explicit study field mappings. - `tests/` contains mathematical, data-integrity, and pipeline checks. - `assets/` contains the overview figure and Ubuntu font license. Upload these files directly to the root of a Hugging Face model repository. Training runs locally or on your compute environment after installation. No server process or Space configuration is required. Downloaded datasets and trained checkpoints are excluded from source control by default. ## Sources - [Tahoe-100M data and metadata](https://huggingface.co/datasets/tahoebio/Tahoe-100M) - [Corrected sci-Plex3 release](https://doi.org/10.5281/zenodo.13350497) - [scPerturb publication](https://doi.org/10.1038/s41592-023-02144-y) - [MolFormer checkpoint](https://huggingface.co/ibm/MoLFormer-XL-both-10pct) - [UCE implementation](https://github.com/snap-stanford/UCE) Dataset and encoder licenses remain those of their respective releases. Ubuntu is distributed under its included font license.