| """SimpleP2G-cross: target tokens and spatially positioned neighbor tokens.""" |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
| from torch.utils.checkpoint import checkpoint |
|
|
| from src.simpleP2G import SimpleP2GDecoderLayer, SimpleP2GRMSNorm |
|
|
|
|
| class SimpleP2GNeigh(nn.Module): |
| """The target-plus-neighbors model used in the main comparison.""" |
|
|
| def __init__(self, config): |
| super().__init__() |
| self.config = config |
| self.decoder_depth = config.decoder_depth |
| self.num_genes = config.num_genes |
| self.hidden_size = config.hidden_size |
| self.cond_size = config.cond_size |
| self.num_cond_tokens = getattr(config, "num_cond_tokens", 4) |
| self.num_neighbor_tokens = int(getattr(config, "num_neighbor_tokens", 24)) |
| if str(getattr(config, "neighbor_position_mode", "relative")) != "relative": |
| raise ValueError("SimpleP2G-cross requires relative neighbor positions") |
| for key in ("neighbor_cond_size", "global_cond_size"): |
| if int(getattr(config, key, self.cond_size)) != self.cond_size: |
| raise ValueError("All image embeddings must use the same encoder dimension") |
| if ( |
| not getattr(config, "use_target_context", True) |
| or not getattr(config, "use_neighbor_context", True) |
| or getattr(config, "use_slide_context", False) |
| or getattr(config, "use_local_summary", False) |
| or getattr(config, "context_source_balance", 0.0) != 0 |
| ): |
| raise ValueError( |
| "This model requires target + neighbor context; use the separate experiment entry point for other variants" |
| ) |
| self.gene_queries = nn.Embedding(self.num_genes, self.hidden_size) |
| self.label_embed = nn.Sequential( |
| nn.Linear(self.cond_size, self.cond_size), |
| nn.SiLU(), |
| nn.Linear(self.cond_size, self.hidden_size * self.num_cond_tokens), |
| ) |
| self.neighbor_proj = nn.Sequential( |
| nn.Linear(self.cond_size, self.hidden_size), |
| nn.SiLU(), |
| nn.Linear(self.hidden_size, self.hidden_size), |
| ) |
| self.neighbor_pos_proj = nn.Sequential( |
| nn.Linear(2, self.hidden_size), |
| nn.SiLU(), |
| nn.Linear(self.hidden_size, self.hidden_size), |
| ) |
| self.decoder_blocks = nn.ModuleList([SimpleP2GDecoderLayer(config) for _ in range(self.decoder_depth)]) |
| self.final_norm = SimpleP2GRMSNorm(self.hidden_size) |
| self.mean_head = nn.Linear(self.hidden_size, 1) |
| if config.loss_name in ("gaussian_nll", "beta_nll"): |
| self.log_var_head = nn.Linear(self.hidden_size, 1) |
| self.gradient_checkpointing = False |
| self.emb_dropout = nn.Dropout(getattr(config, "dropout", 0.0)) |
|
|
| def forward(self, c, neighbor_emb=None, neighbor_mask=None, neighbor_pos=None, global_emb=None): |
| """Predict from target (B,D), neighbors (B,K,D) and relative positions. |
| |
| global_emb is accepted for the shared slide-dataset interface; the |
| main-comparison model does not consume a slide token. |
| """ |
| if neighbor_emb is None or neighbor_pos is None: |
| raise ValueError("neighbor_emb and neighbor_pos are required") |
| batch_size = c.shape[0] |
| output = self.emb_dropout(self.gene_queries.weight.unsqueeze(0).expand(batch_size, -1, -1)) |
| target = self.label_embed(c).view(batch_size, self.num_cond_tokens, self.hidden_size) |
| neighbors = self.neighbor_proj(neighbor_emb) |
| neighbors = neighbors + self.neighbor_pos_proj(neighbor_pos.to(dtype=neighbors.dtype)) |
| condition = torch.cat((target, neighbors), dim=1) |
| if neighbor_mask is None: |
| neighbor_mask = torch.ones(neighbor_emb.shape[:2], device=c.device, dtype=torch.bool) |
| target_mask = torch.ones(batch_size, self.num_cond_tokens, device=c.device, dtype=torch.bool) |
| mask = torch.cat((target_mask, neighbor_mask.bool()), dim=1).unsqueeze(1).unsqueeze(1) |
| for block in self.decoder_blocks: |
| if self.gradient_checkpointing and self.training: |
| output = checkpoint(block, output, condition, None, mask, use_reentrant=False) |
| else: |
| output = block(output, encoder_hidden_states=condition, cross_attention_mask=mask) |
| output = self.final_norm(output) |
| mean = F.softplus(self.mean_head(output).squeeze(-1)) |
| log_var = ( |
| self.log_var_head(output).squeeze(-1) |
| if self.config.loss_name in ("gaussian_nll", "beta_nll") |
| else torch.zeros_like(mean) |
| ) |
| return mean, log_var |
|
|
