| """Calibration, retrospective evaluation, and catalog generation with locked inputs.""" |
| from pathlib import Path |
| import json |
| import time |
| import joblib |
| import numpy as np |
| import pandas as pd |
| from .io import write_json, file_hash |
| from .models import check_feature_space, predict |
| from .uncertainty import target_draws, joint_losses, cross_block_losses, panel_score, calibration_radius |
| from .optimization import solve_regret, generate_baseline |
| from .metrics import intervention_metrics, grouped_summary, paired_difference |
| from .chemistry import verify_catalog, fingerprints, max_tanimoto |
|
|
| BASELINES = ["uniform", "mean", "gibbs", "marginal", "absolute", "thompson", "finite-minimax"] |
|
|
| def action_groups(data, indices, min_blocks=1): |
| rows, groups = [], [] |
| for _, g in data.obs.loc[indices].groupby(["molecule_id", "dose_um", "context"], sort=True): |
| g = g.sort_values("block").drop_duplicates("block") |
| if len(g) < min_blocks: continue |
| rows.append(g.iloc[0][["molecule_id", "smiles", "dose_um", "context"]].to_dict()) |
| groups.append(g.index.to_list()) |
| return pd.DataFrame(rows), groups |
|
|
| def calibrate(data, model, output, coverage=.9, panel_molecules=6, scenarios=32, seed=0, |
| ablation="none", embeddings=None): |
| """Estimate beta from held-out compound panels with two measurement blocks. |
| |
| Pool repeated contexts by their maximum panel score within each compound |
| group. Save the panel scores and the associated model identity. |
| """ |
| check_feature_space(model, data) |
| obs = data.obs |
| cal = obs[obs.split == "calibration"] |
| molecules = np.sort(cal.molecule_id.unique()) |
| if set(molecules) & set(model["train_molecules"]): raise ValueError("Training/calibration leakage") |
| if len(molecules) < 3: raise ValueError("Calibration needs at least three held-out molecules") |
| rng = np.random.default_rng(seed) |
| rng.shuffle(molecules) |
| records, skips = [], [] |
| |
| for group_id, start in enumerate(range(0, len(molecules), panel_molecules)): |
| group = molecules[start:start+panel_molecules] |
| if len(group) < 3: continue |
| for context in sorted(cal.context.unique()): |
| panel = cal[(cal.molecule_id.isin(group)) & (cal.context == context)] |
| if panel.molecule_id.nunique() < 3: continue |
| target_molecule = panel.molecule_id.iloc[0] |
| target_rows = panel[panel.molecule_id == target_molecule] |
| |
| available_doses = np.sort(target_rows.dose_um.unique()) |
| dose = available_doses[len(available_doses)//2] |
| target_rows = target_rows[target_rows.dose_um == dose].sort_values("block").drop_duplicates("block") |
| actions, groups = action_groups(data, panel[panel.molecule_id != target_molecule].index, 2) |
| if len(actions) < 2 or len(target_rows) < 2: continue |
| try: |
| observed = cross_block_losses(data, groups, target_rows.index.to_list(), model["weight"]) |
| except ValueError as e: |
| skips.append({"group": group_id, "context": context, "reason": str(e)}); continue |
| targets, controls = target_draws(data, target_rows.index[:2], scenarios, rng) |
| joint = joint_losses(model, actions, targets, controls, seed+group_id, |
| embeddings, floor=ablation != "no-floor", ablation=ablation) |
| records.append({"group": group_id, "context": context, "score": panel_score(joint, observed), |
| "actions": len(actions), "target_molecule": target_molecule}) |
| if not records: |
| raise ValueError("No calibration panel has independent target/candidate blocks and distinct control pools") |
| frame = pd.DataFrame(records) |
| scores = frame.groupby("group").score.max().to_numpy() |
| beta = calibration_radius(scores, coverage) |
| output = Path(output); output.mkdir(parents=True, exist_ok=True) |
| frame.to_csv(output/"calibration_panels.csv", index=False) |
| result = {"beta": beta, "coverage": coverage, "independent_groups": len(scores), |
| "scores": scores.tolist(), "ablation": ablation, "seed": seed, |
| "calibration_molecules": molecules.tolist(), "feature_space_id": model["feature_space_id"], |
| "model_training_molecules": model["train_molecules"], "model_digest": joblib.hash(model), |
| "status": "empirical_panel_quantile", |
| "latent_ellipsoid_coverage_guaranteed": False, "skipped": skips} |
| write_json(output/"calibration.json", result) |
| return result |
|
|
| def evaluate(data, model, calibration, output, split="test", scenarios=32, tau=.05, |
| tolerance=1., seed=0, max_queries=30, methods=None, embeddings=None, |
| metric_data=None, ablation="none", external_track="none"): |
| """Evaluate action distributions against separate measured response blocks. |
| |
| Each query excludes its source molecule at every dose from the candidate |
| pool. Return per-query metrics and save probabilities and grouped summaries. |
| """ |
| check_feature_space(model, data) |
| held = data.obs[data.obs.split == split] |
| protected = set(model["train_molecules"]) | set(calibration.get("calibration_molecules", [])) |
| if external_track == "none" and set(held.molecule_id) & protected: |
| raise ValueError("Evaluation molecules overlap training or calibration") |
| if external_track != "none": |
| source = data.metadata.get("study") |
| if not source or not model.get("study") or source == model["study"]: |
| raise ValueError("External transfer requires distinct recorded study identifiers") |
| overlap = set(held.molecule_id) & set(model["train_molecules"]) |
| if external_track == "unseen" and overlap: |
| raise ValueError("Unseen external track includes trained molecular structures") |
| if external_track == "shared" and set(held.molecule_id) != overlap: |
| raise ValueError("Shared external track includes structures absent from response-model training") |
| if calibration.get("model_digest") != joblib.hash(model): |
| raise ValueError("Calibration was produced by a different fitted model") |
| if calibration.get("feature_space_id") != model["feature_space_id"]: |
| raise ValueError("Calibration feature space differs from model") |
| methods = methods or BASELINES+["remedi"] |
| metric_data = data if metric_data is None else metric_data |
| keys = ["molecule_id", "dose_um", "context", "block", "control_id"] |
| if not data.obs[keys].equals(metric_data.obs[keys]): raise ValueError("Evaluation representation must preserve condition alignment") |
| if metric_data is data: weight = model["weight"] |
| else: |
| train = metric_data.response[metric_data.obs.split.to_numpy() == "train"] |
| if not len(train): raise ValueError("Independent evaluation space needs training-derived scaling") |
| weight = 1/np.maximum(train.var(0), 1e-6); weight /= weight.sum() |
| rng = np.random.default_rng(seed) |
| query_groups = list(held.groupby(["molecule_id", "dose_um", "context"], sort=True)) |
| rng.shuffle(query_groups) |
| rows, choices, skipped = [], [], [] |
| query_count = 0 |
| for key, target in query_groups: |
| if query_count >= max_queries: break |
| target = target.sort_values("block").drop_duplicates("block") |
| if len(target) < 2: continue |
| query_idx, eval_idx = int(target.index[0]), int(target.index[1]) |
| if target.iloc[0].control_id == target.iloc[1].control_id: |
| skipped.append({"target": key, "reason": "query/evaluation controls shared"}); continue |
| pool = held[(held.context == key[2]) & (held.molecule_id != key[0])] |
| pool = pool[pool.control_id != data.obs.iloc[query_idx].control_id] |
| actions, groups = action_groups(data, pool.index) |
| if len(actions) < 2: continue |
| |
| targets, controls = target_draws(data, [query_idx], scenarios, rng, |
| sampling=ablation != "no-target-sampling") |
| joint = joint_losses(model, actions, targets, controls, seed+query_count, |
| embeddings, floor=ablation != "no-floor", ablation=ablation) |
| observed = np.array([np.sum((metric_data.response[g].mean(0)-metric_data.response[eval_idx])**2*weight) for g in groups]) |
| for name in methods: |
| begin = time.perf_counter() |
| diagnostics = {} |
| if name == "remedi": |
| decision = solve_regret(joint.mean, joint.factor, joint.diagonal, |
| calibration["beta"], tau) |
| p = decision.probabilities |
| diagnostics = {"solver_status": decision.status, "solver_violation": decision.constraint_violation, |
| "estimated_worst_regret": decision.worst_regret} |
| else: |
| p = generate_baseline(name, joint.mean, joint.factor, joint.diagonal, joint.scenarios, |
| calibration["beta"], tau, seed+query_count) |
| scores = intervention_metrics(p, observed, tolerance, seed+query_count) |
| molecule_p = pd.Series(p).groupby(actions.molecule_id.to_numpy()).sum().to_numpy() |
| scores["unique_molecules_10"] = float(np.sum(1-(1-molecule_p)**10)) |
| dose_regret = np.zeros(len(actions)) |
| for _, g in actions.groupby("molecule_id"): |
| ii = g.index.to_numpy(); dose_regret[ii] = observed[ii]-observed[ii].min() |
| row = {"query_id": query_count, "target_molecule": key[0], "target_dose_um": key[1], |
| "context": key[2], "method": name, "ablation": ablation, "candidates": len(actions), |
| "dose_regret": float(p@dose_regret), "seconds": time.perf_counter()-begin, |
| **scores, **diagnostics} |
| rows.append(row) |
| for i in np.flatnonzero(p > 1e-8): |
| choices.append({"query_id": query_count, "method": name, **actions.iloc[i].to_dict(), |
| "probability": float(p[i]), "observed_loss": float(observed[i])}) |
| query_count += 1 |
| if not rows: raise ValueError("No independent query/evaluation pairs with eligible alternatives") |
| output = Path(output); output.mkdir(parents=True, exist_ok=True) |
| frame = pd.DataFrame(rows); frame.to_csv(output/"per_query.csv", index=False) |
| pd.DataFrame(choices).to_csv(output/"action_probabilities.csv", index=False) |
| summary = pd.concat([grouped_summary(frame, m, seed=seed) for m in ["endpoint_loss", "regret", "success", "best_of_10", "dose_regret", "entropy", "unique_molecules_10"]]) |
| summary.to_csv(output/"summary.csv", index=False) |
| differences = [paired_difference(frame, "remedi", n, seed=seed) for n in methods if n != "remedi"] if "remedi" in methods else [] |
| write_json(output/"evaluation.json", {"queries": query_count, "target_molecules": frame.target_molecule.nunique(), |
| "candidate_split": split, "external_track": external_track, "seed": seed, "tau": tau, "tolerance": tolerance, |
| "beta": calibration["beta"], "scenarios": scenarios, "paired_differences": differences, |
| "metric_space": metric_data.metadata.get("cell_features"), "skipped": skipped, |
| "biological_validation": "retrospective measured endpoints; prototype unless complete registered benchmark"}) |
| return frame |
|
|
| def generate_catalog(data, model, calibration, catalog, target_indices, output, doses=(.05,.5,5.), |
| scenarios=32, tau=.05, shortlist=512, draws=20, seed=0, embeddings=None): |
| """Sample catalog IDs, structures, and concentrations for one target endpoint. |
| |
| target_indices identify replicates of one condition in data.obs. Expand |
| catalog records over doses in uM, shortlist if needed, then optimize and |
| save the distribution together with the sampled molecule-dose pairs. |
| """ |
| check_feature_space(model, data) |
| if calibration.get("model_digest") != joblib.hash(model): |
| raise ValueError("Calibration was produced by a different fitted model") |
| if not target_indices or min(target_indices) < 0 or max(target_indices) >= len(data.obs): |
| raise ValueError("Target rows must be valid condition indices") |
| query = data.obs.iloc[target_indices] |
| if len(query[["molecule_id", "dose_um", "context"]].drop_duplicates()) != 1: |
| raise ValueError("Target rows must describe one molecule, dose, and context") |
| if shortlist < 2 or draws < 1 or scenarios < 2: |
| raise ValueError("Require shortlist >= 2, draws >= 1, and scenarios >= 2") |
| catalog = verify_catalog(catalog) |
| if not len(catalog): raise ValueError("Verified catalog is empty") |
| doses = np.asarray(doses, float) |
| if not np.isfinite(doses).all() or (doses <= 0).any(): raise ValueError("Doses must be finite positive values") |
| actions = catalog.merge(pd.DataFrame({"dose_um": np.unique(doses)}), how="cross") |
| targets, controls = target_draws(data, target_indices, scenarios, np.random.default_rng(seed)) |
| source_count = len(actions) |
| |
| if len(actions) > shortlist: |
| mean_target, mean_control = targets.mean(0), controls.mean(0) |
| rough = [] |
| for start in range(0, len(actions), 1024): |
| a = actions.iloc[start:start+1024] |
| f = predict(model, a.smiles, a.dose_um, np.repeat(mean_control[None, :], len(a), axis=0), embeddings).mean(0) |
| rough.extend(((f-mean_target)**2*model["weight"]).sum(1)) |
| rank = np.argsort(rough, kind="stable") |
| low = rank[:max(1, shortlist*3//4)] |
| remaining = np.setdiff1d(np.arange(len(actions)), low) |
| extra = [] |
| candidate_fp = fingerprints(actions.iloc[remaining].smiles) |
| similarity = max_tanimoto(candidate_fp, fingerprints(actions.iloc[low].smiles)) |
| for _ in range(min(shortlist-len(low), len(remaining))): |
| j = int(np.argmin(similarity)) |
| extra.append(remaining[j]) |
| similarity = np.maximum(similarity, max_tanimoto(candidate_fp, candidate_fp[j:j+1])) |
| similarity[j] = np.inf |
| extra = np.asarray(extra, dtype=int) |
| actions = actions.iloc[np.concatenate([low, extra])].reset_index(drop=True) |
| joint = joint_losses(model, actions, targets, controls, seed, embeddings) |
| result = solve_regret(joint.mean, joint.factor, joint.diagonal, calibration["beta"], tau) |
| frame = actions.copy(); frame["probability"] = result.probabilities; frame["predicted_loss"] = joint.mean |
| frame["chemical_distance"] = 1-max_tanimoto(fingerprints(frame.smiles), model["train_fingerprints"]) |
| selected = np.random.default_rng(seed).choice(len(frame), draws, p=result.probabilities) |
| output = Path(output); output.mkdir(parents=True, exist_ok=True) |
| frame.to_csv(output/"distribution.csv", index=False) |
| frame.iloc[selected].to_csv(output/"generated_pairs.csv", index=False) |
| write_json(output/"generation.json", {"source_actions": source_count, "shortlist_actions": len(frame), |
| "catalog_membership": "exact canonical structure match to supplied verified records", |
| "seconds": result.seconds, "solver_status": result.status, "constraint_violation": result.constraint_violation, |
| "estimated_worst_regret": result.worst_regret, "biological_effect_measured": False, |
| "dose_um": doses.tolist(), "seed": seed, "tau": tau, "beta": calibration["beta"]}) |
| return frame |
|
|
