models/multimae3d_utils.py

"""
MultiMAE 3D Utility Functions
- Patchify / Unpatchify
- Patch shuffling for masking
- 3D sinusoidal positional embeddings
- Dirichlet masking with missing modality support
"""

from typing import Union, Tuple, Dict, List

import torch
import torch.nn as nn
from torch.distributions import Dirichlet
from einops import rearrange


def to_3tuple(x):
    if isinstance(x, (list, tuple)):
        assert len(x) == 3
        return tuple(x)
    return (x, x, x)


def calc_patchified_dim(
    img_size: Union[int, Tuple[int, int, int]],
    patch_size: Union[int, Tuple[int, int, int]],
) -> Tuple[int, int, int]:
    img_size = to_3tuple(img_size)
    patch_size = to_3tuple(patch_size)
    return tuple(img_size[i] // patch_size[i] for i in range(3))


def patchify(
    image: torch.Tensor,
    patch_size: Union[int, Tuple[int, int, int]],
) -> torch.Tensor:
    """
    Convert image to patches.
    image: [B, C, D, H, W]
    returns: [B, num_patches, C, pd, ph, pw]
    """
    patch_size = to_3tuple(patch_size)
    img_size = image.shape[-3:]
    patchified_dim = calc_patchified_dim(img_size, patch_size)
    patches = rearrange(
        image,
        "b c (nd pd) (nh ph) (nw pw) -> b (nd nh nw) c pd ph pw",
        pd=patch_size[0],
        ph=patch_size[1],
        pw=patch_size[2],
        nd=patchified_dim[0],
        nh=patchified_dim[1],
        nw=patchified_dim[2],
    )
    return patches


def unpatchify(
    patches: torch.Tensor,
    img_size: Union[int, Tuple[int, int, int]],
    patch_size: Union[int, Tuple[int, int, int]],
) -> torch.Tensor:
    """
    Convert patches back to image.
    patches: [B, num_patches, C, pd, ph, pw]
    returns: [B, C, D, H, W]
    """
    patch_size = to_3tuple(patch_size)
    img_size = to_3tuple(img_size)
    patchified_dim = calc_patchified_dim(img_size, patch_size)
    image = rearrange(
        patches,
        "b (nd nh nw) c pd ph pw -> b c (nd pd) (nh ph) (nw pw)",
        pd=patch_size[0],
        ph=patch_size[1],
        pw=patch_size[2],
        nd=patchified_dim[0],
        nh=patchified_dim[1],
        nw=patchified_dim[2],
    )
    return image


def shuffle_patches(
    patches: torch.Tensor,
    permutations: torch.Tensor = None,
    mask_probs: torch.Tensor = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Shuffle patches along the patch dimension.

    When mask_probs is None: uniform random shuffle.
    When mask_probs is provided: Gumbel-top-k weighted shuffle.
      Patches with higher mask_probs end up at higher indices (masked),
      patches with lower mask_probs end up at lower indices (visible).

    Args:
        patches: [B, N, ...]
        permutations: optional pre-computed permutation indices [B, N]
        mask_probs: optional [N] per-patch masking probability (sums to 1)

    Returns:
        (shuffled_patches, perm_indices)
    """
    batch_size, num_patches = patches.shape[0], patches.shape[1]
    if permutations is not None:
        perm_idx = permutations
    else:
        rand = torch.rand(batch_size, num_patches, device=patches.device)

        if mask_probs is not None:
            # Gumbel-top-k trick for weighted sampling without replacement.
            # key_i = log(p_i) + Gumbel(0,1)_i
            # Top-k of keys = sample from Multinomial(p, k)
            # After ascending argsort: low keys → visible, high keys → masked.
            mask_probs = mask_probs.to(patches.device)
            gumbel = -torch.log(-torch.log(rand.clamp(1e-20, 1.0 - 1e-20)))
            log_probs = torch.log(mask_probs.clamp(min=1e-20))  # [N]
            keys = gumbel + log_probs.unsqueeze(0)  # [B, N]
            perm_idx = torch.argsort(keys, dim=1)
        else:
            perm_idx = torch.argsort(rand, dim=1)

    shuffled = patches[torch.arange(batch_size, device=patches.device)[:, None], perm_idx]
    return shuffled, perm_idx


def unshuffle_patches(
    patches: torch.Tensor,
    perm_idx: torch.Tensor,
) -> torch.Tensor:
    """
    Inverse of shuffle_patches.
    """
    batch_size = patches.shape[0]
    inv_idx = torch.argsort(perm_idx, dim=1)
    return patches[torch.arange(batch_size, device=patches.device)[:, None], inv_idx]


def build_3d_sincos_position_embedding(
    grid_size: Tuple[int, int, int],
    embed_dim: int,
    temperature: float = 10000.0,
) -> nn.Parameter:
    """
    Build 3D sinusoidal positional embedding.
    returns: [1, num_patches, embed_dim] (frozen parameter)
    """
    grid_size = to_3tuple(grid_size)
    h, w, d = grid_size

    assert embed_dim % 6 == 0, \
        f"embed_dim ({embed_dim}) must be divisible by 6 for 3D sincos pos embed"

    pos_dim = embed_dim // 6
    omega = torch.arange(pos_dim, dtype=torch.float32) / pos_dim
    omega = 1.0 / (temperature ** omega)

    grid_h = torch.arange(h, dtype=torch.float32)
    grid_w = torch.arange(w, dtype=torch.float32)
    grid_d = torch.arange(d, dtype=torch.float32)

    out_h = torch.einsum("m,d->md", grid_h.flatten(), omega)
    out_w = torch.einsum("m,d->md", grid_w.flatten(), omega)
    out_d = torch.einsum("m,d->md", grid_d.flatten(), omega)

    # Expand to full grid: [H*W*D, pos_dim] for each axis
    # Use meshgrid ordering to get correct spatial layout
    grid_h_idx, grid_w_idx, grid_d_idx = torch.meshgrid(
        torch.arange(h), torch.arange(w), torch.arange(d), indexing="ij"
    )
    grid_h_flat = grid_h_idx.flatten()  # [H*W*D]
    grid_w_flat = grid_w_idx.flatten()
    grid_d_flat = grid_d_idx.flatten()

    pos_emb = torch.cat([
        torch.sin(out_h[grid_h_flat]),
        torch.cos(out_h[grid_h_flat]),
        torch.sin(out_w[grid_w_flat]),
        torch.cos(out_w[grid_w_flat]),
        torch.sin(out_d[grid_d_flat]),
        torch.cos(out_d[grid_d_flat]),
    ], dim=1)[None, :, :]  # [1, num_patches, embed_dim]

    pos_emb = nn.Parameter(pos_emb)
    pos_emb.requires_grad = False
    return pos_emb


def generate_dirichlet_mask_ratios(
    num_modalities: int,
    alpha: float,
    overall_mask_ratio: float,
) -> torch.Tensor:
    """
    Sample per-modality mask ratios from a Dirichlet distribution.
    The total visible budget is distributed among modalities.

    Returns: [num_modalities] tensor of per-modality mask ratios
    """
    dirichlet = Dirichlet(torch.tensor([float(alpha)] * num_modalities))
    visible_ratio = 1.0 - overall_mask_ratio
    total_visible = visible_ratio * num_modalities
    visible_per_mod = total_visible * dirichlet.sample()
    # Clamp to [0, 1]
    mask_ratios = (1.0 - visible_per_mod).clamp(0.0, 1.0)
    return mask_ratios


def compute_mask_ratios(
    modality_names: List[str],
    observed: torch.Tensor,
    mask_ratio: float = 0.75,
    use_dirichlet: bool = True,
    dirichlet_alpha: float = 1.0,
) -> Dict[str, float]:
    """
    Compute per-modality mask ratios, respecting observed mask.
    Missing modalities (observed=0) get mask_ratio=1.0.
    Present modalities get Dirichlet or uniform masking.

    Args:
        modality_names: list of modality names, e.g. ['T1', 'T2', 'Flair', 'PET']
        observed: [M] bool/float tensor, 1.0=present, 0.0=missing
            NOTE: This is per-sample, called once per sample in the batch.
            For simplicity, we use the same mask ratio for the whole batch
            (based on which modalities are present in the majority of the batch).
        mask_ratio: overall target mask ratio for present modalities
        use_dirichlet: whether to use Dirichlet distribution
        dirichlet_alpha: Dirichlet concentration parameter

    Returns:
        dict mapping modality_name -> mask_ratio (float)
    """
    ratios = {}
    present_mods = [name for i, name in enumerate(modality_names) if observed[i] > 0.5]
    missing_mods = [name for i, name in enumerate(modality_names) if observed[i] <= 0.5]

    # Missing modalities: fully masked
    for name in missing_mods:
        ratios[name] = 1.0

    # Present modalities: Dirichlet or uniform
    if len(present_mods) > 0:
        if use_dirichlet and len(present_mods) > 1:
            # Dirichlet masking among present modalities
            dir_ratios = generate_dirichlet_mask_ratios(
                num_modalities=len(present_mods),
                alpha=dirichlet_alpha,
                overall_mask_ratio=mask_ratio,
            )
            for i, name in enumerate(present_mods):
                ratios[name] = dir_ratios[i].item()
        else:
            # Uniform masking
            for name in present_mods:
                ratios[name] = mask_ratio

    return ratios


def mask_data(
    batch: Dict[str, torch.Tensor],
    modality_names: List[str],
    observed: torch.Tensor,
    mask_ratio: float = 0.75,
    patch_size: Union[int, Tuple[int, int, int]] = 16,
    use_dirichlet: bool = True,
    dirichlet_alpha: float = 1.0,
    patch_mask_probs: torch.Tensor = None,
) -> Tuple[
    Dict[str, torch.Tensor],
    Dict[str, torch.Tensor],
    Dict[str, torch.Tensor],
    Dict[str, float],
]:
    """
    Core masking function for MultiMAE pretraining.

    For each modality:
    - Patchify the image
    - Shuffle patches (optionally weighted by anatomy importance)
    - Split into selected (visible) and masked based on mask_ratio
    - Missing modalities (observed=0) get 100% masking

    Args:
        batch: dict mapping modality name -> [B, 1, D, H, W] tensor
        modality_names: ordered list of modality names
        observed: [B, M] tensor indicating which modalities are present
        mask_ratio: target mask ratio for present modalities
        patch_size: patch size for patchification
        use_dirichlet: whether to use Dirichlet distribution
        dirichlet_alpha: Dirichlet concentration parameter
        patch_mask_probs: optional [N_patches] per-patch masking probability
            from anatomy-aware masking. When provided, uses Gumbel-top-k
            weighted sampling instead of uniform random shuffling.
            Higher probability = more likely to be masked.

    Returns:
        selected_patches: {modality: [B, num_selected, C, pd, ph, pw]}
        masked_patches: {modality: [B, num_masked, C, pd, ph, pw]}
        perm_indices: {modality: [B, num_patches]}
        mask_ratios: {modality: float}
    """
    patch_size = to_3tuple(patch_size)
    batch_size = observed.shape[0]

    # Union strategy: if ANY sample in the batch has a modality, it gets
    # partial masking. Samples where this modality is missing contribute
    # zero-valued patches (harmless in encoder, excluded from loss).
    # This ensures no information is wasted when modalities are present
    # in a minority of samples.
    batch_observed = (observed.max(dim=0).values > 0.5).float()  # [M]
    mask_ratios = compute_mask_ratios(
        modality_names=modality_names,
        observed=batch_observed,
        mask_ratio=mask_ratio,
        use_dirichlet=use_dirichlet,
        dirichlet_alpha=dirichlet_alpha,
    )

    selected_patches = {}
    masked_patches = {}
    perm_indices = {}

    for mod_name in modality_names:
        # Patchify: [B, 1, D, H, W] -> [B, num_patches, 1, pd, ph, pw]
        patches = patchify(batch[mod_name], patch_size)
        num_patches = patches.shape[1]

        # Shuffle patches (weighted by anatomy importance if provided)
        shuffled, perm_idx = shuffle_patches(patches, mask_probs=patch_mask_probs)
        perm_indices[mod_name] = perm_idx

        # Split into selected and masked
        mod_mask_ratio = mask_ratios[mod_name]
        num_selected = int((1.0 - mod_mask_ratio) * num_patches)
        # Ensure at least 0 selected (for fully masked modalities)
        num_selected = max(0, num_selected)

        selected_patches[mod_name] = shuffled[:, :num_selected]
        masked_patches[mod_name] = shuffled[:, num_selected:]

    return selected_patches, masked_patches, perm_indices, mask_ratios