File size: 5,377 Bytes
07fcdfe | 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 | """Gap encoder: fuses source and endpoint embeddings into a 128-dim CLIP-style gap embedding."""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Dict
class GapEncoder(nn.Module):
"""Encode the 'gap' between source and endpoint population embeddings.
Concatenates four interaction terms:
[z_source, z_target, z_target - z_source, z_source * z_target]
then projects to z_gap via MLP, followed by a CLIP-style projection head
that produces a normalized 128-dim gap_emb for drug matching.
Parameters
----------
input_dim : dimensionality of z_source / z_target (H)
hidden_dim : width of intermediate MLP layers
output_dim : size of z_gap (pre-projection)
n_layers : number of MLP layers
proj_dim : dimensionality of the CLIP-style projection head (default 128)
num_cell_lines : number of distinct cell lines for conditioning embedding;
pass 0 to disable cell-line conditioning
num_genes : number of genes for the auxiliary reconstruction head;
pass 0 to disable the reconstruction head
"""
def __init__(
self,
input_dim: int,
hidden_dim: int = 256,
output_dim: int = 256,
n_layers: int = 2,
proj_dim: int = 128,
num_cell_lines: int = 0,
num_genes: int = 0,
) -> None:
super().__init__()
self.output_dim = output_dim
self.proj_dim = proj_dim
self.num_cell_lines = num_cell_lines
self.num_genes = num_genes
# ---- core MLP: 4*H → hidden → ... → output_dim ----
in_dim = 4 * input_dim
layers: list = []
for _ in range(n_layers):
layers += [nn.Linear(in_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.GELU()]
in_dim = hidden_dim
layers.append(nn.Linear(hidden_dim, output_dim))
self.mlp = nn.Sequential(*layers)
# ---- CLIP-style projection head: output_dim → proj_dim, then L2-norm ----
self.gap_proj_head = nn.Sequential(
nn.Linear(output_dim, proj_dim),
nn.LayerNorm(proj_dim),
)
# ---- optional cell-line conditioning embedding ----
# index 0 is reserved as padding / unknown; real cell-line ids start at 1
if num_cell_lines > 0:
self.cell_line_emb = nn.Embedding(num_cell_lines + 1, proj_dim, padding_idx=0)
else:
self.cell_line_emb = None
# ---- optional auxiliary reconstruction head (Phase-1 MSE loss) ----
if num_genes > 0:
self.recon_head = nn.Linear(proj_dim, num_genes)
else:
self.recon_head = None
# ------------------------------------------------------------------
# helpers
# ------------------------------------------------------------------
def _project(self, z_gap: torch.Tensor) -> torch.Tensor:
"""Project z_gap → normalized 128-dim gap_emb."""
gap_emb = self.gap_proj_head(z_gap) # [B, proj_dim]
gap_emb = F.normalize(gap_emb, p=2, dim=-1) # L2 normalize
return gap_emb
# ------------------------------------------------------------------
# public API
# ------------------------------------------------------------------
def forward(
self,
z_source: torch.Tensor,
z_target: torch.Tensor,
cell_line_ids: Optional[torch.Tensor] = None,
) -> Dict[str, torch.Tensor]:
"""
Parameters
----------
z_source : [B, H]
z_target : [B, H]
cell_line_ids : [B] int64, optional — cell-line indices for conditioning;
if None and cell_line_emb exists, conditioning is skipped.
Returns
-------
dict with keys:
'z_gap' : [B, output_dim] — pre-projection MLP output
'gap_emb' : [B, proj_dim] — L2-normalized CLIP embedding
"""
combined = torch.cat(
[z_source, z_target, z_target - z_source, z_source * z_target],
dim=-1,
)
z_gap = self.mlp(combined) # [B, output_dim]
gap_emb = self._project(z_gap) # [B, proj_dim], L2-normalized
# optional additive cell-line conditioning (re-normalize after addition)
if self.cell_line_emb is not None and cell_line_ids is not None:
cl_emb = self.cell_line_emb(cell_line_ids) # [B, proj_dim]
gap_emb = F.normalize(gap_emb + cl_emb, p=2, dim=-1)
return {
"z_gap": z_gap,
"gap_emb": gap_emb,
}
def reconstruct_expression(self, gap_emb: torch.Tensor) -> torch.Tensor:
"""Decode gap_emb into predicted delta-expression (auxiliary MSE loss).
Parameters
----------
gap_emb : [B, proj_dim] — L2-normalized CLIP embedding
Returns
-------
delta_expr : [B, num_genes]
Raises
------
RuntimeError if the reconstruction head was not built (num_genes == 0).
"""
if self.recon_head is None:
raise RuntimeError(
"reconstruct_expression() requires num_genes > 0 at construction time."
)
return self.recon_head(gap_emb)
|