File size: 6,383 Bytes
3f98d52 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """Fetch a bounded sci-Plex3 prototype by HTTP ranges, with exact-name Tahoe SMILES.
The source is the corrected scPerturb object. The selected cohort is intentionally
small and is an engineering prototype, not the external ICML benchmark.
"""
import argparse
from remedi.download import RangeReader
from pathlib import Path
import json
import hashlib
import time
import numpy as np
import pandas as pd
import requests
import fsspec
import h5py
import anndata as ad
import pyarrow.parquet as pq
from scipy import sparse
URL = "https://zenodo.org/records/13350497/files/SrivatsanTrapnell2020_sciplex3.h5ad?download=1"
DRUG_URL = "https://huggingface.co/datasets/tahoebio/Tahoe-100M/resolve/main/metadata/drug_metadata.parquet"
def read_column(node):
if isinstance(node, h5py.Group):
categories = read_column(node["categories"])
codes = node["codes"][:]
out = categories[np.maximum(codes, 0)].astype(object)
out[codes < 0] = None
return out
values = node[:]
return np.array([v.decode() if isinstance(v, bytes) else v for v in values]) if values.dtype.kind in "SO" else values
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--output", required=True)
parser.add_argument("--source", default=URL, help="Local corrected h5ad or source URL")
parser.add_argument("--cells-per-condition", type=int, default=8)
parser.add_argument("--cell-line", default="A549")
parser.add_argument("--max-molecules", type=int, default=40)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--metadata-cache", help="Optional locally generated metadata cache directory")
args = parser.parse_args()
output = Path(args.output); output.mkdir(parents=True, exist_ok=True)
cache = Path(args.metadata_cache) if args.metadata_cache else output
drug_file = cache/"drug_metadata.parquet"
if not drug_file.exists():
response = requests.get(DRUG_URL, timeout=60); response.raise_for_status()
drug_file.parent.mkdir(parents=True, exist_ok=True); drug_file.write_bytes(response.content)
# Arrow conversion avoids dependence on a pandas-specific parquet extension registry.
drugs = pd.DataFrame(pq.read_table(drug_file).to_pylist())
lookup = {s.strip().casefold(): i for i, s in enumerate(drugs.drug) if isinstance(s, str)}
rng = np.random.default_rng(args.seed)
start = time.perf_counter()
handle = RangeReader(args.source, output/'range_cache', 2526631614) if args.source.startswith('https://') else open(args.source, 'rb')
with handle:
with h5py.File(handle, "r") as h:
names = ["cell_line","dose_value","perturbation","replicate","plate","well","time"]
obs = pd.DataFrame({name: read_column(h["obs"][name]) for name in names})
print(f"Loaded metadata for {len(obs)} cells", flush=True)
matches = sorted(s for s in obs.perturbation.dropna().unique() if s.strip().casefold() in lookup)[:args.max_molecules]
mappings = [{"drug": s, "smiles": drugs.iloc[lookup[s.strip().casefold()]].canonical_smiles,
"tahoe_name": drugs.iloc[lookup[s.strip().casefold()]].drug,
"matching_rule": "exact case-folded name after outer whitespace removal"} for s in matches]
pd.DataFrame(mappings).to_csv(output/"structures.csv", index=False)
eligible = obs[(obs.cell_line == args.cell_line) & (obs.time == 24.)]
treated = eligible[eligible.perturbation.isin(matches) & eligible.dose_value.isin([100.,1000.])]
selected = []
for _, group in treated.groupby(["perturbation","dose_value","replicate","plate"], sort=True):
if len(group) >= args.cells_per_condition:
selected.extend(rng.choice(group.index, args.cells_per_condition, replace=False))
plates = set(obs.loc[selected, "plate"])
controls = eligible[(eligible.perturbation == "control") & eligible.plate.isin(plates)]
for _, group in controls.groupby(["replicate","plate"], sort=True):
selected.extend(rng.choice(group.index, min(args.cells_per_condition*2,len(group)), replace=False))
selected = np.sort(np.unique(selected))
x = h["X"]
if x.attrs.get("encoding-type") != "csr_matrix": raise ValueError("Expected source CSR matrix")
pointers = x["indptr"][:]
values, columns, pointer = [], [], [0]
for j, index in enumerate(selected):
begin, end = int(pointers[index]), int(pointers[index+1])
values.append(x["data"][begin:end]); columns.append(x["indices"][begin:end])
pointer.append(pointer[-1]+end-begin)
if j % 100 == 0: print(f"Read {j}/{len(selected)} selected cell rows", flush=True)
matrix = sparse.csr_matrix((np.concatenate(values),np.concatenate(columns),np.asarray(pointer)),
shape=(len(selected),int(x.attrs["shape"][1])))
var_index = h["var"].attrs.get("_index", "_index")
if isinstance(var_index, bytes): var_index=var_index.decode()
gene_names = read_column(h["var"][var_index]).astype(str)
subset_obs = obs.iloc[selected].copy()
subset_obs["source_row"] = selected
subset_obs.index = [f"sciplex3_row_{i}" for i in selected]
result = ad.AnnData(matrix, obs=subset_obs, var=pd.DataFrame(index=gene_names))
result.write_h5ad(output/"sciplex_prototype.h5ad", compression="gzip")
manifest = {"source_url":URL,"source_record":"scPerturb corrected v1.4 / Zenodo 13350497",
"source_md5":"c9e70629505d98c7ca1a837f62b14e89", "full_source_checksum_recomputed":False,
"smiles_source":DRUG_URL,"molecules":len(matches),"cells":len(selected),
"gene_columns":matrix.shape[1],"nonzero_counts":matrix.nnz,"cell_line":args.cell_line,
"exposure_hours":24,"dose_um":[.1,1.],"seed":args.seed,
"prototype_only":True,"elapsed_seconds":time.perf_counter()-start,
"selected_rows_sha256":hashlib.sha256(selected.tobytes()).hexdigest()}
(output/"source_manifest.json").write_text(json.dumps(manifest, indent=2)+"\n")
print(json.dumps(manifest,indent=2))
if __name__ == "__main__": main()
|