src/models/msm.py

"""
Masked Slice Modeling (MSM) — Self-Supervised Pretraining for ACL-LKNet.

This is the PRIMARY RESEARCH CONTRIBUTION.

Core idea: MRI volumes have ordered slices with anatomical continuity.
We exploit this structure by masking some slices and training the model
to reconstruct their features from the remaining (unmasked) context.

Research question:
    Can a self-supervised objective that explicitly models inter-slice
    anatomical context produce better transferable representations for
    ACL injury detection than established pretraining?

Masking strategies (research axis):
    - random:      Mask 50% of slices uniformly at random
    - contiguous:  Mask contiguous blocks of 3-5 adjacent slices
    - structured:  Preferentially mask central slices (clinically relevant)
    - mixed:       Alternate random and contiguous per batch
"""

import math
import random as py_random
from typing import Tuple, Optional

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


class MaskedSliceModeling(nn.Module):
    """
    Self-supervised pretraining via Masked Slice Modeling.
    
    Architecture:
        1. Encode all slices through the shared backbone → slice features
        2. Replace masked slice features with learnable [MASK] tokens
        3. Add positional encoding (sinusoidal — respects slice ordering)
        4. Pass through lightweight Transformer decoder
        5. Predict the original features of masked slices
    
    Loss: MSE between predicted and actual features of masked slices
    """

    def __init__(
        self,
        feature_dim: int,
        decoder_dim: int = 256,
        decoder_layers: int = 2,
        decoder_heads: int = 4,
        max_slices: int = 48,
        mask_ratio: float = 0.5,
        mask_strategy: str = "random",
    ):
        super().__init__()
        self.feature_dim = feature_dim
        self.decoder_dim = decoder_dim
        self.mask_ratio = mask_ratio
        self.mask_strategy = mask_strategy
        self.max_slices = max_slices

        # Learnable [MASK] token
        self.mask_token = nn.Parameter(torch.zeros(1, 1, feature_dim))
        nn.init.trunc_normal_(self.mask_token, std=0.02)

        # Project encoder features → decoder dimension
        self.encoder_to_decoder = nn.Linear(feature_dim, decoder_dim)

        # Sinusoidal positional encoding (respects spatial ordering of slices)
        self.register_buffer(
            "pos_encoding", self._sinusoidal_encoding(max_slices, decoder_dim)
        )

        # Lightweight Transformer decoder
        decoder_layer = nn.TransformerEncoderLayer(
            d_model=decoder_dim,
            nhead=decoder_heads,
            dim_feedforward=decoder_dim * 4,
            dropout=0.1,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.decoder = nn.TransformerEncoder(
            decoder_layer,
            num_layers=decoder_layers,
        )

        # Predict original features from decoded representations
        self.predictor = nn.Sequential(
            nn.LayerNorm(decoder_dim),
            nn.Linear(decoder_dim, feature_dim),
        )

    @staticmethod
    def _sinusoidal_encoding(max_len: int, dim: int) -> torch.Tensor:
        """Generate sinusoidal positional encoding."""
        pe = torch.zeros(max_len, dim)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        div_term = torch.exp(
            torch.arange(0, dim, 2).float() * (-math.log(10000.0) / dim)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        return pe.unsqueeze(0)  # (1, max_len, dim)

    def generate_mask(
        self, num_slices: int, strategy: Optional[str] = None
    ) -> torch.Tensor:
        """
        Generate a binary mask indicating which slices to mask.
        
        Args:
            num_slices: Number of slices in the volume
            strategy: Override the default masking strategy
            
        Returns:
            mask: (num_slices,) boolean tensor. True = masked (to predict)
        """
        strategy = strategy or self.mask_strategy
        num_mask = max(1, int(num_slices * self.mask_ratio))
        # Always keep at least 2 slices unmasked for context
        num_mask = min(num_mask, num_slices - 2)

        mask = torch.zeros(num_slices, dtype=torch.bool)

        if strategy == "random":
            indices = torch.randperm(num_slices)[:num_mask]
            mask[indices] = True

        elif strategy == "contiguous":
            # Mask contiguous blocks of 3-5 slices
            remaining = num_mask
            while remaining > 0:
                block_size = min(py_random.randint(3, 5), remaining)
                max_start = num_slices - block_size
                if max_start <= 0:
                    start = 0
                else:
                    start = py_random.randint(0, max_start)
                mask[start : start + block_size] = True
                remaining = num_mask - mask.sum().item()

        elif strategy == "structured":
            # Preferentially mask central slices (where ACL is typically visible)
            center = num_slices // 2
            # Create probability distribution peaked at center
            positions = torch.arange(num_slices).float()
            probs = torch.exp(-0.5 * ((positions - center) / (num_slices / 4)) ** 2)
            probs = probs / probs.sum()
            indices = torch.multinomial(probs, num_mask, replacement=False)
            mask[indices] = True

        elif strategy == "mixed":
            # Randomly choose between random and contiguous per call
            sub_strategy = py_random.choice(["random", "contiguous"])
            mask = self.generate_mask(num_slices, strategy=sub_strategy)

        else:
            raise ValueError(f"Unknown mask strategy: {strategy}")

        return mask

    def forward(
        self,
        slice_features: torch.Tensor,
        slice_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Forward pass for MSM pretraining.
        
        Args:
            slice_features: (B, S, D) — encoded slice features from backbone
            slice_mask: (B, S) — optional padding mask (True = valid)
            
        Returns:
            loss: scalar MSE loss on masked slices
            predictions: (B, S, D) — predicted features for all slices
            mask: (B, S) — boolean mask of which slices were masked
        """
        B, S, D = slice_features.shape

        # Generate masks for each sample in the batch
        masks = torch.stack([self.generate_mask(S) for _ in range(B)])  # (B, S)
        masks = masks.to(slice_features.device)

        # Replace masked positions with [MASK] token
        mask_tokens = self.mask_token.expand(B, S, -1)  # (B, S, D)
        masked_features = slice_features.clone()
        masked_features[masks] = mask_tokens[masks]

        # Project to decoder dimension
        x = self.encoder_to_decoder(masked_features)  # (B, S, decoder_dim)

        # Add positional encoding
        x = x + self.pos_encoding[:, :S, :]

        # Transformer decoder
        x = self.decoder(x)  # (B, S, decoder_dim)

        # Predict original features
        predictions = self.predictor(x)  # (B, S, D)

        # Compute loss only on masked positions
        if masks.any():
            pred_masked = predictions[masks]     # (num_masked, D)
            target_masked = slice_features[masks]  # (num_masked, D)
            loss = F.mse_loss(pred_masked, target_masked)
        else:
            loss = torch.tensor(0.0, device=slice_features.device)

        return loss, predictions, masks