| """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": |
| |
| 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 |
|
|
