genesis_rna/heads.py

"""
Prediction Heads for Genesis RNA

Multi-task prediction heads for:
1. Masked Language Modeling (MLM)
2. Secondary structure prediction
3. Base-pair prediction
"""

import torch
import torch.nn as nn
from typing import Optional


class MLMHead(nn.Module):
    """
    Masked Language Modeling (MLM) head.

    Predicts masked nucleotide tokens from hidden states.
    Similar to BERT's MLM head architecture.

    Architecture:
        Hidden → Dense → GELU → LayerNorm → Output projection → Vocab logits
    """

    def __init__(self, d_model: int, vocab_size: int):
        """
        Args:
            d_model: Hidden dimension size
            vocab_size: Size of token vocabulary
        """
        super().__init__()
        self.dense = nn.Linear(d_model, d_model)
        self.activation = nn.GELU()
        self.layer_norm = nn.LayerNorm(d_model)
        self.decoder = nn.Linear(d_model, vocab_size)

        # Initialize weights
        self._init_weights()

    def _init_weights(self):
        """Initialize layer weights"""
        nn.init.normal_(self.dense.weight, std=0.02)
        nn.init.zeros_(self.dense.bias)
        nn.init.normal_(self.decoder.weight, std=0.02)
        nn.init.zeros_(self.decoder.bias)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        Args:
            hidden_states: Encoder outputs [batch_size, seq_len, d_model]

        Returns:
            Logits over vocabulary [batch_size, seq_len, vocab_size]
        """
        x = self.dense(hidden_states)
        x = self.activation(x)
        x = self.layer_norm(x)
        logits = self.decoder(x)
        return logits


class StructureHead(nn.Module):
    """
    Secondary structure prediction head.

    Predicts RNA secondary structure type for each position:
    - NONE (unpaired)
    - STEM (base-paired helix)
    - LOOP (hairpin loop)
    - BULGE (bulge/internal loop)
    - HAIRPIN (hairpin loop)

    This is a token-level classification task.
    """

    def __init__(self, d_model: int, num_labels: int):
        """
        Args:
            d_model: Hidden dimension size
            num_labels: Number of structure types (default: 5)
        """
        super().__init__()
        self.dropout = nn.Dropout(0.1)
        self.classifier = nn.Linear(d_model, num_labels)

        # Initialize weights
        self._init_weights()

    def _init_weights(self):
        """Initialize layer weights"""
        nn.init.normal_(self.classifier.weight, std=0.02)
        nn.init.zeros_(self.classifier.bias)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        Args:
            hidden_states: Encoder outputs [batch_size, seq_len, d_model]

        Returns:
            Structure logits [batch_size, seq_len, num_labels]
        """
        x = self.dropout(hidden_states)
        logits = self.classifier(x)
        return logits


class PairHead(nn.Module):
    """
    Base-pair prediction head.

    Predicts which positions form base pairs (Watson-Crick or wobble pairs).
    Uses a bilinear scoring function to compute pair probabilities.

    Output is a symmetric matrix where entry (i, j) represents the
    probability that positions i and j form a base pair.
    """

    def __init__(self, d_model: int, hidden_dim: Optional[int] = None):
        """
        Args:
            d_model: Hidden dimension size
            hidden_dim: Intermediate projection dimension (default: d_model)
        """
        super().__init__()
        if hidden_dim is None:
            hidden_dim = d_model

        # Project hidden states for pairing scores
        self.proj_left = nn.Linear(d_model, hidden_dim, bias=False)
        self.proj_right = nn.Linear(d_model, hidden_dim, bias=False)

        # Optional: add a learned scaling factor
        self.scale = nn.Parameter(torch.ones(1))

        # Initialize weights
        self._init_weights()

    def _init_weights(self):
        """Initialize layer weights"""
        nn.init.normal_(self.proj_left.weight, std=0.02)
        nn.init.normal_(self.proj_right.weight, std=0.02)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        Args:
            hidden_states: Encoder outputs [batch_size, seq_len, d_model]

        Returns:
            Pair logits [batch_size, seq_len, seq_len]
            Entry (b, i, j) is the score for position i pairing with position j
        """
        # Project hidden states
        h_left = self.proj_left(hidden_states)   # [B, L, H]
        h_right = self.proj_right(hidden_states)  # [B, L, H]

        # Compute bilinear pairing scores: h_left @ h_right^T
        # This gives a matrix where entry (i,j) = h_left[i] · h_right[j]
        pair_scores = torch.matmul(h_left, h_right.transpose(-2, -1))  # [B, L, L]

        # Apply scaling
        pair_scores = pair_scores * self.scale

        # Optionally symmetrize (since base pairing is symmetric)
        # pair_scores = (pair_scores + pair_scores.transpose(-2, -1)) / 2

        return pair_scores


class ContactMapHead(nn.Module):
    """
    Alternative contact map prediction head using outer product.

    This head is more parameter-efficient than PairHead and explicitly
    constructs pairwise features via outer product.
    """

    def __init__(self, d_model: int):
        """
        Args:
            d_model: Hidden dimension size
        """
        super().__init__()

        # Reduce dimension for efficiency
        self.proj = nn.Linear(d_model, d_model // 4)
        self.activation = nn.ReLU()

        # Final classifier for pair
        self.classifier = nn.Linear(1, 1)

        # Initialize weights
        self._init_weights()

    def _init_weights(self):
        """Initialize layer weights"""
        nn.init.normal_(self.proj.weight, std=0.02)
        nn.init.zeros_(self.proj.bias)
        nn.init.normal_(self.classifier.weight, std=0.02)
        nn.init.zeros_(self.classifier.bias)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        Args:
            hidden_states: Encoder outputs [batch_size, seq_len, d_model]

        Returns:
            Contact map logits [batch_size, seq_len, seq_len]
        """
        batch_size, seq_len, d_model = hidden_states.size()

        # Project to lower dimension
        h = self.proj(hidden_states)  # [B, L, d//4]
        h = self.activation(h)

        # Compute outer product for all pairs
        # h[i] * h[j] gives pairwise features
        h_i = h.unsqueeze(2)  # [B, L, 1, d//4]
        h_j = h.unsqueeze(1)  # [B, 1, L, d//4]

        # Element-wise product
        pairwise = h_i * h_j  # [B, L, L, d//4]

        # Reduce to scalar score per pair
        contact_scores = pairwise.sum(dim=-1, keepdim=True)  # [B, L, L, 1]
        contact_scores = self.classifier(contact_scores)      # [B, L, L, 1]
        contact_scores = contact_scores.squeeze(-1)           # [B, L, L]

        return contact_scores


class MutationEffectHead(nn.Module):
    """
    Mutation effect prediction head (for fine-tuning).

    Predicts the effect of mutations on RNA function, stability, or binding.
    Can be used for:
    - Fitness prediction
    - Stability change (ΔΔG)
    - Binding affinity change

    This is typically a sequence-level regression task.
    """

    def __init__(self, d_model: int, pooling: str = "mean"):
        """
        Args:
            d_model: Hidden dimension size
            pooling: How to pool sequence representations ("mean", "max", "cls")
        """
        super().__init__()
        self.pooling = pooling

        self.regressor = nn.Sequential(
            nn.Linear(d_model, d_model // 2),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(d_model // 2, 1),
        )

        # Initialize weights
        self._init_weights()

    def _init_weights(self):
        """Initialize layer weights"""
        for module in self.regressor.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, std=0.02)
                nn.init.zeros_(module.bias)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """
        Args:
            hidden_states: Encoder outputs [batch_size, seq_len, d_model]
            attention_mask: Binary mask [batch_size, seq_len]

        Returns:
            Effect predictions [batch_size, 1]
        """
        # Pool sequence representation
        if self.pooling == "mean":
            if attention_mask is not None:
                # Masked mean pooling
                mask_expanded = attention_mask.unsqueeze(-1).float()
                summed = (hidden_states * mask_expanded).sum(dim=1)
                pooled = summed / mask_expanded.sum(dim=1).clamp(min=1e-9)
            else:
                pooled = hidden_states.mean(dim=1)

        elif self.pooling == "max":
            pooled = hidden_states.max(dim=1)[0]

        elif self.pooling == "cls":
            # Use [CLS] token (first position)
            pooled = hidden_states[:, 0, :]

        else:
            raise ValueError(f"Unknown pooling method: {self.pooling}")

        # Predict effect
        effect = self.regressor(pooled)
        return effect