File size: 2,278 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 | """the pivot model: perturbation encoder ψ_η + flow map X_θ, plus the
inverse-design jacobian used by reward guidance (Alg 4)."""
from __future__ import annotations
import numpy as np
import torch
import torch.nn as nn
from pivot.models.encoders import PerturbationEncoder, build_pert_tensors
from pivot.models.flow_map import FlowMap
class PIVOT(nn.Module):
def __init__(
self,
d_state: int,
n_genes: int,
n_ops: int,
n_perts: int,
d_pert: int = 64,
hidden: int = 512,
depth: int = 4,
rep_mode: str = "gene_op",
gene_pathway=None,
n_pathways: int = 0,
dropout: float = 0.0,
):
super().__init__()
self.encoder = PerturbationEncoder(
n_genes,
n_ops,
n_perts,
emb_dim=d_pert,
mode=rep_mode,
gene_pathway=gene_pathway,
n_pathways=n_pathways,
)
self.flow = FlowMap(
d_state, d_pert, hidden=hidden, depth=depth, dropout=dropout
)
self.d_state, self.d_pert = d_state, d_pert
def encode(self, g, o, mask, pid):
return self.encoder(g, o, mask, pid)
def endpoint_from_e(self, c0, e):
return self.flow.endpoint(c0, e)
def endpoint_from_pert(self, c0, g, o, mask, pid):
e = self.encode(g, o, mask, pid)
if e.shape[0] == 1 and c0.shape[0] > 1:
e = e.expand(c0.shape[0], -1)
return self.flow.endpoint(c0, e)
# --- inverse-design jacobian (Alg 4 building block) ---
def endpoint_jacobian_e(self, c0_single, e_single):
"""∇_e X_{0,1}(c0, e) for a single (c0, e). returns (d_state, d_pert)."""
e = e_single.detach().clone().requires_grad_(True)
c0 = c0_single.unsqueeze(0)
def f(ev):
return self.flow.endpoint(c0, ev.unsqueeze(0)).squeeze(0)
return torch.autograd.functional.jacobian(f, e, create_graph=False)
@torch.no_grad()
def candidate_embeddings(model: PIVOT, data, labels, device) -> np.ndarray:
"""ψ(u) for every candidate perturbation label -> (n_cand, d_pert)."""
g, o, mask, pid = build_pert_tensors(data, labels, device=device)
return model.encode(g, o, mask, pid).cpu().numpy()
|