File size: 5,488 Bytes
6fa9282 | 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 128 129 130 131 132 133 134 135 136 137 | """Reference predictors fitted with the same training cells and feature basis."""
from __future__ import annotations
import copy
import numpy as np
import torch
from sklearn.linear_model import Ridge
from pivot.training.train import make_model, TrainConfig, validation_loss
from pivot.models.encoders import build_pert_tensors
from pivot.utils.common import set_seed
class Baseline:
"""Population predictor. `predict(c0, label)` returns (n_controls, d)."""
def __init__(
self, data, kind="ridge", alpha=1.0, epochs=60, hidden=512, seed=0, device="cpu"
):
if kind not in (
"mean_control",
"average_effect",
"additive",
"ridge",
"endpoint_mlp",
"conditional_mlp",
):
raise ValueError(kind)
self.data = data
self.kind = kind
self.device = device
tr = data.indices("train", False)
ctrl = data.indices("train", True)
self.control = data.emb[ctrl].mean(0)
self.effects = {
p: data.emb[np.intersect1d(tr, data.pert_to_idx[p])].mean(0) - self.control
for p in data.labels("train")
}
self.avg = np.mean(list(self.effects.values()), axis=0)
self.training_info = {
"kind": kind,
"alpha": alpha,
"seed": seed,
"train_cell_count": len(tr),
}
if kind == "ridge":
labels = list(self.effects)
self.reg = Ridge(alpha=alpha, fit_intercept=False).fit(
self.features(labels), np.stack(list(self.effects.values()))
)
if kind in ("endpoint_mlp", "conditional_mlp"):
set_seed(seed)
torch.set_num_threads(4)
cfg = TrainConfig(hidden=hidden, depth=4, seed=seed, device=device)
self.model = make_model(data, cfg)
opt = torch.optim.AdamW(self.model.parameters(), lr=1e-3, weight_decay=1e-5)
rng = np.random.default_rng(seed)
best = float("inf")
beststate = None
history = []
for epoch in range(epochs):
self.model.train()
for ids in np.array_split(
rng.permutation(tr), max(1, int(np.ceil(len(tr) / 1024)))
):
c = data.emb[rng.choice(ctrl, len(ids))]
if kind == "endpoint_mlp":
c = np.broadcast_to(self.control, c.shape).copy()
c = torch.as_tensor(c, device=device)
y = torch.as_tensor(data.emb[ids], device=device)
labels = data.obs.iloc[ids].perturbation.tolist()
g, o, m, p = build_pert_tensors(data, labels, device)
pred = self.model.endpoint_from_pert(c, g, o, m, p)
loss = (pred - y).square().sum(-1).mean()
opt.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5.0)
opt.step()
self.model.eval()
vals = []
vc = data.emb[data.indices("val", True)][:128]
for p in data.labels("val"):
ids = np.intersect1d(
data.indices("val", False), data.pert_to_idx[p]
)
vals.append(
np.mean(
(self.predict(vc, p).mean(0) - data.emb[ids].mean(0)) ** 2
)
)
score = float(np.mean(vals))
history.append(score)
if score < best:
best = score
beststate = copy.deepcopy(self.model.state_dict())
self.model.load_state_dict(beststate)
self.training_info.update(
validation_mse=history,
best_validation_mse=best,
epochs=epochs,
hidden=hidden,
)
def features(self, labels):
a = np.zeros((len(labels), len(self.data.genes_vocab)), np.float32)
for i, p in enumerate(labels):
for g in self.data.parse(p):
a[i, self.data.gene_to_id[g]] = 1.0
return a
def predict(self, c0, label):
if self.kind in ("endpoint_mlp", "conditional_mlp"):
c = np.asarray(c0, dtype=np.float32)
if self.kind == "endpoint_mlp":
c = np.broadcast_to(self.control, c.shape).copy()
g, o, m, p = build_pert_tensors(self.data, [label], self.device)
with torch.no_grad():
return (
self.model.endpoint_from_pert(
torch.as_tensor(c, device=self.device), g, o, m, p
)
.cpu()
.numpy()
)
if self.kind == "mean_control":
delta = np.zeros(self.data.d)
elif self.kind == "average_effect":
delta = self.avg
elif self.kind == "additive":
# Unsupported single genes receive the training-average effect.
delta = sum(
(self.effects.get(g, self.avg) for g in self.data.parse(label)),
start=np.zeros(self.data.d),
)
else:
delta = self.reg.predict(self.features([label]))[0]
return np.asarray(c0) + delta
|