File size: 3,793 Bytes
b67979a | 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 | """perturbation encoder ψ_η plus a builder that turns perturbation labels into
padded (gene_id, op_id, mask) tensors.
ψ separates gene identity from operation and is permutation-invariant over the
gene set (mean pooling), so one encoder handles single-gene and combinatorial
perturbations:
e_u = (1/M) Σ_j ( W_g[g_j] + W_o[o_j] )
rep modes: gene_op (default), op_only, gene_only, random_id, gene_pathway_op.
"""
from __future__ import annotations
import numpy as np
import torch
import torch.nn as nn
REP_MODES = ("gene_op", "op_only", "gene_only", "random_id", "gene_pathway_op")
class PerturbationEncoder(nn.Module):
def __init__(
self,
n_genes: int,
n_ops: int,
n_perts: int,
emb_dim: int = 64,
mode: str = "gene_op",
gene_pathway: np.ndarray | None = None, # (n_genes,) cluster id, for gene_pathway_op mode
n_pathways: int = 0,
):
super().__init__()
assert mode in REP_MODES, mode
self.mode = mode
self.emb_dim = emb_dim
self.W_g = nn.Embedding(n_genes, emb_dim)
self.W_o = nn.Embedding(n_ops, emb_dim)
nn.init.normal_(self.W_g.weight, std=0.02)
nn.init.normal_(self.W_o.weight, std=0.02)
if mode == "random_id":
self.W_id = nn.Embedding(n_perts + 1, emb_dim) # +1 for control/unknown
nn.init.normal_(self.W_id.weight, std=0.02)
if mode == "gene_pathway_op":
assert gene_pathway is not None and n_pathways > 0
self.register_buffer("gene_pathway", torch.as_tensor(gene_pathway, dtype=torch.long))
self.W_path = nn.Embedding(n_pathways, emb_dim)
nn.init.normal_(self.W_path.weight, std=0.02)
def forward(self, gene_ids: torch.Tensor, op_ids: torch.Tensor, mask: torch.Tensor,
pert_ids: torch.Tensor | None = None) -> torch.Tensor:
"""gene_ids/op_ids/mask: (B, Lmax) padded; mask 1 for valid gene slots.
returns e_u: (B, emb_dim). control rows (mask all zero) map to the zero vector.
"""
if self.mode == "random_id":
return self.W_id(pert_ids)
m = mask.unsqueeze(-1).float() # (b, l, 1)
if self.mode == "op_only":
per_gene = self.W_o(op_ids)
elif self.mode == "gene_only":
per_gene = self.W_g(gene_ids)
elif self.mode == "gene_pathway_op":
path_ids = self.gene_pathway[gene_ids.clamp(min=0)]
per_gene = self.W_g(gene_ids) + self.W_o(op_ids) + self.W_path(path_ids)
else: # gene_op
per_gene = self.W_g(gene_ids) + self.W_o(op_ids)
per_gene = per_gene * m
denom = m.sum(dim=1).clamp(min=1.0) # (b,1); control -> 1 to avoid /0 (sum is 0 anyway)
return per_gene.sum(dim=1) / denom
def build_pert_tensors(data, labels, device="cpu"):
"""convert a list of perturbation labels into padded (gene_ids, op_ids, mask, pert_ids).
pert_ids index into ``data.perturbations`` (len = control/unknown sentinel) for random_id mode.
"""
pert_index = {p: i for i, p in enumerate(data.perturbations)}
L = max((len(data.parse(p)) for p in labels), default=1)
L = max(L, 1)
B = len(labels)
g = np.zeros((B, L), dtype=np.int64)
o = np.zeros((B, L), dtype=np.int64)
mask = np.zeros((B, L), dtype=np.float32)
pid = np.full(B, len(data.perturbations), dtype=np.int64) # sentinel
for i, p in enumerate(labels):
gids, oids = data.pert_gene_op_ids(p)
n = len(gids)
if n:
g[i, :n] = gids
o[i, :n] = oids
mask[i, :n] = 1.0
if p in pert_index:
pid[i] = pert_index[p]
t = lambda a: torch.as_tensor(a, device=device)
return t(g), t(o), t(mask), t(pid)
|