ReMEDi

ReMEDi

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:

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:

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:

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. The complete file is approximately 2.53 GB.

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:

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:

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. Per-cell embeddings are aggregated after extraction.

Cache molecular embeddings with a pinned MolFormer checkpoint:

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.

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 for generation baselines, ablations, frozen external transfer, and figure exports. See executed checks 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
• Corrected sci-Plex3 release
• scPerturb publication
• MolFormer checkpoint
• UCE implementation

Dataset and encoder licenses remain those of their respective releases. Ubuntu is distributed under its included font license.