Blame for src/opentslm/model/encoder/TransformerCNNEncoder2D.py

# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
#
# SPDX-License-Identifier: MIT

import math
from typing import List, Tuple
import torch
import torch.nn as nn

from opentslm.model_config import TRANSFORMER_INPUT_DIM, ENCODER_OUTPUT_DIM, PATCH_SIZE
from opentslm.model.encoder.TimeSeriesEncoderBase import TimeSeriesEncoderBase


def _get_1d_sincos_pos_embed(embed_dim: int, pos: torch.Tensor, base: float = 10000.0) -> torch.Tensor:
    """
    MAE-style 1D sin-cos embedding.
    pos: (M,) float tensor of positions
    returns: (M, embed_dim)
    """
    assert embed_dim % 2 == 0
    device, dtype = pos.device, pos.dtype

    omega = torch.arange(embed_dim // 2, device=device, dtype=dtype)
    omega = omega / (embed_dim / 2.0)
    omega = 1.0 / (base ** omega)  # (D/2,)

    out = pos.unsqueeze(1) * omega.unsqueeze(0)  # (M, D/2)
    emb = torch.cat([torch.sin(out), torch.cos(out)], dim=1)  # (M, D)
    return emb


class TransformerCNNEncoder2D(TimeSeriesEncoderBase):
    """
    Causal patch encoder — each patch attends only to itself and all
    preceding patches (causal self-attention mask).

    Runs ONCE on the full history window and produces a fixed memory M
    that the causal forecast decoder cross-attends to at every generation
    step. Causal attention ensures patch representation i encodes only
    patches 0..i, preventing the LLM from reading future patch information
    out of earlier positions and making the next-patch prediction task
    non-trivial.

    - 24 causal self-attention layers (TransformerEncoderLayer, causal mask)
    - d_model = 768, ff_dim = 3072
    - 12 attention heads

    Input:
      x: (C, L)  — C channels (total across batch), L timesteps
                    (padded to a multiple of patch_size)

    Output:
      (C, T, D)  — T = L / patch_size patch representations
    """

    def __init__(
        self,
        args=None,
        output_dim: int = ENCODER_OUTPUT_DIM,
        dropout: float = 0.1,
        transformer_input_dim: int = 768,
        num_heads: int = 12,
        num_layers: int = 24,
        patch_size: int = PATCH_SIZE,
        ff_dim: int = 3072,
        pe_base: float = 10000.0,        # legacy single-base fallback
        pe_base_t: float | None = None,  # overrides pe_base for TIME axis
        pe_base_c: float | None = None,  # overrides pe_base for CHANNEL axis
        return_4d: bool = False,
    ):
        super().__init__(output_dim, dropout)
        self.patch_size = patch_size
        # Separate bases for time vs channel. CLI args (`args.pe_base_t/c`) take
        # precedence; constructor kwargs are next; legacy `pe_base` is final
        # fallback. With separate bases, channel PE can use a small base (e.g.
        # 10) so c∈{0..C-1} produces visibly distinct sin-cos signatures, while
        # time PE keeps the standard 10000 (good for positions up to ~10k).
        _pe_t_arg = getattr(args, "pe_base_t", None) if args is not None else None
        _pe_c_arg = getattr(args, "pe_base_c", None) if args is not None else None
        self.pe_base_t = float(_pe_t_arg if _pe_t_arg is not None else (pe_base_t if pe_base_t is not None else pe_base))
        self.pe_base_c = float(_pe_c_arg if _pe_c_arg is not None else (pe_base_c if pe_base_c is not None else pe_base))
        self.pe_base = float(pe_base)   # legacy attr retained
        self.return_4d = return_4d
        self.transformer_input_dim = transformer_input_dim

        # ── Patch embedding: (B, 1, L) → (B, D, T) ──────────────────────────
        self.patch_embed = nn.Conv1d(
            in_channels=1,
            out_channels=transformer_input_dim,
            kernel_size=patch_size,
            stride=patch_size,
            bias=False,
        )

        # ── Frequency patch embedding (optional) ─────────────────────────────
        # rfft on a patch of length P gives P//2+1 complex values → 2*(P//2+1) reals
        self._use_freq_embed = getattr(args, "use_freq_embed", False)
        # See MLPEncoder2D for why "linear" is an exact linear reparameterisation
        # of patch_embed and therefore expressively inert; "spectral" is not.
        self._freq_mode = str(getattr(args, "freq_mode", "linear")).lower()
        if self._freq_mode not in ("linear", "spectral"):
            raise ValueError(
                f"freq_mode must be 'linear' or 'spectral', got {self._freq_mode!r}"
            )
        if self._use_freq_embed:
            freq_input_dim = (patch_size // 2 + 1) * 2
            self.freq_embed = nn.Linear(freq_input_dim, transformer_input_dim, bias=False)
            self.concat_proj = nn.Linear(2 * transformer_input_dim, transformer_input_dim, bias=False)

        self.input_norm = nn.LayerNorm(transformer_input_dim)
        self.input_dropout = nn.Dropout(self.dropout)

        # ── Transformer encoder (24 layers) ──────────────────────────────────
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=transformer_input_dim,
            nhead=num_heads,
            dim_feedforward=ff_dim,
            dropout=self.dropout,
            batch_first=True,
            activation="gelu",
            norm_first=True,  # Pre-LN for stable deep training
        )
        self.encoder = nn.TransformerEncoder(
            encoder_layer,
            num_layers=num_layers,
            norm=nn.LayerNorm(transformer_input_dim),  # final norm
        )

        if transformer_input_dim != output_dim:
            self.output_proj = nn.Linear(transformer_input_dim, output_dim)
        else:
            self.output_proj = nn.Identity()

        self.args = args

        self._print_param_count()

    def _print_param_count(self):
        total = sum(p.numel() for p in self.parameters())
        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        print(f"🔢 Encoder: {total:,} params ({trainable:,} trainable)")

    # ------------------------------------------------------------------
    # Incremental (KV-cache) helpers
    # ------------------------------------------------------------------

    def _compute_tok(
        self,
        x: torch.Tensor,          # [C, T_new * patch_size]
        chan_id: torch.Tensor,     # [C]
        t_offset: int = 0,        # time-position index of the first new patch
    ) -> torch.Tensor:
        """Patch-embed x and add 2D sin-cos positional encoding.

        Returns [C, T_new, transformer_input_dim] after input_norm + dropout.
        Mirrors the embedding block in forward() but supports an arbitrary
        starting time position (t_offset) so that new patches receive the
        correct absolute time PE.
        """
        C = x.size(0)
        T_new = x.size(1) // self.patch_size
        D = self.transformer_input_dim

        tok = self.patch_embed(x.reshape(C, 1, -1)).transpose(1, 2)  # [C, T_new, D]

        if self._use_freq_embed:
            x_patches = x.reshape(C, T_new, self.patch_size)
            tok = self._fuse_freq(tok, x_patches)

        t_pos = torch.arange(t_offset, t_offset + T_new, device=tok.device, dtype=torch.float32)
        c_pos = chan_id.to(device=tok.device, dtype=torch.float32)

        if getattr(self.args, "use_only_time_pe", False):
            pe = _get_1d_sincos_pos_embed(D, t_pos, base=self.pe_base_t).to(tok.dtype)
            pe = pe.unsqueeze(0).expand(C, -1, -1)
        elif getattr(self.args, "use_only_channel_pe", False):
            pe = _get_1d_sincos_pos_embed(D, c_pos, base=self.pe_base_c).to(tok.dtype)
            pe = pe.unsqueeze(1).expand(-1, T_new, -1)
        else:
            if D % 4 == 0:
                pe_t = _get_1d_sincos_pos_embed(D // 2, t_pos, base=self.pe_base_t).to(tok.dtype)
                pe_c = _get_1d_sincos_pos_embed(D // 2, c_pos, base=self.pe_base_c).to(tok.dtype)
                pe = torch.cat(
                    [pe_t.unsqueeze(0).expand(C, -1, -1),
                     pe_c.unsqueeze(1).expand(-1, T_new, -1)],
                    dim=-1,
                )
            else:
                pe_t = _get_1d_sincos_pos_embed(D, t_pos, base=self.pe_base_t).to(tok.dtype)
                pe_c = _get_1d_sincos_pos_embed(D, c_pos, base=self.pe_base_c).to(tok.dtype)
                pe = pe_t.unsqueeze(0).expand(C, -1, -1) + pe_c.unsqueeze(1).expand(-1, T_new, -1)

        tok = tok + pe
        tok = self.input_norm(tok)
        tok = self.input_dropout(tok)
        return tok  # [C, T_new, D]

    def tokenize_patches(
        self,
        x: torch.Tensor,
        chan_id: torch.Tensor,
        t_offset: int = 0,
    ) -> torch.Tensor:
        """
        Return patch tokens before the encoder stack.

        This exposes the encoder's patch embedding and positional encoding
        without running the contextual encoder stack.
        """
        if x.size(1) % self.patch_size != 0:
            raise ValueError(f"L={x.size(1)} not divisible by patch_size={self.patch_size}")
        return self._compute_tok(x, chan_id, t_offset=t_offset)

    def build_kv_cache(
        self,
        x: torch.Tensor,      # [C, L]  L divisible by patch_size
        chan_id: torch.Tensor, # [C]
    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        """Full encoder forward that also returns per-layer input caches.

        The cache stores the un-normed input to each transformer layer for
        every history patch.  forward_incremental() uses these as K/V so
        that a new patch can attend to all history without re-encoding it.

        Returns:
            output    : [C, T, output_dim]  (same as forward())
            kv_caches : list of num_layers tensors, each [C, T, transformer_dim]
        """
        if x.size(1) % self.patch_size != 0:
            raise ValueError(f"L={x.size(1)} not divisible by patch_size={self.patch_size}")

        tok = self._compute_tok(x, chan_id, t_offset=0)  # [C, T, D]
        T = tok.size(1)
        causal_mask = nn.Transformer.generate_square_subsequent_mask(
            T, device=tok.device, dtype=tok.dtype
        )

        kv_caches: List[torch.Tensor] = []
        h = tok
        for layer in self.encoder.layers:
            kv_caches.append(h)                             # store pre-layer activation
            h = layer(h, src_mask=causal_mask, is_causal=True)

        if self.encoder.norm is not None:
            h = self.encoder.norm(h)
        output = self.output_proj(h)
        return output, kv_caches

    def forward_incremental(
        self,
        x_new: torch.Tensor,           # [C, patch_size]  one new patch per channel
        chan_id: torch.Tensor,          # [C]
        t_offset: int,                  # absolute time index of the new patch
        kv_caches: List[torch.Tensor],  # per-layer [C, T_past, D]; updated in-place
    ) -> torch.Tensor:
        """Encode one new patch using cached K/V activations from all past patches.

        Because the encoder is causal, position t can only attend to 0..t, so
        past positions' outputs never change.  Only the new position's output
        must be computed, attending to the cached history via the saved
        per-layer pre-norm activations.

        kv_caches is updated in-place: each entry grows by one column so
        successive calls build up the growing sequence correctly.

        Returns: [C, 1, output_dim]
        """
        # 1. Embed the new patch with correct absolute time PE
        h_new = self._compute_tok(x_new, chan_id, t_offset=t_offset)  # [C, 1, D]

        # 2. Propagate through each transformer layer with cached K/V
        for l, layer in enumerate(self.encoder.layers):
            h_past = kv_caches[l]               # [C, T_past, D]  — pre-layer inputs

            # norm_first=True: K/V are the normed layer inputs
            kv_past  = layer.norm1(h_past)       # [C, T_past, D]
            q_kv_new = layer.norm1(h_new)        # [C, 1, D]

            # K and V: all past + new token (new token is the last so no causal
            # masking needed — it can attend to every preceding position)
            k = torch.cat([kv_past, q_kv_new], dim=1)   # [C, T_past+1, D]
            v = k

            attn_out = layer.self_attn(q_kv_new, k, v, need_weights=False)[0]  # [C, 1, D]

            # Store the pre-layer input for the new position before overwriting h_new
            kv_caches[l] = torch.cat([h_past, h_new], dim=1)   # grow cache by 1

            # Apply residual + feed-forward (norm_first path)
            h_new = h_new + layer.dropout1(attn_out)
            h_new = h_new + layer._ff_block(layer.norm2(h_new))

        # 3. Final norm + projection
        if self.encoder.norm is not None:
            h_new = self.encoder.norm(h_new)
        return self.output_proj(h_new)   # [C, 1, output_dim]

    def _fuse_freq(self, tok, x_patches):
        """Fuse a per-patch spectral embedding into the time-domain patch tokens.

        See the `_freq_mode` note in MLPEncoder2D.__init__: "linear" reproduces the
        legacy behaviour and is expressively inert; "spectral" is genuinely
        nonlinear. Shapes are identical, so keep the mode consistent with the
        checkpoint being loaded.
        """
        import torch as _torch
        x_fft = _torch.fft.rfft(x_patches, dim=-1)
        if self._freq_mode == "spectral":
            feats = _torch.cat([x_fft.abs(), x_fft.angle()], dim=-1)
            freq_tok = _torch.nn.functional.gelu(self.freq_embed(feats.to(tok.dtype)))
        else:
            feats = _torch.cat([x_fft.real, x_fft.imag], dim=-1)
            freq_tok = self.freq_embed(feats.to(tok.dtype))
        return self.concat_proj(_torch.cat([tok, freq_tok], dim=-1))

    def forward(self, x: torch.Tensor, chan_id: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
        """
        NOTE: chan_id is needed to reconstruct the 2D positional encodings. During processing, the time series are batched into shape (C,L), where C is the TOTAL number of channels across the batched examples
        For example, if the batch contains 2 examples, each with 3 channels, then C = 6. chan_id is then [0, 1, 2, 0, 1, 2], which results in similar channel positional encoding for similar channel. L is the length 
        (padded to the same length across all series in the batch). Now, we will use chan_id to identify which of the 6 channels correspond to the same original channel across the 2 examples, so that we can assign 
        them the same channel positional encoding.
        """
        C, L = x.shape
        D = self.transformer_input_dim
        # chan_id is 0-indexed and resets to 0 at the start of every sample —
        # e.g. [0,1,2, 0,1,2] for two 3-channel samples, or [0,1,2, 0..9] for a
        # 3-channel + 10-channel mixed batch. Recover per-sample channel counts
        # by finding the zero positions; this also tells us whether the batch
        # is uniform-C (fast path) or variable-C (per-sample fallback).
        if chan_id.numel() == 0:
            sample_starts = [0]
            sample_lengths = [C]
        else:
            sample_starts = (chan_id == 0).nonzero(as_tuple=True)[0].tolist()
            if not sample_starts or sample_starts[0] != 0:
                # Defensive: if for some reason chan_id doesn't start at 0,
                # treat the whole batch as one sample.
                sample_starts = [0]
                sample_lengths = [C]
            else:
                sample_lengths = [
                    (sample_starts[i + 1] if i + 1 < len(sample_starts) else C) - sample_starts[i]
                    for i in range(len(sample_starts))
                ]
        is_uniform_c = len(set(sample_lengths)) == 1
        n_ch = sample_lengths[0]                     # only meaningful when uniform
        B = len(sample_lengths)                       # true batch size, even when variable C

        if L % self.patch_size != 0:
            raise ValueError(f"L={L} not divisible by patch_size={self.patch_size}")

        T = L // self.patch_size

        x_bc = x.reshape(C, 1, L)
        tok = self.patch_embed(x_bc)           # (C, D, T)
        tok = tok.transpose(1, 2)              # (C, T, D)

        # ── Frequency embedding branch ────────────────────────────────────────
        if self._use_freq_embed:
            # reshape into patches, apply rfft, stack real+imag
            x_patches = x.reshape(C, T, self.patch_size)          # (C, T, P)
            tok = self._fuse_freq(tok, x_patches)

        t_pos = torch.arange(T, device=tok.device, dtype=torch.float32)
        c_pos = chan_id.to(device=tok.device, dtype=torch.float32)

        if getattr(self.args, "use_only_time_pe", False):
            pe_t = _get_1d_sincos_pos_embed(D, t_pos, base=self.pe_base_t).to(tok.dtype)
            pe = pe_t.unsqueeze(0).expand(C, -1, -1)
        elif getattr(self.args, "use_only_channel_pe", False):
            pe_c = _get_1d_sincos_pos_embed(D, c_pos, base=self.pe_base_c).to(tok.dtype)
            pe = pe_c.unsqueeze(1).expand(-1, T, -1)
        else:
            if D % 4 == 0:
                pe_t = _get_1d_sincos_pos_embed(D // 2, t_pos, base=self.pe_base_t).to(tok.dtype)
                pe_c = _get_1d_sincos_pos_embed(D // 2, c_pos, base=self.pe_base_c).to(tok.dtype)
                pe = torch.cat(
                    [
                        pe_t.unsqueeze(0).expand(C, -1, -1),
                        pe_c.unsqueeze(1).expand(-1, T, -1),
                    ],
                    dim=-1,
                )
            else:
                pe_t = _get_1d_sincos_pos_embed(D, t_pos, base=self.pe_base_t).to(tok.dtype)
                pe_c = _get_1d_sincos_pos_embed(D, c_pos, base=self.pe_base_c).to(tok.dtype)
                pe = pe_t.unsqueeze(0) + pe_c.unsqueeze(1)

    
        tok = tok + pe

        tok = self.input_norm(tok)
        tok = self.input_dropout(tok)

        # ── Cross-channel attention ──────────────────────────────────────────
        # Fast path: every sample in the batch has the same channel count, so
        # we can batch all B samples together as (B, T*n_ch, D).
        # Fallback: variable-C batches (e.g. synthetic MV with C∈[3..10] mixed)
        # require per-sample iteration — the reshape (B, n_ch, T, D) is
        # invalid when n_ch isn't constant across the batch.
        if is_uniform_c:
            # (C, T, D) = (B*n_ch, T, D) → (B, n_ch, T, D) → (B, T, n_ch, D) → (B, T*n_ch, D)
            tok = tok.reshape(B, n_ch, T, D)
            tok = tok.permute(0, 2, 1, 3).reshape(B, T * n_ch, D)

            seq_len = T * n_ch
            t_idx = torch.arange(seq_len, device=tok.device)
            can_attend = (t_idx // n_ch).unsqueeze(1) >= (t_idx // n_ch).unsqueeze(0)
            causal_mask = torch.zeros(seq_len, seq_len, dtype=tok.dtype, device=tok.device)
            causal_mask.masked_fill_(~can_attend, float('-inf'))

            src_key_padding_mask = None
            if attention_mask is not None:
                bool_pad = ~attention_mask.bool()                          # (C, T) = (B*n_ch, T)
                bool_pad = bool_pad.reshape(B, n_ch, T).permute(0, 2, 1)  # (B, T, n_ch)
                bool_pad = bool_pad.reshape(B, T * n_ch)
                src_key_padding_mask = torch.zeros(
                    bool_pad.shape, dtype=tok.dtype, device=tok.device
                ).masked_fill_(bool_pad, float('-inf'))

            tok = self.encoder(
                tok,
                mask=causal_mask,
                src_key_padding_mask=src_key_padding_mask,
                is_causal=False,
            )

            # (B, T*n_ch, D) → (B, T, n_ch, D) → (B, n_ch, T, D) → (B*n_ch, T, D)
            tok = tok.reshape(B, T, n_ch, D)
            tok = tok.permute(0, 2, 1, 3).reshape(C, T, D)
        else:
            # Variable-C fallback: process each sample's MV block separately.
            # Each sample has its own (1, T*n_ch_j, D) sequence and its own
            # causal mask. Slower than the batched path (B kernel launches
            # instead of 1) but correct for arbitrary per-sample channel counts.
            out_parts = []
            for start, n_ch_j in zip(sample_starts, sample_lengths):
                sample_tok = tok[start:start + n_ch_j].unsqueeze(0)       # (1, n_ch_j, T, D)
                sample_tok = sample_tok.permute(0, 2, 1, 3).reshape(1, T * n_ch_j, D)

                seq_len_j = T * n_ch_j
                t_idx_j = torch.arange(seq_len_j, device=tok.device)
                can_attend_j = (t_idx_j // n_ch_j).unsqueeze(1) >= (t_idx_j // n_ch_j).unsqueeze(0)
                causal_mask_j = torch.zeros(
                    seq_len_j, seq_len_j, dtype=tok.dtype, device=tok.device
                ).masked_fill_(~can_attend_j, float('-inf'))

                src_key_padding_mask_j = None
                if attention_mask is not None:
                    sample_mask = attention_mask[start:start + n_ch_j]    # (n_ch_j, T)
                    bool_pad_j = ~sample_mask.bool()
                    bool_pad_j = bool_pad_j.reshape(1, n_ch_j, T).permute(0, 2, 1)
                    bool_pad_j = bool_pad_j.reshape(1, T * n_ch_j)
                    src_key_padding_mask_j = torch.zeros(
                        bool_pad_j.shape, dtype=tok.dtype, device=tok.device
                    ).masked_fill_(bool_pad_j, float('-inf'))

                sample_tok = self.encoder(
                    sample_tok,
                    mask=causal_mask_j,
                    src_key_padding_mask=src_key_padding_mask_j,
                    is_causal=False,
                )

                # (1, T*n_ch_j, D) → (1, T, n_ch_j, D) → (1, n_ch_j, T, D) → (n_ch_j, T, D)
                sample_tok = sample_tok.reshape(1, T, n_ch_j, D)
                sample_tok = sample_tok.permute(0, 2, 1, 3).reshape(n_ch_j, T, D)
                out_parts.append(sample_tok)
            tok = torch.cat(out_parts, dim=0)     # (C, T, D)

        # ── Output projection ────────────────────────────────────────────────
        tok = self.output_proj(tok)          # (C, T, output_dim)

        return tok