PDGrapher/src/gidflow/models/gap_encoder.py

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