| """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") |
| |
| 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) |
| |
| factor = (losses-mean).T/np.sqrt(s-1) |
| diagonal = np.zeros(n) |
| if floor: |
| residual = model["residuals"] |
| |
| 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)) |
| |
| 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)") |
| |
| return float(np.quantile(scores, coverage, method="higher")) |
|
|
