PDGrapher/src/gidflow/models/population_encoder.py

"""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]