Blame for model.py

#!/usr/bin/env python3
"""Dual-head dilated-CNN for per-nucleotide Ribo-seq P-site prediction.

Design (see methods.md section 5B): the per-nt FM embedding already carries long-range
sequence context, so the network body only needs an O(L) local operator to mix in the
RNAseq coverage and sharpen to a base-resolution, frame-aware profile. A stack of residual
dilated 1-D convolutions (BPNet-style, Avsec et al. 2021) does this. Two heads:

  - profile head: per-nt logits -> (at loss time) log-softmax over the transcript ->
    multinomial NLL against per-nt counts. Scale-free, captures 3-nt periodicity and
    CDS localization, and is the transferable target across datasets.
  - count head: a scalar log-total from masked-mean body features + log1p(total coverage),
    regressed to log1p(total P-sites). Magnitude; dataset-depth-dependent.

Input per nt is [FM emb (D_emb) ; log1p(coverage) (1)] -> D_emb+1 channels.
Convolutions operate on (B, C, L); the dataset provides (B, L, C) so we transpose.
"""
from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F


class ResidualDilatedBlock(nn.Module):
    """Conv -> norm -> GELU -> Conv -> norm, added back to the input (same length)."""

    def __init__(self, channels: int, kernel_size: int, dilation: int, dropout: float):
        super().__init__()
        pad = dilation * (kernel_size - 1) // 2
        self.conv1 = nn.Conv1d(channels, channels, kernel_size,
                               padding=pad, dilation=dilation)
        self.conv2 = nn.Conv1d(channels, channels, kernel_size,
                               padding=pad, dilation=dilation)
        # GroupNorm is batch-size independent (we use token-budget batches of varying B).
        self.norm1 = nn.GroupNorm(8, channels)
        self.norm2 = nn.GroupNorm(8, channels)
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        h = self.drop(F.gelu(self.norm1(self.conv1(x))))
        h = self.norm2(self.conv2(h))
        return F.gelu(x + h)


class BiMambaBlock(nn.Module):
    """Pre-norm residual block with a BIDIRECTIONAL Mamba mixer (forward + reversed scan).

    seq2ribo's polisher (Kaynar & Kingsford 2026) refines an sTASEP-simulated profile with a
    single CAUSAL Mamba, which suffices there because the simulator already injects downstream
    (exclusion-zone) context. This model has no simulator, so a causal scan alone would lose the
    downstream context that the bidirectional transformer provides; we run Mamba both directions
    and sum, matching the transformer's full-transcript receptive field at O(L) instead of
    O(L^2). mamba_ssm is imported lazily so the default transformer path needs no CUDA Mamba build.
    """

    def __init__(self, d_model: int, d_state: int = 16, d_conv: int = 4,
                 expand: int = 2, dropout: float = 0.1):
        super().__init__()
        from mamba_ssm import Mamba  # lazy: only imported when a mamba mixer is actually built
        self.norm = nn.LayerNorm(d_model)
        self.fwd = Mamba(d_model=d_model, d_state=d_state, d_conv=d_conv, expand=expand)
        self.bwd = Mamba(d_model=d_model, d_state=d_state, d_conv=d_conv, expand=expand)
        self.drop = nn.Dropout(dropout)

    def forward(self, x, mask):
        # x (B, L, d); mask (B, L) bool. Zero padded positions before the scan so pads do not
        # seed the SSM state (they would otherwise start the reversed scan). Both scans read the
        # same normed input; the assignment to h happens only after the RHS is fully evaluated.
        residual = x
        h = self.norm(x) * mask.unsqueeze(-1).to(x.dtype)
        h = self.fwd(h) + self.bwd(h.flip(1)).flip(1)
        return residual + self.drop(h)


class MambaBody(nn.Module):
    """Stack of bidirectional Mamba blocks; a drop-in alternative to the TransformerEncoder."""

    def __init__(self, d_model: int, n_layers: int, d_state: int = 16, d_conv: int = 4,
                 expand: int = 2, dropout: float = 0.1):
        super().__init__()
        self.blocks = nn.ModuleList([
            BiMambaBlock(d_model, d_state, d_conv, expand, dropout) for _ in range(n_layers)
        ])

    def forward(self, x, mask):
        for blk in self.blocks:
            x = blk(x, mask)
        return x


class RiboSignalModel(nn.Module):
    def __init__(self, d_emb: int = 1280, channels: int = 256, n_blocks: int = 10,
                 kernel_size: int = 3, dropout: float = 0.1,
                 dilations: tuple[int, ...] | None = None,
                 extra_in: int = 0, n_attn_layers: int = 0, n_heads: int = 8,
                 mixer: str = "transformer", mamba_d_state: int = 16,
                 mamba_d_conv: int = 4, mamba_expand: int = 2,
                 learn_start_context: bool = False, fm_to_mixer: bool = False):
        super().__init__()
        self.d_emb = d_emb
        self.extra_in = extra_in            # e.g. sequence-derived ORF-candidate channels
        self.n_attn_layers = n_attn_layers
        self.mixer = mixer                  # global-context mixer: "transformer" or "mamba"
        # Q3 (Kozak ablation): replace the hand-picked start-context heuristic with a LEARNED gate.
        # A conv over the one-hot window [-6,+4] around each nt yields a per-nt gate in (0,1); the
        # (no-Kozak) codon-weight start channel is remodulated w*(0.5+0.5*gate) before in_proj, so the
        # 4x10 kernel is a learned nucleotide x position Kozak matrix (extracted post-training). The
        # start channel is ORF-track channel 3 -> feats index d_emb+3. One-hot backend only.
        self.learn_start_context = learn_start_context
        self.start_ch_idx = d_emb + 3
        if learn_start_context:
            assert d_emb == 4 and extra_in >= 4, \
                "learn_start_context requires the onehot backend (d_emb=4) + the ORF track"
            self.start_ctx_conv = nn.Conv1d(4, 1, kernel_size=10, bias=True)
            self.start_ctx_pad = (6, 3)     # transcript window [i-6, i+3] = Kozak -6..+4
        if dilations is None:
            dilations = tuple(2 ** (i % 10) for i in range(n_blocks))  # 1,2,...,512,1,...
        # input per nt is [FM emb (d_emb) ; ORF track (extra_in) ; log1p(coverage) (1)];
        # coverage stays the LAST channel so the count head can read feats[..., -1].
        self.in_proj = nn.Conv1d(d_emb + extra_in + 1, channels, 1)
        self.blocks = nn.ModuleList([
            ResidualDilatedBlock(channels, kernel_size, d, dropout) for d in dilations
        ])
        # Optional global-context stack: the dilated body has a ~4 kb receptive field, so a
        # position deep in a long CDS cannot see the start codon that sets its frame. Self-
        # attention gives full-transcript context (torch 2.x uses a memory-efficient kernel).
        if n_attn_layers > 0:
            if mixer == "mamba":
                # bidirectional Mamba stack (O(L)); same layer count as the transformer for a
                # controlled swap. Built only here, so mamba_ssm is required just for this path.
                self.mamba_body = MambaBody(channels, n_attn_layers, mamba_d_state,
                                            mamba_d_conv, mamba_expand, dropout)
            else:
                layer = nn.TransformerEncoderLayer(
                    d_model=channels, nhead=n_heads, dim_feedforward=channels * 2,
                    dropout=dropout, activation="gelu", batch_first=True, norm_first=True)
                self.attn = nn.TransformerEncoder(layer, num_layers=n_attn_layers)
        # Direct FM -> mixer path. By default the FM embedding reaches the mixer only after being
        # (a) linearly squeezed to `channels` by the SHARED in_proj, jointly with the ORF track and
        # coverage, and (b) rewritten by all n_blocks dilated convs. For RiNALMo that is 1280 -> 256
        # in a single 1x1 conv before any spatial processing -- a 5x bottleneck the 4-d one-hot
        # backend never pays, which confounds "does the FM help" with "does the FM survive fusion".
        #
        # This adds a dedicated projection of the embedding straight onto the mixer input:
        #     x_mixer = cnn_out + gate * GroupNorm(fm_proj(emb))
        # `gate` is a scalar initialised to ZERO (LayerScale / ReZero style), so the network starts
        # numerically identical to the baseline and must actively open the path. The learned |gate|
        # is therefore a direct readout of how much the model wants the un-convolved embedding --
        # a gate that stays ~0 is itself an interpretable negative result.
        self.fm_to_mixer = fm_to_mixer
        if fm_to_mixer:
            assert n_attn_layers > 0, "fm_to_mixer requires a mixer stack (n_attn_layers > 0)"
            self.fm_proj = nn.Conv1d(d_emb, channels, 1)
            self.fm_norm = nn.GroupNorm(8, channels)
            self.fm_gate = nn.Parameter(torch.zeros(1))
        self.profile_head = nn.Conv1d(channels, 1, 1)
        # count head: masked-mean body features (channels) + log1p(total coverage) (1)
        self.count_head = nn.Sequential(
            nn.Linear(channels + 1, channels), nn.GELU(), nn.Dropout(dropout),
            nn.Linear(channels, 1),
        )

    def forward(self, feats, mask):
        """feats (B, L, d_emb+extra_in+1) float; mask (B, L) bool (True = real nt).

        Returns profile_logits (B, L) with -1e9 at padded positions, and pred_logcount
        (B,) the predicted log1p(total P-sites).
        """
        if self.learn_start_context:
            # learned Kozak: gate the codon-weight start channel by a conv over the one-hot context
            oh = feats[:, :, :4].transpose(1, 2)                  # (B, 4, L) one-hot (onehot backend)
            gate = torch.sigmoid(
                self.start_ctx_conv(F.pad(oh, self.start_ctx_pad))).squeeze(1)   # (B, L) in (0,1)
            idx = self.start_ch_idx
            w = feats[:, :, idx]                                  # (B, L) codon-weight start (no Kozak)
            new_start = (w * (0.5 + 0.5 * gate)).unsqueeze(-1)    # (B, L, 1)
            feats = torch.cat([feats[:, :, :idx], new_start, feats[:, :, idx + 1:]], dim=2)
        x = feats.transpose(1, 2)                 # (B, C_in, L)
        x = self.in_proj(x)
        m = mask.unsqueeze(1).to(x.dtype)         # (B, 1, L)
        for blk in self.blocks:
            x = blk(x * m)                        # zero padded positions before each conv
        x = x * m
        if self.n_attn_layers > 0:
            if self.fm_to_mixer:
                # un-convolved embedding, own projection, zero-init gate (see __init__)
                fm = feats[:, :, :self.d_emb].transpose(1, 2)      # (B, d_emb, L)
                x = x + self.fm_gate * self.fm_norm(self.fm_proj(fm)) * m
            xt = x.transpose(1, 2)                # (B, L, C)
            if self.mixer == "mamba":
                xt = self.mamba_body(xt, mask)
            else:
                xt = self.attn(xt, src_key_padding_mask=~mask)
            x = xt.transpose(1, 2) * m           # (B, C, L), re-zero padded positions
        logits = self.profile_head(x).squeeze(1)  # (B, L)
        logits = logits.masked_fill(~mask, -1e9)

        # count head: masked mean over valid positions
        denom = mask.sum(dim=1, keepdim=True).clamp(min=1).to(x.dtype)   # (B,1)
        feat_mean = (x * m).sum(dim=2) / denom                          # (B, C)
        logcov = torch.log1p(feats[:, :, -1].clamp(min=0).sum(dim=1, keepdim=True))
        pred_logcount = self.count_head(torch.cat([feat_mean, logcov], dim=1)).squeeze(1)
        return logits, pred_logcount


def profile_multinomial_nll(logits, counts, mask):
    """Per-transcript multinomial NLL, equal-weighted across transcripts.

    logits, counts, mask: (B, L). Returns scalar mean over the batch of
    -sum_i (c_i / N) * log softmax(logits)_i  (cross-entropy of the empirical profile),
    which is scale-free (independent of a transcript's total depth N).
    """
    log_p = F.log_softmax(logits, dim=1)               # -inf at padded (logit -1e9)
    log_p = torch.where(mask, log_p, torch.zeros_like(log_p))
    counts = counts * mask                             # zero padded counts
    n = counts.sum(dim=1).clamp(min=1.0)               # (B,)
    per_tx = -(counts * log_p).sum(dim=1) / n          # (B,)
    return per_tx.mean()


def count_mse(pred_logcount, counts, mask):
    """MSE between predicted log-count and log1p(observed total P-sites)."""
    target = torch.log1p((counts * mask).sum(dim=1))
    return F.mse_loss(pred_logcount, target)


def profile_entropy_gap(logits, counts, mask):
    """Anti-smoothing penalty (O3): how much SMOOTHER the predicted profile is than the
    observed one, per transcript, averaged over the batch.

    The multinomial NLL is minimized by matching the empirical profile, but a smooth
    prediction still scores well because cross-entropy tolerates a hedged distribution.
    That smoothness is why RiboCode over-calls short non-canonical ORFs on the predicted
    density (the shape features that separate real from spurious ORFs are washed out).

    This term penalizes the Shannon-entropy excess H(p) - H(q), where p = softmax(logits)
    is the predicted profile and q = counts / N the empirical one, both over the masked
    transcript positions. relu() keeps only over-smoothing (H(p) > H(q)); a prediction
    peakier than the target is never pushed back (avoids chasing shot noise). Because the
    pooled Chothani targets are deep (~6.6 P-sites/nt), H(q) reflects real 3-nt periodicity
    and CDS localization, not sparse-sampling spikes. Returns a scalar >= 0 (nats).
    """
    log_p = F.log_softmax(logits, dim=1)                  # padded logits (-1e9) -> ~ -inf
    zero = torch.zeros_like(log_p)
    p = torch.where(mask, log_p.exp(), zero)              # zero the padded positions
    ent_p = -(p * torch.where(mask, log_p, zero)).sum(dim=1)          # (B,) H(p)

    counts = counts * mask
    n = counts.sum(dim=1).clamp(min=1.0)                  # (B,)
    q = counts / n.unsqueeze(1)
    log_q = torch.where(q > 0, torch.log(q.clamp(min=1e-12)), zero)
    ent_q = -(q * log_q).sum(dim=1)                       # (B,) H(q)

    return F.relu(ent_p - ent_q).mean()


if __name__ == "__main__":
    # smoke: shapes + finite loss on random data
    torch.manual_seed(0)
    B, L, D = 3, 512, 1280
    feats = torch.randn(B, L, D + 1)
    feats[:, :, -1] = feats[:, :, -1].abs()            # coverage >= 0
    mask = torch.ones(B, L, dtype=torch.bool)
    mask[0, 300:] = False                              # a shorter transcript
    counts = (torch.rand(B, L) * 5).floor() * mask
    m = RiboSignalModel(d_emb=D, channels=64, n_blocks=6)
    logits, pc = m(feats, mask)
    loss = profile_multinomial_nll(logits, counts, mask) + 0.1 * count_mse(pc, counts, mask)
    print("logits", tuple(logits.shape), "pred_logcount", tuple(pc.shape),
          "loss", float(loss))
    loss.backward()
    print("backward OK; params:", sum(p.numel() for p in m.parameters()))