source/project/models/close_distribution_v2.py

"""Continuous-horizon sequence models for forecasting-v4."""

from __future__ import annotations

import math
from collections.abc import Sequence

import torch
from torch import nn


class PooledLinear(nn.Module):
    def __init__(
        self,
        *,
        context_length: int,
        output_horizons: Sequence[float],
        classes: int,
        bins: int = 8,
    ) -> None:
        super().__init__()
        if context_length % bins:
            raise ValueError("context length must be divisible by pool bins")
        self.bins = bins
        self.horizon_count = len(output_horizons)
        self.classes = classes
        outputs = self.horizon_count * classes
        self.output = nn.Linear(bins * 5, outputs)
        self.estimated_madds = (
            context_length * 4 + bins * 5 * outputs
        )

    def forward(
        self,
        tokens: torch.Tensor,
        padding_mask: torch.Tensor,
        horizons: torch.Tensor | None = None,
    ) -> torch.Tensor:
        del padding_mask
        if horizons is not None:
            raise ValueError("fixed pooled control cannot query horizons")
        if tokens.shape[2] != 4:
            raise ValueError("pooled control requires baseline plus padding")
        batch, length, _channels = tokens.shape
        segment = tokens.reshape(
            batch,
            self.bins,
            length // self.bins,
            tokens.shape[2],
        )
        observed = segment[..., 1]
        valid = segment[..., 3]
        observed_count = observed.sum(dim=2).clamp_min(1.0)
        valid_count = valid.sum(dim=2).clamp_min(1.0)
        return_mean = (segment[..., 0] * observed).sum(dim=2) / observed_count
        return_rms = torch.sqrt(
            (segment[..., 0].square() * observed).sum(dim=2)
            / observed_count
            + 1e-8
        )
        state = torch.stack(
            (
                return_mean,
                return_rms,
                observed.mean(dim=2),
                (segment[..., 2] * valid).sum(dim=2) / valid_count,
                valid.mean(dim=2),
            ),
            dim=2,
        ).reshape(batch, -1)
        return self.output(state).reshape(
            batch,
            self.horizon_count,
            self.classes,
        )


class CausalResidualBlock(nn.Module):
    def __init__(self, width: int, dilation: int) -> None:
        super().__init__()
        self.padding = 2 * dilation
        self.convolution = nn.Conv1d(
            width,
            width,
            kernel_size=3,
            dilation=dilation,
            padding=self.padding,
        )
        self.normalization = nn.GroupNorm(1, width)
        self.activation = nn.GELU()

    def forward(self, values: torch.Tensor) -> torch.Tensor:
        residual = values
        values = self.convolution(values)
        values = values[:, :, : -self.padding]
        return self.activation(self.normalization(values)) + residual


class CausalTCN(nn.Module):
    def __init__(
        self,
        *,
        context_length: int,
        channels: int,
        output_horizons: Sequence[float],
        classes: int,
        target_parameters: int,
    ) -> None:
        super().__init__()
        blocks = max(1, math.ceil(math.log2(context_length)))
        outputs = len(output_horizons) * classes
        width = 8
        for candidate in range(8, 513):
            estimate = (
                channels * candidate
                + candidate
                + blocks * (3 * candidate * candidate + 3 * candidate)
                + candidate * outputs
                + outputs
            )
            if estimate > target_parameters:
                break
            width = candidate
        self.width = width
        self.block_count = blocks
        self.input_projection = nn.Conv1d(channels, width, kernel_size=1)
        self.blocks = nn.Sequential(
            *[
                CausalResidualBlock(width, 2**index)
                for index in range(blocks)
            ]
        )
        self.output = nn.Linear(width, outputs)
        self.horizon_count = len(output_horizons)
        self.classes = classes
        self.estimated_madds = (
            context_length * channels * width
            + blocks * context_length * 3 * width * width
            + width * outputs
        )

    def forward(
        self,
        tokens: torch.Tensor,
        padding_mask: torch.Tensor,
        horizons: torch.Tensor | None = None,
    ) -> torch.Tensor:
        del padding_mask
        if horizons is not None:
            raise ValueError("fixed-head TCN cannot query custom horizons")
        values = self.input_projection(tokens.transpose(1, 2))
        values = self.blocks(values)
        result = self.output(values[:, :, -1])
        return result.reshape(-1, self.horizon_count, self.classes)


class ContinuousHorizonTransformer(nn.Module):
    def __init__(
        self,
        *,
        context_length: int,
        channels: int,
        output_horizons: Sequence[float],
        classes: int,
        target_parameters: int,
        query_mode: str,
        rotary_base: float,
    ) -> None:
        super().__init__()
        if query_mode not in ("learned", "scalar_mlp", "rotary"):
            raise ValueError("invalid horizon query mode")
        if rotary_base <= 1.0:
            raise ValueError("rotary base must exceed one")
        output_horizons = tuple(float(value) for value in output_horizons)
        if any(value <= 0.0 for value in output_horizons):
            raise ValueError("output horizons must be positive")

        width = 16
        closest = float("inf")
        for candidate in range(16, 257, 4):
            estimate = (
                20 * candidate * candidate
                + candidate
                * (
                    context_length
                    + channels
                    + len(output_horizons)
                    + classes
                    + 27
                )
                + classes
            )
            distance = abs(estimate - target_parameters)
            if distance < closest:
                width = candidate
                closest = distance
        if width % 2:
            raise ValueError("continuous query width must be even")
        heads = 4 if width % 4 == 0 else 2
        self.width = width
        self.query_mode = query_mode
        self.rotary_base = float(rotary_base)
        self.classes = classes
        self.input_projection = nn.Linear(channels, width)
        self.position = nn.Parameter(torch.empty(context_length, width))
        layer = nn.TransformerEncoderLayer(
            d_model=width,
            nhead=heads,
            dim_feedforward=width * 2,
            dropout=0.0,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=2)
        self.query_attention = nn.MultiheadAttention(
            width,
            heads,
            dropout=0.0,
            batch_first=True,
        )
        self.output = nn.Linear(width, classes)
        self.register_buffer(
            "output_horizons",
            torch.tensor(output_horizons, dtype=torch.float32),
            persistent=True,
        )
        if query_mode == "learned":
            self.learned_queries = nn.Parameter(
                torch.empty(len(output_horizons), width)
            )
        elif query_mode == "scalar_mlp":
            scalar_width = max(4, width // 4)
            self.scalar_query = nn.Sequential(
                nn.Linear(1, scalar_width),
                nn.GELU(),
                nn.Linear(scalar_width, width),
            )
        else:
            self.rotary_query = nn.Parameter(torch.empty(width))
            self.linear_query = nn.Parameter(torch.empty(width))
            pair = torch.arange(width // 2, dtype=torch.float32)
            self.register_buffer(
                "rotary_frequency",
                self.rotary_base ** (-2.0 * pair / width),
                persistent=True,
            )
        nn.init.normal_(self.position, std=0.02)
        if query_mode == "learned":
            nn.init.normal_(self.learned_queries, std=0.02)
        elif query_mode == "rotary":
            nn.init.normal_(self.rotary_query, std=0.02)
            nn.init.normal_(self.linear_query, std=0.02)

        horizon_count = len(output_horizons)
        encoder_madds = 2 * (
            8 * context_length * width * width
            + 2 * context_length * context_length * width
        )
        query_madds = (
            2 * (horizon_count + context_length) * width * width
            + 2 * horizon_count * context_length * width
        )
        self.estimated_madds = (
            context_length * channels * width
            + encoder_madds
            + query_madds
            + horizon_count * width * classes
        )

    def _queries(self, horizons: torch.Tensor) -> torch.Tensor:
        values = horizons.to(
            device=self.position.device,
            dtype=self.position.dtype,
        )
        coordinate = torch.log2(values / 2.0)
        if self.query_mode == "learned":
            if (
                len(values) != len(self.output_horizons)
                or not torch.allclose(
                    values,
                    self.output_horizons.to(values),
                )
            ):
                raise ValueError(
                    "learned queries only support configured horizons"
                )
            return self.learned_queries
        if self.query_mode == "scalar_mlp":
            return self.scalar_query(coordinate[:, None])

        angles = coordinate[:, None] * self.rotary_frequency[None, :]
        base = self.rotary_query.reshape(-1, 2)
        even = base[:, 0][None, :]
        odd = base[:, 1][None, :]
        cosine = torch.cos(angles)
        sine = torch.sin(angles)
        rotated_even = even * cosine - odd * sine
        rotated_odd = even * sine + odd * cosine
        rotated = torch.stack(
            (rotated_even, rotated_odd),
            dim=2,
        ).reshape(len(values), self.width)
        return rotated + coordinate[:, None] * self.linear_query[None, :]

    def forward(
        self,
        tokens: torch.Tensor,
        padding_mask: torch.Tensor,
        horizons: torch.Tensor | None = None,
    ) -> torch.Tensor:
        values = self.input_projection(tokens) + self.position[None, :, :]
        encoded = self.encoder(
            values,
            src_key_padding_mask=padding_mask,
        )
        query_horizons = (
            self.output_horizons
            if horizons is None
            else horizons
        )
        queries = self._queries(query_horizons)
        queries = queries[None, :, :].expand(len(tokens), -1, -1)
        decoded, _weights = self.query_attention(
            queries,
            encoded,
            encoded,
            key_padding_mask=padding_mask,
            need_weights=False,
        )
        return self.output(decoded)


def build_model(
    name: str,
    *,
    context_length: int,
    channels: int,
    output_horizons: Sequence[float],
    classes: int,
    target_parameters: int,
    rotary_base: float,
) -> nn.Module:
    if name == "pooled_linear":
        return PooledLinear(
            context_length=context_length,
            output_horizons=output_horizons,
            classes=classes,
        )
    if name == "tcn":
        return CausalTCN(
            context_length=context_length,
            channels=channels,
            output_horizons=output_horizons,
            classes=classes,
            target_parameters=target_parameters,
        )
    query_modes = {
        "transformer_learned": "learned",
        "transformer_scalar": "scalar_mlp",
        "transformer_rotary": "rotary",
    }
    if name in query_modes:
        return ContinuousHorizonTransformer(
            context_length=context_length,
            channels=channels,
            output_horizons=output_horizons,
            classes=classes,
            target_parameters=target_parameters,
            query_mode=query_modes[name],
            rotary_base=rotary_base,
        )
    raise ValueError(f"unknown forecasting-v4 model: {name}")