File size: 6,634 Bytes
3f98d52 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """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"))
|