PerturbReason/PerturbReason_dataset_code
012
1#!/usr/bin/env python32"""3Official Data Generator v5 — Unified Chemical + Genetic Perturbation Data4 5Changes from v4:6 - Train/valid JSONL now sourced from a flat directory (vanilla_with_basal/),7 with 4 files (chemical train, chemical valid, genetic train, genetic valid).8 New JSONL format: {"id": ..., "input": {...}, "basal_context": {...}}9 No 'path_ambiguity' or 'is_combo' fields; new fields: chebi, smiles, lfc, pval.10 - Test JSONL sourced from dataset_0426_final/noisy_context/ (was dataset_0419_path_ambiguity).11 - Outputs 3 train/valid adata: chemical_train_valid.h5ad, genetic_train_valid.h5ad,12 merged_train_valid.h5ad (inner join of the two).13 - All 11 test adata (8 standard + chemical2genetic + genetic_combo + genetic2chemical).14 15Genetic datasets:16 hepg2, jurkat: Nadig_2025 (log1tp10k)17 k562, rpe1: Replogle_2022 (log1tp10k)18 mcf7: jiang24_processed.h5ad (other_test / chemical2genetic)19 k562 (combo): norman19_processed.h5ad (other_test / genetic_combo)20 21Chemical datasets:22 sciplex3 (mcf7, a549, k562)23 op3 (B cells, Myeloid cells, NK cells, T cells)24 25genetic2chemical dataset:26 hepg2: Demeo_2025/hepg2_log1tp10k.h5ad27 28Output:29 chemical_train_valid.h5ad30 genetic_train_valid.h5ad31 merged_train_valid.h5ad32 <scenario_name>.h5ad (11 test files)33 <scenario_name>_de_info.json (sidecar for each adata)34"""35 36import pandas as pd37import scanpy as sc38import numpy as np39import glob40import os41import json42import time43import argparse44import anndata as ad45from collections import defaultdict46 47# ── Timing helper ─────────────────────────────────────────────────────────────48_STEP_START: float = 0.049 50 51def _ts() -> str:52 """Return elapsed time since last _tick() call as a formatted string."""53 return f"{time.time() - _STEP_START:.1f}s"54 55 56def _tick(label: str = "") -> None:57 global _STEP_START58 _STEP_START = time.time()59 if label:60 print(f"\n[{time.strftime('%H:%M:%S')}] >>> {label}")61 62# ============================================================================63# Configuration64# ============================================================================65 66 67OMICS_DATA_ROOT = os.environ.get("OMICS_DATA_ROOT")68if not OMICS_DATA_ROOT:69 for _candidate in (70 "/home/USER_NAME/Desktop/state/OMICS_DATA",71 "/home/USER_NAME/scratch/state/OMICS_DATA",72 "/scratch/group/PROJECT_NAME/OMICS_DATA/",73 ""74 ):75 if os.path.exists(_candidate):76 OMICS_DATA_ROOT = _candidate77 break78 else:79 OMICS_DATA_ROOT = "/home/USER_NAME/Desktop/state/OMICS_DATA"80 81 82def omics_path(*parts):83 return os.path.join(OMICS_DATA_ROOT, *parts)84 85 86# --- Train/valid JSONL (new flat directory, 4 files) ---87TRAIN_VALID_JSONL_DIR = (88 omics_path("baseline_collections", "gears_combined_new", "train_valid")89)90 91# --- Test JSONL (updated path) ---92TEST_JSONL_BASE = (93 omics_path("baseline_collections", "dataset_0426_final", "noisy_context")94)95 96OUTPUT_DIR = "processed_data_v5"97 98# --- Genetic datasets ---99GENETIC_DATASETS = {100 'hepg2': omics_path("Perturb_Seq_Datasets", "Nadig_2025", "hepg2_log1tp10k_singlecell_01.h5ad"),101 'jurkat': omics_path("Perturb_Seq_Datasets", "Nadig_2025", "jurkat_log1tp10k_singlecell_01.h5ad"),102 'k562': omics_path("Perturb_Seq_Datasets", "Replogle_2022", "K562_essential_log1tp10k_singlecell_01.h5ad"),103 'rpe1': omics_path("Perturb_Seq_Datasets", "Replogle_2022", "rpe1_log1tp10k_singlecell_01.h5ad"),104}105GENETIC_PERT_COL = 'gene'106GENETIC_CONTROL_PATTERNS = frozenset(['non-targeting', 'control', 'unperturbed', 'ctrl'])107 108# --- Other-test genetic datasets ---109JIANG24_PATH = omics_path("Perturb_Seq_Datasets", "PerturBench_PY", "jiang24_processed.h5ad")110NORMAN19_PATH = omics_path("Perturb_Seq_Datasets", "PerturBench_PY", "norman19_processed.h5ad")111 112# --- genetic2chemical dataset (Demeo 2025) ---113DEMEO_H5AD = omics_path("Perturb_Seq_Datasets", "Demeo_2025", "hepg2_log1tp10k.h5ad")114# Mapping: compound_name in h5ad -> drug-map key used in JSONL pert field115# Only compounds actually present in the h5ad are listed here.116DEMEO_COMPOUND_TO_DRUG_KEY = {117 "KI-8751": "Ki8751", # ✅ confirmed in h5ad118 "Sunitinib Malate": "Sunitinib", # ✅ confirmed in h5ad (not plain "Sunitinib")119}120DEMEO_CONTROL_NAME = "DMSO"121DEMEO_CELL_TYPE = "hepg2"122 123# --- Chemical datasets ---124SCIPLEX_PATH = omics_path("Perturb_Seq_Datasets", "PerturBench_PY", "sciplex3_processed.h5ad")125OP3_PATH = omics_path("Perturb_Seq_Datasets", "PerturBench_PY", "op3_processed.h5ad")126 127CELL_TYPE_TO_CHEM_DATASET = {128 'mcf7': 'sciplex3', 'a549': 'sciplex3', 'k562': 'sciplex3',129 'MCF7': 'sciplex3', 'A549': 'sciplex3', 'K562': 'sciplex3',130 'B cells': 'op3', 'Myeloid cells': 'op3',131 'NK cells': 'op3', 'T cells': 'op3',132}133SCIPLEX_PERT_COL = 'condition'134SCIPLEX_CELL_COL = 'cell_type'135OP3_PERT_COL = 'sm_name'136OP3_CELL_COL = 'cell_type'137CHEM_CONTROL_PATTERNS = frozenset(['control', 'vehicle', 'dmso', 'dimethyl sulfoxide', 'unperturbed', 'ctrl'])138 139 140# ============================================================================141# Train/valid scenario definitions (new flat JSONL directory)142# ============================================================================143# Each entry: output_name -> list of glob patterns relative to TRAIN_VALID_JSONL_DIR144TRAIN_VALID_SCENARIOS = {145 "chemical_train_valid": [146 "*pert_type_chemical*train*",147 "*pert_type_chemical*valid*",148 ],149 "genetic_train_valid": [150 "*pert_type_geneticRNv3*train*",151 "*pert_type_geneticRNv3*valid*",152 ],153}154 155# ============================================================================156# Test scenario definitions (updated JSONL base path)157# ============================================================================158# Each entry: output_name -> (subfolder, glob_pattern)159TEST_SCENARIOS = {160 "genetic_id_test": ("id_test", "*pert_type_geneticRNv3*"),161 "genetic_pert_ood_test": ("pert_ood", "*pert_type_geneticRNv3*"),162 "genetic_cell_ood_test": ("cell_ood", "*pert_type_geneticRNv3*"),163 "genetic_pert_ood_cell_ood_test": ("double_ood", "*pert_type_geneticRNv3*"),164 "chem_id_test": ("id_test", "*pert_type_chemical*"),165 "chem_pert_ood_test": ("pert_ood", "*pert_type_chemical*"),166 "chem_cell_ood_test": ("cell_ood", "*pert_type_chemical*"),167 "chem_pert_ood_cell_ood_test": ("double_ood", "*pert_type_chemical*"),168 # other_test169 # "other_chemical2genetic_test": ("other_test", "*pert_type_s_chemical2genetic*"), # NOTE: temp remove170 "other_genetic_combo_test": ("other_test", "*pert_type_s_genetic_combo*"),171 "other_genetic2chemical_test": ("other_test", "*pert_type_s_genetic2chemical*"),172}173 174SCENARIO_TO_PERT_TYPE = {175 "genetic_id_test": "genetic",176 "genetic_pert_ood_test": "genetic",177 "genetic_cell_ood_test": "genetic",178 "genetic_pert_ood_cell_ood_test": "genetic",179 "chem_id_test": "chemical",180 "chem_pert_ood_test": "chemical",181 "chem_cell_ood_test": "chemical",182 "chem_pert_ood_cell_ood_test": "chemical",183 "other_chemical2genetic_test": "genetic",184 "other_genetic_combo_test": "genetic",185 "other_genetic2chemical_test": "chemical",186}187 188 189# ============================================================================190# JSONL loading (new format: {"id":..., "input":{...}, "basal_context":{...}})191# ============================================================================192def load_jsonl_records(jsonl_path):193 """Load JSONL and return list of record dicts.194 195 New format fields extracted:196 cell_type, pert, gene, label, pert_type, split_pert, split_cell197 No path_ambiguity or is_combo in new format.198 """199 records = []200 n_bad = 0201 with open(jsonl_path, 'r') as f:202 for line in f:203 line = line.strip()204 if not line:205 continue206 try:207 data = json.loads(line)208 inp = data.get('input', {})209 rec = {210 'cell_type': inp.get('cell_type'),211 'pert': inp.get('pert'),212 'gene': inp.get('gene'),213 'label': inp.get('label'),214 'pert_type': inp.get('pert_type'),215 'split_pert': inp.get('split_pert'),216 'split_cell': inp.get('split_cell'),217 }218 records.append(rec)219 except json.JSONDecodeError:220 n_bad += 1221 continue222 if n_bad:223 print(f" WARNING: {n_bad} malformed JSON lines skipped in {os.path.basename(jsonl_path)}")224 return records225 226 227def collect_train_valid_files(patterns):228 """Find JSONL files matching any of the given glob patterns in TRAIN_VALID_JSONL_DIR."""229 files = []230 for pat in patterns:231 found = glob.glob(os.path.join(TRAIN_VALID_JSONL_DIR, pat + ".jsonl"))232 if not found:233 found = glob.glob(os.path.join(TRAIN_VALID_JSONL_DIR, pat))234 files.extend(found)235 return sorted(set(files))236 237 238def collect_test_files(subfolder, pattern):239 """Find JSONL files matching pattern in a subfolder of TEST_JSONL_BASE."""240 search = os.path.join(TEST_JSONL_BASE, subfolder, pattern + ".jsonl")241 files = glob.glob(search)242 if not files:243 files = glob.glob(os.path.join(TEST_JSONL_BASE, subfolder, pattern))244 return sorted(files)245 246 247def load_all_records_from_files(jsonl_files):248 """Load and concatenate records from a list of JSONL files."""249 all_records = []250 for jf in jsonl_files:251 t0 = time.time()252 print(f" Loading {os.path.basename(jf)}...")253 recs = load_jsonl_records(jf)254 elapsed = time.time() - t0255 # Quick label distribution256 label_counts: dict = {}257 for r in recs:258 lbl = r.get('label', '?')259 label_counts[lbl] = label_counts.get(lbl, 0) + 1260 print(f" -> {len(recs):,} records in {elapsed:.1f}s | labels: {label_counts}")261 all_records.extend(recs)262 return all_records263 264 265# ============================================================================266# AnnData loading helpers267# ============================================================================268def load_genetic_adata(cell_type, path):269 """Load a standard genetic perturbation dataset (Nadig/Replogle)."""270 t0 = time.time()271 print(f" Loading genetic [{cell_type}] ...")272 print(f" Path: {path}")273 adata = sc.read_h5ad(path)274 print(f" Raw shape: {adata.shape} ({time.time()-t0:.1f}s)")275 276 if 'gene_name' in adata.var.columns:277 adata.var_names = adata.var['gene_name'].astype(str)278 adata.var_names_make_unique()279 adata.var.index.name = None280 print(f" var_names set from 'gene_name' column")281 282 if GENETIC_PERT_COL not in adata.obs.columns:283 print(f" ERROR: '{GENETIC_PERT_COL}' not in obs. Available: {adata.obs.columns.tolist()}")284 return None, {}285 286 n_perts = adata.obs[GENETIC_PERT_COL].nunique()287 pert_lower = adata.obs[GENETIC_PERT_COL].astype(str).str.lower()288 n_ctrl = pert_lower.isin(GENETIC_CONTROL_PATTERNS).sum()289 print(f" Unique perts: {n_perts} | Control cells: {n_ctrl:,}")290 291 mapping = defaultdict(list)292 for idx, pert in enumerate(adata.obs[GENETIC_PERT_COL]):293 mapping[str(pert)].append(idx)294 295 print(f" Mapping built: {len(mapping)} keys ({time.time()-t0:.1f}s total)")296 return adata, mapping297 298 299def load_jiang24():300 """Load jiang24 (mcf7 CRISPRi). Uses obs['gene'] for pert."""301 t0 = time.time()302 print(f" Loading jiang24 (chemical2genetic) ...")303 print(f" Path: {JIANG24_PATH}")304 adata = sc.read_h5ad(JIANG24_PATH)305 if 'gene_name' not in adata.var.columns:306 adata.var['gene_name'] = adata.var.index.astype(str)307 308 cell_types = adata.obs['cell_type'].unique().tolist()309 n_perts = adata.obs['gene'].nunique()310 print(f" Shape: {adata.shape} | cell_types: {cell_types}")311 print(f" Unique gene perts: {n_perts} ({time.time()-t0:.1f}s)")312 313 obs = adata.obs314 obs_temp = pd.DataFrame({315 'ct': obs['cell_type'].astype(str).values,316 'gene': obs['gene'].astype(str).values,317 })318 grouped = obs_temp.groupby(['ct', 'gene'])319 mapping = {k: v.tolist() for k, v in grouped.indices.items()}320 321 print(f" Mapping built: {len(mapping)} (cell_type, gene) keys ({time.time()-t0:.1f}s total)")322 return adata, mapping323 324 325def load_norman19():326 """Load norman19 (k562 combo). Uses obs['condition'] for pert."""327 t0 = time.time()328 print(f" Loading norman19 (genetic_combo) ...")329 print(f" Path: {NORMAN19_PATH}")330 adata = sc.read_h5ad(NORMAN19_PATH)331 if 'gene_name' not in adata.var.columns:332 if 'gene_symbol' in adata.var.columns:333 adata.var['gene_name'] = adata.var['gene_symbol'].astype(str)334 else:335 adata.var['gene_name'] = adata.var.index.astype(str)336 337 cell_types = adata.obs['cell_type'].unique().tolist()338 n_conds = adata.obs['condition'].nunique()339 print(f" Shape: {adata.shape} | cell_types: {cell_types}")340 print(f" Unique conditions: {n_conds} ({time.time()-t0:.1f}s)")341 342 cond_series = pd.Series(adata.obs['condition'].astype(str).values)343 grouped = cond_series.groupby(cond_series)344 mapping = {k: v.tolist() for k, v in grouped.indices.items()}345 346 print(f" Mapping built: {len(mapping)} condition keys ({time.time()-t0:.1f}s total)")347 return adata, mapping348 349 350def load_demeo():351 """Load Demeo 2025 hepg2 chemical dataset for genetic2chemical test."""352 t0 = time.time()353 print(f" Loading Demeo 2025 hepg2 (genetic2chemical) ...")354 print(f" Path: {DEMEO_H5AD}")355 adata = sc.read_h5ad(DEMEO_H5AD)356 print(f" Shape: {adata.shape} ({time.time()-t0:.1f}s)")357 358 compound_col = adata.obs['compound_name'].astype(str)359 all_compounds = compound_col.unique().tolist()360 print(f" All compound_name values: {all_compounds}")361 362 drug_key_to_compound = {v: k for k, v in DEMEO_COMPOUND_TO_DRUG_KEY.items()}363 364 # Build mapping: drug_key -> [indices] (treatment cells only)365 mapping = defaultdict(list)366 for idx, compound in enumerate(compound_col):367 drug_key = DEMEO_COMPOUND_TO_DRUG_KEY.get(compound)368 if drug_key is not None:369 mapping[drug_key].append(idx)370 371 # Control indices (DMSO)372 ctrl_mask = (compound_col == DEMEO_CONTROL_NAME).values373 ctrl_indices = np.where(ctrl_mask)[0].tolist()374 375 print(f" Drug key -> n_cells: { {k: len(v) for k, v in mapping.items()} }")376 print(f" Control (DMSO) cells: {len(ctrl_indices):,} ({time.time()-t0:.1f}s total)")377 378 return adata, mapping, ctrl_indices, drug_key_to_compound379 380 381def load_chemical_adata(path, dataset_name):382 """Load chemical dataset (sciplex3 or op3)."""383 t0 = time.time()384 print(f" Loading chemical [{dataset_name}] ...")385 print(f" Path: {path}")386 adata = sc.read_h5ad(path)387 388 if dataset_name == 'sciplex3':389 pert_col, cell_col = SCIPLEX_PERT_COL, SCIPLEX_CELL_COL390 elif dataset_name == 'op3':391 pert_col, cell_col = OP3_PERT_COL, OP3_CELL_COL392 else:393 raise ValueError(f"Unknown chemical dataset: {dataset_name}")394 395 if 'gene_name' not in adata.var.columns:396 if 'gene_symbol' in adata.var.columns:397 adata.var['gene_name'] = adata.var['gene_symbol'].astype(str)398 else:399 adata.var['gene_name'] = adata.var.index.astype(str)400 401 print(f" Shape: {adata.shape} | pert_col='{pert_col}' cell_col='{cell_col}' ({time.time()-t0:.1f}s)")402 403 # Cell type distribution404 cell_types = adata.obs[cell_col].value_counts().to_dict()405 print(f" Cell type counts: {cell_types}")406 407 mapping = defaultdict(list)408 for idx, (ct, pt) in enumerate(zip(adata.obs[cell_col], adata.obs[pert_col])):409 mapping[(str(ct), str(pt))].append(idx)410 411 # Control cell count412 pert_lower = adata.obs[pert_col].astype(str).str.lower()413 n_ctrl = pert_lower.isin(CHEM_CONTROL_PATTERNS).sum()414 print(f" Unique (cell_type, pert): {len(mapping)} | Control cells: {n_ctrl:,} ({time.time()-t0:.1f}s total)")415 416 adata.uns['_pert_col'] = pert_col417 adata.uns['_cell_col'] = cell_col418 419 return adata, mapping420 421 422# ============================================================================423# de_info builder (no path_ambiguity in new format)424# ============================================================================425def build_de_info(records, pert_key_fn=None):426 """Build de_info dict from records.427 428 pert_key_fn: optional callable(rec) -> str for the pert key part.429 Defaults to rec['pert'].430 Returns: {"cell_type|pert": [{"gene": ..., "label": ...}, ...]}431 """432 de_info = defaultdict(list)433 for rec in records:434 pert = pert_key_fn(rec) if pert_key_fn else str(rec['pert'])435 key = f"{rec['cell_type']}|{pert}"436 de_info[key].append({'gene': rec['gene'], 'label': rec['label']})437 return dict(de_info)438 439 440# ============================================================================441# Core processing - Genetic (standard: hepg2, jurkat, k562, rpe1)442# ============================================================================443def process_genetic_scenario(444 scenario_name,445 records,446 adata_dict,447 mapping_dict,448 join_cell_type='inner',449 fill_value=-1.0,450):451 """Process standard genetic perturbation data."""452 print(f"\n{'='*60}")453 print(f"Processing genetic scenario: {scenario_name}")454 print(f"{'='*60}")455 456 if not records:457 print(" No records. Skipping.")458 return None459 460 records_df = pd.DataFrame(records)461 print(f" Total records: {len(records_df)}")462 463 # Build (cell_type_lower, pert) -> split_pert for the 'split' obs column464 pert_split_map: dict = {}465 for rec in records:466 sp = rec.get('split_pert')467 if sp:468 pert_split_map[(str(rec['cell_type']).lower(), str(rec['pert']))] = sp469 470 unique_conds = records_df[['cell_type', 'pert']].drop_duplicates()471 print(f" Unique (cell_type, pert) conditions: {len(unique_conds)}")472 473 cells_to_include = []474 matched = 0475 unmatched = []476 477 for _, row in unique_conds.iterrows():478 ct = str(row['cell_type']).lower()479 pert = str(row['pert'])480 481 if ct not in mapping_dict:482 unmatched.append((row['cell_type'], pert, f"Cell type '{ct}' not loaded"))483 continue484 485 mapping = mapping_dict[ct]486 if pert in mapping:487 matched += 1488 cells_to_include.append({'cell_type': ct, 'pert': pert, 'indices': mapping[pert]})489 else:490 unmatched.append((row['cell_type'], pert, "Pert not found in AnnData"))491 492 print(f" Matched: {matched}/{len(unique_conds)}")493 if unmatched:494 print(f" Unmatched (first 10):")495 for ct, p, reason in unmatched[:10]:496 print(f" {ct} | {p} | {reason}")497 498 if not cells_to_include:499 print(" No cells matched. Skipping.")500 return None501 502 dataset_subsets = defaultdict(list)503 for item in cells_to_include:504 dataset_subsets[item['cell_type']].append(item)505 506 processed_adatas = []507 for cell_type, items in dataset_subsets.items():508 t_ct = time.time()509 adata = adata_dict[cell_type]510 all_indices = []511 for item in items:512 all_indices.extend(item['indices'])513 n_pert_cells = len(all_indices)514 515 # Add control cells516 pert_lower = adata.obs[GENETIC_PERT_COL].astype(str).str.lower()517 ctrl_mask = pert_lower.isin(GENETIC_CONTROL_PATTERNS)518 control_indices = np.where(ctrl_mask.values)[0].tolist()519 if control_indices:520 print(f" [{cell_type}] Adding {len(control_indices):,} control cells")521 all_indices.extend(control_indices)522 523 all_indices = sorted(set(all_indices))524 print(f" [{cell_type}] Extracting {len(all_indices):,} cells "525 f"({n_pert_cells:,} pert + {len(control_indices):,} ctrl) ...")526 527 subset = adata[all_indices].copy()528 print(f" [{cell_type}] Subset copied: {subset.shape} ({time.time()-t_ct:.1f}s)")529 530 pert_vals = subset.obs[GENETIC_PERT_COL].astype(str)531 is_ctrl = pert_vals.str.lower().isin(GENETIC_CONTROL_PATTERNS)532 533 new_obs = pd.DataFrame(index=subset.obs.index)534 new_obs['cell_type'] = cell_type535 new_obs['pert'] = pert_vals.values536 new_obs['pert_type'] = 'genetic'537 new_obs['control'] = np.where(is_ctrl.values, 'True', 'False')538 new_obs['is_normalized'] = 'True'539 new_obs['split'] = [540 'ctrl' if is_c else pert_split_map.get((cell_type, str(p)), 'train')541 for is_c, p in zip(is_ctrl.values, pert_vals.values)542 ]543 544 subset.obs = new_obs545 for col in subset.obs.columns:546 subset.obs[col] = subset.obs[col].astype(str)547 processed_adatas.append(subset)548 549 if not processed_adatas:550 return None551 552 if len(processed_adatas) > 1:553 print(f" Concatenating {len(processed_adatas)} cell-type subsets ({join_cell_type} join) ...")554 t_cat = time.time()555 concat_kwargs = {"join": join_cell_type}556 if join_cell_type == "outer":557 concat_kwargs["fill_value"] = fill_value558 final = ad.concat(processed_adatas, **concat_kwargs)559 if join_cell_type == "outer":560 final.uns["missing_gene_fill_value"] = float(fill_value)561 print(f" Concat done: {final.shape} ({time.time()-t_cat:.1f}s)")562 else:563 final = processed_adatas[0]564 565 de_info = build_de_info(records)566 final.uns['de_info_json'] = json.dumps(de_info)567 568 # Label distribution summary569 label_counts: dict = {}570 for entries in de_info.values():571 for e in entries:572 lbl = e.get('label', '?')573 label_counts[lbl] = label_counts.get(lbl, 0) + 1574 print(f" de_info: {len(de_info)} (cell_type|pert) keys | label dist: {label_counts}")575 print(f" Final shape: {final.shape}")576 return final, de_info, len(processed_adatas) > 1577 578 579# ============================================================================580# Core processing - Chemical581# ============================================================================582def process_chemical_scenario(583 scenario_name,584 records,585 adata_dict,586 mapping_dict,587 join_cell_type='inner',588 fill_value=-1.0,589):590 """Process chemical perturbation data (sciplex3 / op3)."""591 # Build (cell_type_lower, pert) -> split_pert for the 'split' obs column592 pert_split_map: dict = {}593 for rec in records:594 sp = rec.get('split_pert')595 if sp:596 pert_split_map[(str(rec['cell_type']).lower(), str(rec['pert']))] = sp597 598 print(f"\n{'='*60}")599 print(f"Processing chemical scenario: {scenario_name}")600 print(f"{'='*60}")601 602 if not records:603 print(" No records. Skipping.")604 return None605 606 records_df = pd.DataFrame(records)607 print(f" Total records: {len(records_df)}")608 609 unique_conds = records_df[['cell_type', 'pert']].drop_duplicates()610 print(f" Unique (cell_type, pert) conditions: {len(unique_conds)}")611 612 cells_to_include = []613 matched = 0614 unmatched = []615 616 for _, row in unique_conds.iterrows():617 ct = str(row['cell_type'])618 pert = str(row['pert'])619 620 dataset = CELL_TYPE_TO_CHEM_DATASET.get(ct)621 if dataset is None:622 unmatched.append((ct, pert, "Unknown cell type"))623 continue624 if dataset not in mapping_dict:625 unmatched.append((ct, pert, f"Dataset {dataset} not loaded"))626 continue627 628 mapping = mapping_dict[dataset]629 key = (ct, pert)630 if key in mapping:631 matched += 1632 cells_to_include.append({'dataset': dataset, 'indices': mapping[key], 'cell_type': ct, 'pert': pert})633 else:634 unmatched.append((ct, pert, "Not found in AnnData"))635 636 print(f" Matched: {matched}/{len(unique_conds)}")637 if unmatched:638 print(f" Unmatched (first 10):")639 for ct, p, reason in unmatched[:10]:640 print(f" {ct} | {p} | {reason}")641 642 if not cells_to_include:643 print(" No cells matched. Skipping.")644 return None645 646 dataset_subsets = defaultdict(list)647 for item in cells_to_include:648 dataset_subsets[item['dataset']].append(item)649 650 processed_adatas = []651 for dataset_name, items in dataset_subsets.items():652 t_ds = time.time()653 adata = adata_dict[dataset_name]654 pert_col = adata.uns['_pert_col']655 cell_col = adata.uns['_cell_col']656 657 all_indices = []658 for item in items:659 all_indices.extend(item['indices'])660 n_pert_cells = len(all_indices)661 662 # Add control cells663 pert_lower = adata.obs[pert_col].astype(str).str.lower()664 ctrl_mask = pert_lower.isin(CHEM_CONTROL_PATTERNS)665 control_indices = np.where(ctrl_mask.values)[0].tolist()666 if control_indices:667 print(f" [{dataset_name}] Adding {len(control_indices):,} control cells")668 all_indices.extend(control_indices)669 670 all_indices = sorted(set(all_indices))671 print(f" [{dataset_name}] Extracting {len(all_indices):,} cells "672 f"({n_pert_cells:,} pert + {len(control_indices):,} ctrl) ...")673 674 subset = adata[all_indices].copy()675 print(f" [{dataset_name}] Subset copied: {subset.shape} ({time.time()-t_ds:.1f}s)")676 677 pert_vals = subset.obs[pert_col].astype(str)678 is_ctrl = pert_vals.str.lower().isin(CHEM_CONTROL_PATTERNS)679 ct_vals = subset.obs[cell_col].astype(str)680 681 new_obs = pd.DataFrame(index=subset.obs.index)682 new_obs['cell_type'] = ct_vals.values683 new_obs['pert'] = pert_vals.values684 new_obs['pert_type'] = 'chemical'685 new_obs['control'] = np.where(is_ctrl.values, 'True', 'False')686 new_obs['is_normalized'] = 'False'687 new_obs['split'] = [688 'ctrl' if is_c else pert_split_map.get((str(ct).lower(), str(p)), 'train')689 for is_c, ct, p in zip(is_ctrl.values, ct_vals.values, pert_vals.values)690 ]691 692 subset.obs = new_obs693 for col in subset.obs.columns:694 subset.obs[col] = subset.obs[col].astype(str)695 processed_adatas.append(subset)696 697 if not processed_adatas:698 return None699 700 if len(processed_adatas) > 1:701 print(f" Concatenating {len(processed_adatas)} dataset subsets ({join_cell_type} join) ...")702 t_cat = time.time()703 concat_kwargs = {"join": join_cell_type}704 if join_cell_type == "outer":705 concat_kwargs["fill_value"] = fill_value706 final = ad.concat(processed_adatas, **concat_kwargs)707 if join_cell_type == "outer":708 final.uns["missing_gene_fill_value"] = float(fill_value)709 print(f" Concat done: {final.shape} ({time.time()-t_cat:.1f}s)")710 else:711 final = processed_adatas[0]712 713 de_info = build_de_info(records)714 final.uns['de_info_json'] = json.dumps(de_info)715 716 # Label distribution summary717 label_counts: dict = {}718 for entries in de_info.values():719 for e in entries:720 lbl = e.get('label', '?')721 label_counts[lbl] = label_counts.get(lbl, 0) + 1722 print(f" de_info: {len(de_info)} (cell_type|pert) keys | label dist: {label_counts}")723 print(f" Final shape: {final.shape}")724 return final, de_info, len(processed_adatas) > 1725 726 727# ============================================================================728# Core processing - chemical2genetic (jiang24 / mcf7 CRISPRi)729# ============================================================================730def process_chemical2genetic_scenario(scenario_name, records, adata, mapping):731 """Process chemical2genetic test data from jiang24."""732 print(f"\n{'='*60}")733 print(f"Processing chemical2genetic scenario: {scenario_name}")734 print(f"{'='*60}")735 736 if not records:737 print(" No records. Skipping.")738 return None739 740 records_df = pd.DataFrame(records)741 print(f" Total records: {len(records_df)}")742 743 unique_conds = records_df[['cell_type', 'pert']].drop_duplicates()744 print(f" Unique (cell_type, pert) conditions: {len(unique_conds)}")745 746 cells_to_include = []747 matched = 0748 unmatched = []749 750 for _, row in unique_conds.iterrows():751 ct = str(row['cell_type']).lower()752 pert = str(row['pert'])753 key = (ct, pert)754 if key in mapping:755 matched += 1756 cells_to_include.append({'cell_type': ct, 'pert': pert, 'indices': mapping[key]})757 else:758 unmatched.append((ct, pert, "Not found in jiang24"))759 760 print(f" Matched: {matched}/{len(unique_conds)}")761 if unmatched:762 print(f" Unmatched (first 10):")763 for ct, p, reason in unmatched[:10]:764 print(f" {ct} | {p} | {reason}")765 766 if not cells_to_include:767 print(" No cells matched. Skipping.")768 return None769 770 all_indices = []771 for item in cells_to_include:772 all_indices.extend(item['indices'])773 774 # Add control cells775 ctrl_col = adata.obs['control'].astype(str)776 ctrl_mask = ctrl_col.isin(['1', '1.0', 'True', 'true'])777 if 'condition' in adata.obs.columns:778 ctrl_mask = ctrl_mask | (adata.obs['condition'].astype(str).str.lower() == 'control')779 control_indices = np.where(ctrl_mask.values)[0].tolist()780 if control_indices:781 print(f" Adding {len(control_indices)} control cells from jiang24")782 all_indices.extend(control_indices)783 784 all_indices = sorted(set(all_indices))785 print(f" Extracting {len(all_indices)} cells")786 787 subset = adata[all_indices].copy()788 789 ctrl_col_sub = subset.obs['control'].astype(str)790 is_ctrl = ctrl_col_sub.isin(['1', '1.0', 'True', 'true'])791 if 'condition' in subset.obs.columns:792 is_ctrl = is_ctrl | (subset.obs['condition'].astype(str).str.lower() == 'control')793 794 new_obs = pd.DataFrame(index=subset.obs.index)795 new_obs['cell_type'] = subset.obs['cell_type'].astype(str).values796 new_obs['pert'] = subset.obs['gene'].astype(str).values797 new_obs['pert_type'] = 's_chemical2genetic'798 new_obs['control'] = np.where(is_ctrl.values, 'True', 'False')799 new_obs['is_normalized'] = 'True'800 801 subset.obs = new_obs802 for col in subset.obs.columns:803 subset.obs[col] = subset.obs[col].astype(str)804 805 de_info = build_de_info(records)806 subset.uns['de_info_json'] = json.dumps(de_info)807 808 print(f" Final shape: {subset.shape}")809 return subset, de_info810 811 812# ============================================================================813# Core processing - genetic_combo (norman19 / k562)814# ============================================================================815def process_genetic_combo_scenario(scenario_name, records, adata, mapping):816 """Process genetic combo test data from norman19."""817 print(f"\n{'='*60}")818 print(f"Processing genetic_combo scenario: {scenario_name}")819 print(f"{'='*60}")820 821 if not records:822 print(" No records. Skipping.")823 return None824 825 records_df = pd.DataFrame(records)826 print(f" Total records: {len(records_df)}")827 828 # Build canonical pert keys (sorted "A+B")829 def canonical_pert(pert):830 if isinstance(pert, list):831 return "+".join(sorted(pert))832 return str(pert)833 834 unique_perts = set(canonical_pert(rec['pert']) for rec in records)835 print(f" Unique perturbation conditions: {len(unique_perts)}")836 837 cells_to_include = []838 matched = 0839 unmatched = []840 841 for pert_key in unique_perts:842 if pert_key in mapping:843 matched += 1844 cells_to_include.extend(mapping[pert_key])845 else:846 parts = pert_key.split('+')847 if len(parts) == 2:848 alt_key = f"{parts[1]}+{parts[0]}"849 if alt_key in mapping:850 matched += 1851 cells_to_include.extend(mapping[alt_key])852 else:853 unmatched.append(pert_key)854 else:855 unmatched.append(pert_key)856 857 print(f" Matched: {matched}/{len(unique_perts)}")858 if unmatched:859 print(f" Unmatched (first 10): {unmatched[:10]}")860 861 if not cells_to_include:862 print(" No cells matched. Skipping.")863 return None864 865 # Add control cells866 ctrl_col = adata.obs['control'].astype(str)867 cond_col = adata.obs['condition'].astype(str)868 ctrl_mask = ctrl_col.isin(['1', '1.0', 'True', 'true']) | (cond_col.str.lower() == 'control')869 control_indices = np.where(ctrl_mask.values)[0].tolist()870 if control_indices:871 print(f" Adding {len(control_indices)} control cells from norman19")872 cells_to_include.extend(control_indices)873 874 all_indices = sorted(set(cells_to_include))875 print(f" Extracting {len(all_indices)} cells")876 877 subset = adata[all_indices].copy()878 879 ctrl_col_sub = subset.obs['control'].astype(str)880 cond_col_sub = subset.obs['condition'].astype(str)881 is_ctrl = ctrl_col_sub.isin(['1', '1.0', 'True', 'true']) | (cond_col_sub.str.lower() == 'control')882 883 new_obs = pd.DataFrame(index=subset.obs.index)884 new_obs['cell_type'] = subset.obs['cell_type'].astype(str).values885 new_obs['pert'] = cond_col_sub.values886 new_obs['pert_type'] = 's_genetic_combo'887 new_obs['control'] = np.where(is_ctrl.values, 'True', 'False')888 new_obs['is_normalized'] = 'True'889 890 subset.obs = new_obs891 for col in subset.obs.columns:892 subset.obs[col] = subset.obs[col].astype(str)893 894 de_info = build_de_info(records, pert_key_fn=lambda r: canonical_pert(r['pert']))895 subset.uns['de_info_json'] = json.dumps(de_info)896 897 print(f" Final shape: {subset.shape}")898 return subset, de_info899 900 901# ============================================================================902# Core processing - genetic2chemical (Demeo 2025 / hepg2)903# ============================================================================904def process_genetic2chemical_scenario(scenario_name, records, demeo_adata,905 demeo_mapping, demeo_ctrl_indices,906 drug_key_to_compound):907 """Process genetic2chemical test data from Demeo 2025 hepg2 dataset.908 909 The JSONL pert field contains drug-map keys (e.g. 'Doxorubicin', 'Ki8751').910 demeo_mapping: drug_key -> [indices in backed adata]911 demeo_ctrl_indices: indices of DMSO control cells912 drug_key_to_compound: drug_key -> compound_name in h5ad913 """914 print(f"\n{'='*60}")915 print(f"Processing genetic2chemical scenario: {scenario_name}")916 print(f"{'='*60}")917 918 if not records:919 print(" No records. Skipping.")920 return None921 922 records_df = pd.DataFrame(records)923 print(f" Total records: {len(records_df)}")924 925 unique_perts = records_df['pert'].unique().tolist()926 print(f" Unique drug keys needed: {unique_perts}")927 928 all_indices = list(demeo_ctrl_indices)929 matched = 0930 unmatched = []931 932 for pert in unique_perts:933 if pert in demeo_mapping:934 matched += 1935 all_indices.extend(demeo_mapping[pert])936 else:937 unmatched.append(pert)938 939 print(f" Matched: {matched}/{len(unique_perts)}")940 if unmatched:941 print(f" Unmatched: {unmatched}")942 943 if matched == 0:944 print(" No treatment cells matched. Skipping.")945 return None946 947 all_indices = sorted(set(all_indices))948 print(f" Extracting {len(all_indices)} cells (treatment + DMSO control)")949 950 # Bring subset into memory951 subset = demeo_adata[all_indices].copy()952 953 # Ensure gene_name in var954 if 'gene_name' not in subset.var.columns:955 subset.var['gene_name'] = subset.var_names.astype(str)956 957 compound_col = subset.obs['compound_name'].astype(str)958 is_ctrl = (compound_col == DEMEO_CONTROL_NAME)959 960 # Map compound_name -> drug_key for pert column961 compound_to_drug_key = {v: k for k, v in drug_key_to_compound.items()}962 pert_vals = []963 for compound in compound_col:964 if compound == DEMEO_CONTROL_NAME:965 pert_vals.append('ctrl')966 else:967 pert_vals.append(compound_to_drug_key.get(compound, compound))968 969 new_obs = pd.DataFrame(index=subset.obs.index)970 new_obs['cell_type'] = DEMEO_CELL_TYPE971 new_obs['pert'] = pert_vals972 new_obs['pert_type'] = 'chemical' # GEARS pipeline applies drug map973 new_obs['control'] = np.where(is_ctrl.values, 'True', 'False')974 new_obs['is_normalized'] = 'True' # hepg2_log1tp10k is pre-normalised975 976 subset.obs = new_obs977 for col in subset.obs.columns:978 subset.obs[col] = subset.obs[col].astype(str)979 980 de_info = build_de_info(records)981 subset.uns['de_info_json'] = json.dumps(de_info)982 983 print(f" Final shape: {subset.shape}")984 return subset, de_info985 986 987# ============================================================================988# Save helper989# ============================================================================990def save_result(result, scenario_name, output_dir):991 """Save AnnData (.h5ad) and DE info sidecar (.json)."""992 if result is None:993 print(f" {scenario_name}: no result to save.")994 return False995 996 h5ad_path = os.path.join(output_dir, f"{scenario_name}.h5ad")997 if os.path.exists(h5ad_path):998 print(f" [SKIP] {h5ad_path} already exists.")999 return False1000 1001 result_adata, de_info = result[:2]1002 1003 print(f" Saving {h5ad_path} ...")1004 result_adata.write_h5ad(h5ad_path)1005 1006 json_path = os.path.join(output_dir, f"{scenario_name}_de_info.json")1007 with open(json_path, 'w') as f:1008 json.dump(de_info, f)1009 print(f" Saved {json_path}")1010 return True1011 1012 1013def parse_args():1014 parser = argparse.ArgumentParser()1015 parser.add_argument("--join_cell_type", default="inner", choices=["inner", "outer"])1016 parser.add_argument(1017 "--fill_mask",1018 type=float,1019 default=-1.0,1020 help=(1021 "Fill value used for missing genes when --join_cell_type=outer. "1022 "Default -1 marks missing genes; use 0 to recover the previous zero-fill behavior."1023 ),1024 )1025 return parser.parse_args()1026 1027 1028def result_joined_cell_type(result):1029 return len(result) > 2 and result[2]1030 1031 1032def output_name_for(scenario_name, join_cell_type, joined_cell_type=False):1033 if join_cell_type == "inner" or not joined_cell_type:1034 return scenario_name1035 return f"{scenario_name}_{join_cell_type}"1036 1037 1038def genetic_will_join_cell_type(records, mapping_dict):1039 cell_types = set()1040 seen = set()1041 for rec in records:1042 ct = str(rec['cell_type']).lower()1043 pert = str(rec['pert'])1044 key = (ct, pert)1045 if key in seen:1046 continue1047 seen.add(key)1048 if ct in mapping_dict and pert in mapping_dict[ct]:1049 cell_types.add(ct)1050 return len(cell_types) > 11051 1052 1053def chemical_will_join_cell_type(records, mapping_dict):1054 datasets = set()1055 seen = set()1056 for rec in records:1057 ct = str(rec['cell_type'])1058 pert = str(rec['pert'])1059 key = (ct, pert)1060 if key in seen:1061 continue1062 seen.add(key)1063 dataset = CELL_TYPE_TO_CHEM_DATASET.get(ct)1064 if dataset in mapping_dict and key in mapping_dict[dataset]:1065 datasets.add(dataset)1066 return len(datasets) > 11067 1068 1069# ============================================================================1070# Main pipeline1071# ============================================================================1072def run_pipeline():1073 args = parse_args()1074 join_cell_type = args.join_cell_type1075 fill_value = args.fill_mask1076 pipeline_start = time.time()1077 os.makedirs(OUTPUT_DIR, exist_ok=True)1078 1079 # ── Step 1: Load AnnData needed for Steps 2 + 3 ──────────────────────────1080 # NOTE: jiang24 (89 GB), norman19, and demeo are loaded LAZILY in Step 3 to1081 # avoid loading all large files into memory simultaneously (OOM risk).1082 _tick("STEP 1 / 3 — Loading AnnData files (standard genetic + chemical)")1083 print("=" * 70)1084 1085 # Standard genetic datasets (hepg2, jurkat, k562, rpe1)1086 genetic_adata = {}1087 genetic_mapping = {}1088 print(f"\n[1a] Standard genetic datasets ({len(GENETIC_DATASETS)} cell types):")1089 for ct, path in GENETIC_DATASETS.items():1090 if os.path.exists(path):1091 ad_obj, mp = load_genetic_adata(ct, path)1092 if ad_obj is not None:1093 genetic_adata[ct] = ad_obj1094 genetic_mapping[ct] = mp1095 else:1096 print(f" WARNING: [{ct}] not found at {path}")1097 print(f" Loaded {len(genetic_adata)}/{len(GENETIC_DATASETS)} genetic datasets")1098 1099 # NOTE: jiang24 / norman19 / demeo loaded lazily in Step 3 (see below)1100 print(f"\n[1b] Other-test datasets (jiang24, norman19, demeo) → deferred to Step 3")1101 1102 # Chemical datasets1103 print(f"\n[1c] Chemical datasets:")1104 chem_adata = {}1105 chem_mapping = {}1106 if os.path.exists(SCIPLEX_PATH):1107 ad_obj, mp = load_chemical_adata(SCIPLEX_PATH, 'sciplex3')1108 if ad_obj is not None:1109 chem_adata['sciplex3'] = ad_obj1110 chem_mapping['sciplex3'] = mp1111 else:1112 print(f" WARNING: sciplex3 not found at {SCIPLEX_PATH}")1113 1114 if os.path.exists(OP3_PATH):1115 ad_obj, mp = load_chemical_adata(OP3_PATH, 'op3')1116 if ad_obj is not None:1117 chem_adata['op3'] = ad_obj1118 chem_mapping['op3'] = mp1119 else:1120 print(f" WARNING: op3 not found at {OP3_PATH}")1121 1122 print(f"\n Step 1 complete in {time.time()-pipeline_start:.1f}s")1123 print(f" Loaded: {len(genetic_adata)} genetic | {len(chem_adata)} chemical")1124 1125 # ── Step 2: Process train/valid scenarios ─────────────────────────────────1126 _tick("STEP 2 / 3 — Processing TRAIN/VALID scenarios")1127 print("=" * 70)1128 print(f" Train/valid JSONL dir: {TRAIN_VALID_JSONL_DIR}")1129 print(f" Scenarios to generate: {list(TRAIN_VALID_SCENARIOS.keys())}")1130 if join_cell_type == "outer":1131 print(f" Outer-join missing gene fill value: {fill_value:g}")1132 1133 train_results = {}1134 skipped_train_paths = {}1135 for scenario_name, patterns in TRAIN_VALID_SCENARIOS.items():1136 t_scen = time.time()1137 if join_cell_type == "inner":1138 output_name = output_name_for(scenario_name, join_cell_type)1139 output_path = os.path.join(OUTPUT_DIR, f"{output_name}.h5ad")1140 if os.path.exists(output_path):1141 print(f"\n [SKIP] {output_path} already exists.")1142 skipped_train_paths[scenario_name] = (output_path, False)1143 continue1144 1145 files = collect_train_valid_files(patterns)1146 if not files:1147 print(f"\n [SKIP] No JSONL files for '{scenario_name}' (patterns: {patterns})")1148 continue1149 1150 print(f"\n --- {scenario_name} ---")1151 print(f" Found {len(files)} JSONL file(s):")1152 for f in files:1153 sz = os.path.getsize(f) / (1024 * 1024)1154 print(f" {os.path.basename(f)} ({sz:.1f} MB)")1155 1156 records = load_all_records_from_files(files)1157 print(f" Total records loaded: {len(records):,}")1158 1159 if scenario_name.startswith("genetic_"):1160 joined_cell_type = genetic_will_join_cell_type(records, genetic_mapping)1161 output_name = output_name_for(scenario_name, join_cell_type, joined_cell_type)1162 output_path = os.path.join(OUTPUT_DIR, f"{output_name}.h5ad")1163 if os.path.exists(output_path):1164 print(f" [SKIP] {output_path} already exists.")1165 skipped_train_paths[scenario_name] = (output_path, joined_cell_type)1166 continue1167 result = process_genetic_scenario(1168 scenario_name,1169 records,1170 genetic_adata,1171 genetic_mapping,1172 join_cell_type,1173 fill_value=fill_value,1174 )1175 elif scenario_name.startswith("chemical_"):1176 joined_cell_type = chemical_will_join_cell_type(records, chem_mapping)1177 output_name = output_name_for(scenario_name, join_cell_type, joined_cell_type)1178 output_path = os.path.join(OUTPUT_DIR, f"{output_name}.h5ad")1179 if os.path.exists(output_path):1180 print(f" [SKIP] {output_path} already exists.")1181 skipped_train_paths[scenario_name] = (output_path, joined_cell_type)1182 continue1183 result = process_chemical_scenario(1184 scenario_name,1185 records,1186 chem_adata,1187 chem_mapping,1188 join_cell_type,1189 fill_value=fill_value,1190 )1191 else:1192 print(f" [SKIP] Unknown scenario type: {scenario_name}")1193 continue1194 1195 if result is not None:1196 output_name = output_name_for(1197 scenario_name, join_cell_type, result_joined_cell_type(result))1198 if save_result(result, output_name, OUTPUT_DIR):1199 train_results[scenario_name] = result1200 print(f" '{output_name}' done in {time.time()-t_scen:.1f}s")