src/remedi/pipeline.py

"""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 = [], []
    # A group is disjoint in molecule identity; repeated contexts contribute one maximum score.
    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]
            # One dose per group/context prevents dependent target-dose panels being counted independently.
            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
        # All candidate responses below remain evaluation-only; no observed value enters joint_losses.
        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)
    # Shared bounded shortlist. This stage is explicitly outside the finite-pool regret guarantee.
    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