File size: 1,581 Bytes
3b4941f | 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 | """uniform predictor interface so pivot and every baseline run through the same
forward/inverse code. a predictor maps (perturbation label, control-cell embeddings
c0) -> predicted post-perturbation population ĉ1 in embedding space.
"""
from __future__ import annotations
import numpy as np
import torch
from src.evaluation import inference as inf
class PivotPredictor:
name = "PIVOT"
def __init__(self, model, data, device):
self.model, self.data, self.device = model, data, device
model.eval()
def population(self, label: str, c0: np.ndarray) -> np.ndarray:
c0t = torch.as_tensor(c0, dtype=torch.float32, device=self.device)
e = inf.encode_label(self.model, self.data, label, self.device)
return inf.forward_predict(self.model, c0t, e).cpu().numpy()
def population_from_genes(self, genes: list[str], c0: np.ndarray) -> np.ndarray:
c0t = torch.as_tensor(c0, dtype=torch.float32, device=self.device)
e = inf.encode_gene_set(self.model, self.data, genes, self.device)
return inf.forward_predict(self.model, c0t, e).cpu().numpy()
class BaselinePredictor:
def __init__(self, baseline):
self.bl = baseline
self.name = baseline.name
def population(self, label: str, c0: np.ndarray) -> np.ndarray:
return self.bl.predict_endpoint(label, c0)
def population_from_genes(self, genes: list[str], c0: np.ndarray) -> np.ndarray:
label = self.bl.data.sep.join(genes) if genes else self.bl.data.control_label
return self.bl.predict_endpoint(label, c0)
|