"""Population-level encoder: [B, N, G] → [B, H].""" from typing import Optional import torch import torch.nn as nn from ..metrics.distribution_metrics import masked_mean, masked_var class PopulationEncoder(nn.Module): """Encode a variable-size cell population into a fixed-size vector. Uses a **four-way pooling** strategy to capture population structure: 1. **Mean pooling**: per-gene mean expression across cells 2. **Variance pooling**: per-gene variance (captures subpopulation heterogeneity) 3. **Attention pooling**: learnable query attends over per-cell embeddings, upweighting cells near the population center 4. **Skewness pooling** (optional): third moment for tail behavior The four streams are projected to a common dimension, concatenated, and processed through an MLP to produce the final embedding. Parameters ---------- num_genes : input feature dimension G hidden_dim : width of MLP layers (default 256) output_dim : size of output embedding H (default 128) n_layers : number of MLP layers (default 2) use_var : if True, include variance pooling use_attention : if True, include attention pooling use_skewness : if True, include skewness pooling (clamped for stability) cell_encoder_hidden : hidden dim for per-cell encoder (default 128) """ def __init__( self, num_genes: int, hidden_dim: int = 256, output_dim: int = 128, n_layers: int = 2, use_var: bool = True, use_attention: bool = True, use_skewness: bool = False, cell_encoder_hidden: int = 128, ) -> None: super().__init__() self.num_genes = num_genes self.output_dim = output_dim self.use_var = use_var self.use_attention = use_attention self.use_skewness = use_skewness # ── Per-cell encoder (for attention pooling) ────────────────── # Maps each cell's gene expression to a compact embedding if use_attention: enc_layers = [] in_dim = num_genes for i in range(2): # 2-layer small encoder out_dim = cell_encoder_hidden if i == 0 else cell_encoder_hidden enc_layers += [nn.Linear(in_dim, out_dim), nn.LayerNorm(out_dim), nn.GELU()] in_dim = out_dim self.cell_encoder = nn.Sequential(*enc_layers) self.cell_emb_dim = cell_encoder_hidden # Learnable query for attention pooling [1, H_cell] self.query = nn.Parameter(torch.randn(1, cell_encoder_hidden) * 0.02) # Projection for mean/var to cell_emb_dim for concatenation self.mean_proj = nn.Linear(num_genes, cell_encoder_hidden) if use_var: self.var_proj = nn.Linear(num_genes, cell_encoder_hidden) if use_skewness: self.skew_proj = nn.Linear(num_genes, cell_encoder_hidden) else: self.cell_encoder = None self.cell_emb_dim = 0 # ── Output MLP ──────────────────────────────────────────────── # Compute total input dimension mlp_in = 0 if use_var: mlp_in += num_genes * 2 # mean + var else: mlp_in += num_genes # mean only if use_attention: mlp_in += self.cell_emb_dim # attention pooling if use_skewness: mlp_in += num_genes # skewness layers = [] in_dim = mlp_in 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) def forward( self, cells: torch.Tensor, mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Encode a cell population into a fixed-size vector. Parameters ---------- cells : [B, N, G] raw gene expression (log1p normalized) mask : [B, N] optional boolean mask (True = valid cell) Returns ------- z : [B, output_dim] population embedding """ B, N, G = cells.shape device = cells.device # ── Stream 1: Mean + Var ───────────────────────────────────── mu = masked_mean(cells, mask) # [B, G] streams = [mu] if self.use_var: var = masked_var(cells, mask) # [B, G] streams.append(var) # ── Stream 2: Attention pooling ────────────────────────────── if self.use_attention and self.cell_encoder is not None: # Encode each cell independently flat_cells = cells.reshape(B * N, G) cell_emb = self.cell_encoder(flat_cells) # [B*N, H_cell] cell_emb = cell_emb.reshape(B, N, -1) # [B, N, H_cell] # Learnable query attends over all cells # query: [1, H_cell] → [B, 1, H_cell] q = self.query.unsqueeze(0).expand(B, -1, -1) # [B, 1, H_cell] scores = torch.bmm(q, cell_emb.transpose(1, 2)) # [B, 1, N] scores = scores.squeeze(1) # [B, N] # Apply mask: set invalid cells to -inf before softmax if mask is not None: scores = scores.masked_fill(~mask, float("-inf")) weights = torch.softmax(scores, dim=-1) # [B, N] attn_pooled = torch.bmm(weights.unsqueeze(1), cell_emb).squeeze(1) # [B, H_cell] streams.append(attn_pooled) # ── Stream 3: Skewness (optional) ──────────────────────────── if self.use_skewness: # Third central moment, clamped for numerical stability if mask is not None: masked_cells = cells * mask.unsqueeze(-1) # [B, N, G] n_valid = mask.sum(dim=-1, keepdim=True).clamp(min=1) # [B, 1] mean_expanded = mu.unsqueeze(1) # [B, 1, G] diff = masked_cells - mean_expanded # [B, N, G] skew_num = (diff ** 3).sum(dim=1) # [B, G] var_val = masked_var(cells, mask).clamp(min=1e-6) # [B, G] skewness = skew_num / (n_valid * var_val ** 1.5) # [B, G] skewness = skewness.clamp(-3.0, 3.0) # Clamp for stability else: mean_expanded = mu.unsqueeze(1) diff = cells - mean_expanded var_val = cells.var(dim=1, keepdim=False).clamp(min=1e-6) skew_num = (diff ** 3).mean(dim=1) skewness = (skew_num / (var_val ** 1.5)).clamp(-3.0, 3.0) streams.append(skewness) # ── Concatenate & project ──────────────────────────────────── pooled = torch.cat(streams, dim=-1) # [B, mlp_in] return self.mlp(pooled) # [B, output_dim]