src/wjad/modules/pos_encoding.py

"""3D RoPE（仅作用于视觉 token）。

12 头按 4+4+4 拆为三组：
  - 头 0-3：射线 RoPE，编码自车系下的单位射线方向 ``(dx, dy, dz)``。
  - 头 4-7：H/W/T RoPE，编码归一化的空间-时间索引 ``(h_norm, w_norm, t_norm)``。
  - 头 8-11：零频段 RoPE，cos=1 / sin=0 → 旋转矩阵恒为 I（identity）。

为减少分支与显存通信，全部 12 头统一走同一份 RoPE 算子（不写 if/else），
零频段头自然变为恒等映射。

将 ``head_dim=64`` 切成 32 个 (cos, sin) 对（两两一组旋转）。每组头内部再按
3 个分量（dx,dy,dz 或 h,w,t）平均分配 32/3 ≈ 10 对（最后 2 对补 0 频）。
"""

from __future__ import annotations

import torch
import torch.nn as nn


def _split_head_dim_for_components(half: int, num_components: int) -> list[int]:
    """把 head_dim/2 个旋转对均匀分给若干个分量；剩余补 0 频。

    返回每个分量分到的旋转对数，最后一项是 ``half - sum(其它)``。
    若 ``num_components == 0``（零频段头），则返回 ``[0, 0, ..., half]``，最后
    一项视为"零频段"——它的频率会被显式置为 0。
    """
    if num_components == 0:
        return [0, half]
    base = half // num_components
    splits = [base] * num_components
    splits[-1] += half - base * num_components  # 余数全归到最后一个分量
    return splits


def build_rope_freqs(
    rays: torch.Tensor,
    hwt_grid: torch.Tensor,
    num_heads: int = 12,
    head_dim: int = 64,
    rope_theta: float = 10000.0,
    device: torch.device | None = None,
    dtype: torch.dtype = torch.float32,
) -> tuple[torch.Tensor, torch.Tensor]:
    """构造 3D RoPE 的 cos / sin 表。

    参数
    ----
    rays : Tensor, shape ``[B, N_v, 3]``
        每个视觉 token 在自车系下的单位射线方向 ``(dx, dy, dz)``。
    hwt_grid : Tensor, shape ``[B, N_v, 3]``
        归一化的空间-时间坐标 ``(h_norm, w_norm, t_norm)`` ∈ [-1, 1]。
    num_heads : int
        总头数（默认 12）。
    head_dim : int
        每头维度（默认 64，必须为偶数）。

    返回
    ----
    cos, sin : Tensor, shape ``[B, N_v, num_heads, head_dim // 2]``
        每个旋转对的 cos / sin 值，已就绪可送入 ``apply_rope``。
    """
    assert head_dim % 2 == 0, "head_dim 必须为偶数"
    assert num_heads % 3 == 0, "num_heads 需被 3 整除以便 4+4+4 分组"

    half = head_dim // 2
    heads_per_group = num_heads // 3
    bsz, n_v, _ = rays.shape
    if device is None:
        device = rays.device

    # === 三组分量值 ===
    # group 0: rays (3 components)
    # group 1: hwt  (3 components)
    # group 2: zero (0 components -> 全部 half 视为零频段)
    splits_g0 = _split_head_dim_for_components(half, 3)  # 用于 rays
    splits_g1 = _split_head_dim_for_components(half, 3)  # 用于 hwt
    splits_g2 = _split_head_dim_for_components(half, 0)  # [0, half]

    # === 频率向量（沿 head_dim 半轴）===
    # 经典 RoPE: theta_i = base ^ (-2i / d)
    #   这里我们对每个分量独立排布频率
    def _freqs(num_pairs: int) -> torch.Tensor:
        # 前 num_pairs 个用 RoPE 频率，剩余补 0
        idx = torch.arange(num_pairs, device=device, dtype=dtype)
        freqs = rope_theta ** (-2.0 * idx / head_dim)
        return freqs  # [num_pairs]

    # 把分量值与频率张量逐头展开为 [B, N_v, num_heads, half]
    angles = torch.zeros(bsz, n_v, num_heads, half, device=device, dtype=dtype)

    # ---- 第 0 组（4 头）：射线 ----
    base_offset = 0
    h0_start = 0
    h0_end = h0_start + heads_per_group
    cursor = 0
    for c in range(3):  # dx, dy, dz
        n_pairs = splits_g0[c]
        if n_pairs > 0:
            f = _freqs(n_pairs)  # [n_pairs]
            comp_val = rays[..., c : c + 1]  # [B, N_v, 1]
            ang = comp_val * f  # 广播 -> [B, N_v, n_pairs]
            angles[:, :, h0_start:h0_end, cursor : cursor + n_pairs] = ang.unsqueeze(2)
        cursor += n_pairs
    # 余数（splits_g0 最后一项的"补足"部分由 _split 已并入最后分量），无需置 0

    # ---- 第 1 组（4 头）：HWT ----
    h1_start = heads_per_group
    h1_end = h1_start + heads_per_group
    cursor = 0
    for c in range(3):  # h, w, t
        n_pairs = splits_g1[c]
        if n_pairs > 0:
            f = _freqs(n_pairs)
            comp_val = hwt_grid[..., c : c + 1]
            ang = comp_val * f
            angles[:, :, h1_start:h1_end, cursor : cursor + n_pairs] = ang.unsqueeze(2)
        cursor += n_pairs

    # ---- 第 2 组（4 头）：零频段 ----
    # 角度恒为 0 → cos=1, sin=0 → 等价 identity；不需要再赋值（已是零）

    cos = torch.cos(angles)
    sin = torch.sin(angles)
    return cos, sin


def apply_rope(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """对 ``q`` ``k`` 的视觉 token 部分应用 3D RoPE。

    所有 12 头一视同仁地走同一段代码（零频段头 cos=1/sin=0 → identity）。

    参数
    ----
    q, k : Tensor, shape ``[B, H, N_v, head_dim]``
    cos, sin : Tensor, shape ``[B, N_v, H, head_dim // 2]``

    返回
    ----
    旋转后的 q, k，形状不变。
    """
    # 把 cos/sin 转成 [B, H, N_v, half]
    cos_e = cos.permute(0, 2, 1, 3)
    sin_e = sin.permute(0, 2, 1, 3)

    # 把 head_dim 维度按 (even, odd) 拆开成 [..., half]
    q_even = q[..., 0::2]
    q_odd = q[..., 1::2]
    k_even = k[..., 0::2]
    k_odd = k[..., 1::2]

    q_rot_even = q_even * cos_e - q_odd * sin_e
    q_rot_odd = q_even * sin_e + q_odd * cos_e
    k_rot_even = k_even * cos_e - k_odd * sin_e
    k_rot_odd = k_even * sin_e + k_odd * cos_e

    q_out = torch.empty_like(q)
    k_out = torch.empty_like(k)
    q_out[..., 0::2] = q_rot_even
    q_out[..., 1::2] = q_rot_odd
    k_out[..., 0::2] = k_rot_even
    k_out[..., 1::2] = k_rot_odd
    return q_out, k_out


class RoPE3D(nn.Module):
    """3D RoPE 工具模块：缓存 hwt_grid（视觉 token 网格上不变），动态计算 rays。

    使用方式：
        rope = RoPE3D(num_heads=12, head_dim=64, T=4, H=12, W=32)
        cos, sin = rope.compute_freqs(rays)            # rays: [B, N_v, 3]
        q, k = apply_rope(q_visual_only, k_visual_only, cos, sin)
    """

    def __init__(
        self,
        num_heads: int = 12,
        head_dim: int = 64,
        time_size: int = 4,
        height_size: int = 12,
        width_size: int = 32,
        rope_theta: float = 10000.0,
    ) -> None:
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.rope_theta = rope_theta
        self.T = time_size
        self.H = height_size
        self.W = width_size

        # 预计算并缓存归一化 H/W/T 网格 [N_v, 3]，N_v = T*H*W
        t = torch.linspace(-1.0, 1.0, steps=time_size) if time_size > 1 else torch.zeros(1)
        h = torch.linspace(-1.0, 1.0, steps=height_size) if height_size > 1 else torch.zeros(1)
        w = torch.linspace(-1.0, 1.0, steps=width_size) if width_size > 1 else torch.zeros(1)
        # 顺序：t -> h -> w（与 Conv3D 输出展平顺序一致）
        T_, H_, W_ = torch.meshgrid(t, h, w, indexing="ij")
        hwt = torch.stack([H_, W_, T_], dim=-1).reshape(-1, 3)  # [N_v, 3]
        self.register_buffer("hwt_grid", hwt, persistent=False)

    @property
    def num_visual_tokens(self) -> int:
        return self.T * self.H * self.W

    def compute_freqs(self, rays: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """根据每 token 的射线方向计算 cos/sin。

        ``rays`` shape: ``[B, N_v, 3]``。
        """
        bsz = rays.shape[0]
        hwt = self.hwt_grid.unsqueeze(0).expand(bsz, -1, -1)  # [B, N_v, 3]
        return build_rope_freqs(
            rays=rays,
            hwt_grid=hwt,
            num_heads=self.num_heads,
            head_dim=self.head_dim,
            rope_theta=self.rope_theta,
            dtype=rays.dtype,
        )