"""Coherent target draws, joint loss covariance, and observed-comparison calibration.""" from dataclasses import dataclass import numpy as np from .models import design_features, _predict_head from .chemistry import encode, fingerprints, max_tanimoto @dataclass class JointLoss: mean: np.ndarray factor: np.ndarray diagonal: np.ndarray scenarios: np.ndarray predictions: np.ndarray def target_draws(data, indices, count, rng, sampling=True): """Draw target-response and control means, each with shape (count, features). Sample blocks with replacement. Reuse control noise within each control pool so the response difference and context retain their shared variation. """ indices = np.asarray(indices, dtype=int) if not len(indices): raise ValueError("A query needs target observations") draws, controls = [], [] for _ in range(count): chosen = rng.choice(indices, len(indices), replace=True) if sampling else indices control_noise = {} treated_noise = {} rr, cc = [], [] for i in chosen: cid = data.obs.iloc[i].control_id if cid not in control_noise: control_noise[cid] = rng.normal(size=data.response.shape[1])*np.sqrt(data.control_var[i]) if sampling else np.zeros(data.response.shape[1]) cn = control_noise[cid] if i not in treated_noise: treated_noise[i] = rng.normal(size=data.response.shape[1])*np.sqrt(data.treated_var[i]) if sampling else 0 tn = treated_noise[i] rr.append(data.response[i]+tn-cn) cc.append(data.control[i]+cn) draws.append(np.mean(rr, axis=0)); controls.append(np.mean(cc, axis=0)) return np.asarray(draws), np.asarray(controls) def joint_losses(model, actions, targets, controls, seed=0, embeddings=None, floor=True, ablation="none"): """Estimate joint endpoint losses for a common set of candidate actions. targets and controls are aligned (S, D) draws. Return loss scenarios (S, N), their mean (N,), covariance factor (N, S), and residual floor (N,). """ rng = np.random.default_rng(seed) n, s = len(actions), len(targets) if s < 2: raise ValueError("At least two scenarios are required") if np.asarray(controls).shape != np.asarray(targets).shape: raise ValueError("Target and control draws must align") # A scenario uses one head AND one query realization across every action. losses, predictions = [], [] chemical = encode(actions.smiles, model["encoder"], embeddings or model["embeddings"]) for i in range(s): x = design_features(chemical, actions.dose_um, np.repeat(controls[i:i+1], n, axis=0), model["kind"]) f = _predict_head(model["heads"][int(rng.integers(len(model["heads"])))], x) losses.append(((f-targets[i])**2*model["weight"]).sum(1)) predictions.append(f) losses, predictions = np.asarray(losses), np.asarray(predictions) if ablation == "permuted": losses = np.stack([rng.permutation(losses[:, a]) for a in range(n)], axis=1) mean = losses.mean(0) # This factor reproduces the sample covariance without a dense N-by-N matrix. factor = (losses-mean).T/np.sqrt(s-1) diagonal = np.zeros(n) if floor: residual = model["residuals"] # Empirical response residuals induce query-specific scalar loss errors. f = predictions.mean(0); t = targets.mean(0) distance = 1-max_tanimoto(fingerprints(actions.smiles), model["train_fingerprints"]) bins = np.asarray([.25, .5, .75]) residual_bins = np.digitize(model["residual_distance"], bins) for a in range(n): stratum = np.flatnonzero(residual_bins == np.digitize(distance[a], bins)) # Blend sparse chemical-distance strata with the pooled residual sample. shrinkage = len(stratum)/(len(stratum)+20.) indices = rng.choice(len(residual), size=min(128, len(residual)), replace=True) use_stratum = rng.random(len(indices)) < shrinkage if len(stratum): indices[use_stratum] = rng.choice(stratum, int(use_stratum.sum())) sampled = residual[indices] perturbed = (((f[a]+sampled-t)**2-(f[a]-t)**2)*model["weight"]).sum(1) diagonal[a] = max(float(np.var(perturbed, ddof=1)), 1e-10) if ablation == "diagonal": diagonal += np.sum(factor**2, axis=1) factor = np.zeros((n, 0)) return JointLoss(mean, factor, diagonal, losses, predictions.mean(0)) def cross_block_losses(data, action_indices, target_indices, weight): """Unbiased cross-product losses only for two distinct supplied blocks.""" target_indices = list(target_indices) if len(target_indices) < 2: raise ValueError("Two target blocks are needed") u, v = target_indices[:2] if data.obs.iloc[u].block == data.obs.iloc[v].block: raise ValueError("Target blocks are not distinct") left_controls = {data.obs.iloc[u].control_id}; right_controls = {data.obs.iloc[v].control_id} values = [] for indices in action_indices: if len(indices) < 2: raise ValueError("Two blocks per candidate are needed") a, b = indices[:2] if data.obs.iloc[a].block == data.obs.iloc[b].block: raise ValueError("Candidate blocks are not distinct") left_controls.add(data.obs.iloc[a].control_id); right_controls.add(data.obs.iloc[b].control_id) values.append(np.sum((data.response[a]-data.response[u])*(data.response[b]-data.response[v])*weight)) if left_controls & right_controls: raise ValueError("Cross-block comparisons reuse a control pool across blocks") return np.asarray(values) def panel_score(joint, observed_losses, epsilon=1e-8): error = joint.mean-np.asarray(observed_losses) cov_diag = np.sum(joint.factor**2, axis=1)+joint.diagonal variance = cov_diag[:, None]+cov_diag[None, :]-2*joint.factor@joint.factor.T np.fill_diagonal(variance, 0) standardized = np.abs(error[:, None]-error[None, :])/(np.sqrt(np.maximum(variance, 0))+epsilon) np.fill_diagonal(standardized, 0) return float(standardized.max()) def calibration_radius(scores, coverage=.9): scores = np.asarray(scores, float) if not 0 < coverage < 1 or not len(scores) or not np.isfinite(scores).all(): raise ValueError("Calibration needs finite grouped scores and coverage in (0,1)") # Use one score per molecule group to estimate the held-out panel quantile. return float(np.quantile(scores, coverage, method="higher"))