File size: 5,472 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 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | """inference procedures, algorithms 2-6.
Alg 2 forward_predict predict endpoint, aggregate population
Alg 3 endpoint_ranking score = reward of predicted endpoint, return topk
Alg 4 reward_guidance jacobian pull-back ascent in embedding space
Alg 5 project_and_rerank project e* to admissible candidates, rerank k-nn
Alg 6 greedy_combinatorial sequentially add the best gene (mean-pool embed)
"""
from __future__ import annotations
import numpy as np
import torch
from pivot.models.encoders import build_pert_tensors
from pivot.models.pivot import PIVOT, candidate_embeddings
def encode_label(model: PIVOT, data, label: str, device) -> torch.Tensor:
g, o, mask, pid = build_pert_tensors(data, [label], device=device)
return model.encode(g, o, mask, pid) # (1, m)
def encode_gene_set(model: PIVOT, data, genes: list[str], device) -> torch.Tensor:
"""encode an arbitrary gene set via the dataset operation (mean-pool)."""
label = data.sep.join(genes) if genes else data.control_label
return encode_label(model, data, label, device)
@torch.no_grad()
def forward_predict(model: PIVOT, c0: torch.Tensor, e_u: torch.Tensor) -> torch.Tensor:
"""Alg 2. c0: (N,d) control embeddings; e_u: (1,m) or (N,m). returns endpoint (N,d)."""
if e_u.shape[0] == 1 and c0.shape[0] > 1:
e_u = e_u.expand(c0.shape[0], -1)
return model.endpoint_from_e(c0, e_u)
@torch.no_grad()
def endpoint_ranking(
model: PIVOT,
data,
candidate_labels,
c0: torch.Tensor,
reward,
topk: int | None = None,
device="cpu",
chunk: int = 4,
):
"""Alg 3. score each candidate by reward of its predicted endpoint population.
batched: all candidates x control cells go through the flow map in one
(chunked) forward pass, then per-row rewards are averaged per candidate. works
for any per-row reward (centroid/cosine). returns list of (label, score) desc.
"""
E = torch.as_tensor(
candidate_embeddings(model, data, candidate_labels, device), device=device
)
C, N = E.shape[0], c0.shape[0]
scores = np.empty(C, dtype=np.float64)
# chunk over candidates to bound memory (c*n rows per chunk)
cands_per_chunk = max(1, chunk * 1024 // max(N, 1))
for s in range(0, C, cands_per_chunk):
e = E[s : s + cands_per_chunk]
c = e.shape[0]
c0_rep = c0.unsqueeze(0).expand(c, N, -1).reshape(c * N, -1)
e_rep = e.unsqueeze(1).expand(c, N, -1).reshape(c * N, -1)
chat = model.endpoint_from_e(c0_rep, e_rep)
if reward.kind in ("mmd", "wasserstein"):
pops = chat.view(c, N, -1)
r = torch.stack([reward(pop).mean() for pop in pops])
else:
r = reward(chat).view(c, N).mean(1)
scores[s : s + c] = r.detach().cpu().numpy()
order = np.argsort(-scores)
ranked = [(candidate_labels[i], float(scores[i])) for i in order]
return ranked if topk is None else ranked[:topk]
def reward_guidance(
model: PIVOT,
c0: torch.Tensor,
reward,
e_init: torch.Tensor,
steps: int = 25,
step_size: float = 0.5,
eps: float = 1e-8,
normalize: bool = True,
) -> torch.Tensor:
"""Alg 4. gradient ascent on reward in embedding space via the flow-map jacobian.
e_{l+1} = e_l + gamma * g_e/(norm(g_e)+eps), g_e from autograd."""
e = e_init.detach().clone()
for _ in range(steps):
e.requires_grad_(True)
chat = model.endpoint_from_e(c0, e.expand(c0.shape[0], -1))
r = reward(chat).mean()
(g_e,) = torch.autograd.grad(r, e)
with torch.no_grad():
step = g_e / (g_e.norm() + eps) if normalize else g_e
e = e + step_size * step
e = e.detach()
return e
@torch.no_grad()
def project_and_rerank(
model: PIVOT,
data,
candidate_labels,
e_star: torch.Tensor,
c0: torch.Tensor,
reward,
k_nearest: int = 10,
topk: int = 5,
device="cpu",
):
"""Alg 5. project e* to the k nearest admissible candidates, rerank by endpoint reward."""
E = torch.as_tensor(
candidate_embeddings(model, data, candidate_labels, device), device=device
)
d = torch.cdist(e_star.view(1, -1), E).squeeze(0)
knn = torch.argsort(d)[:k_nearest].cpu().numpy()
sub = [candidate_labels[i] for i in knn]
ranked = endpoint_ranking(model, data, sub, c0, reward, device=device)
return ranked[:topk]
@torch.no_grad()
def greedy_combinatorial(
model: PIVOT,
data,
gene_pool,
c0: torch.Tensor,
reward,
max_size: int = 2,
min_gain: float = 1e-4,
device="cpu",
):
"""Alg 6. greedily add the gene that most improves the reward of the mean-pooled endpoint."""
chosen: list[str] = []
best_score = -np.inf
history = []
while len(chosen) < max_size:
remaining = [g for g in gene_pool if g not in chosen]
if not remaining:
break
# score all candidate gene-sets {chosen + g} in one batched pass
labels = [data.sep.join(chosen + [g]) for g in remaining]
ranked = endpoint_ranking(model, data, labels, c0, reward, device=device)
best_label, best_local = ranked[0]
best_gene = remaining[labels.index(best_label)]
if best_local - best_score < min_gain:
break
chosen.append(best_gene)
best_score = best_local
history.append((list(chosen), best_score))
return chosen, best_score, history
|