Temporal Fusion Transformer

"""The three distinctive pieces of TFT, transcribed literally.

  1. Gated Residual Network - the component that can switch itself off
  2. Variable Selection Network - a softmax over INPUTS, per timestep
  3. Quantile (pinball) loss - what makes the output a distribution

Attention and the LSTM are standard and covered elsewhere in this reference;
what is specific to TFT is these three, so these three are what is written
out. Plain loops, no libraries.
"""

import math


def sigmoid(value):
    if value >= 0.0:
        return 1.0 / (1.0 + math.exp(-value))
    positive = math.exp(value)
    return positive / (1.0 + positive)


def elu(value):
    """ELU rather than ReLU, deliberately: a GRN that has gated itself off
    should pass small negative values through smoothly rather than clipping
    them, so the residual path stays informative."""
    return value if value > 0.0 else math.exp(value) - 1.0


def matvec(weight, vector, bias):
    """weight is out_dim rows of in_dim."""
    return [
        bias[i] + sum(w * v for w, v in zip(row, vector))
        for i, row in enumerate(weight)
    ]


def layer_norm(vector, gain, shift, epsilon=1e-6):
    mean = sum(vector) / len(vector)
    variance = sum((value - mean) ** 2 for value in vector) / len(vector)
    denominator = math.sqrt(variance + epsilon)
    return [
        (value - mean) / denominator * gain[i] + shift[i]
        for i, value in enumerate(vector)
    ]


def gated_linear_unit(vector, weight_value, bias_value, weight_gate, bias_gate):
    """GLU: one branch is the value, the other is a gate in [0, 1].

    The gate CAN go to zero, and that is the point - it lets the network
    delete a whole sublayer rather than having to learn the identity through
    it. On the small datasets forecasting actually has, switching capacity
    off is frequently the right answer.
    """
    values = matvec(weight_value, vector, bias_value)
    gates = matvec(weight_gate, vector, bias_gate)
    return [value * sigmoid(gate) for value, gate in zip(values, gates)]


def gated_residual_network(primary, context, params):
    """GRN(a, c) = LayerNorm(a + GLU(eta)), eta = ELU(W1 a + W2 c).

    `context` is optional static conditioning - a static covariate enters
    here rather than being repeated at every timestep, which is both cheaper
    and, more importantly, keeps it recognizably static.
    """
    hidden = matvec(params['w1'], primary, params['b1'])
    if context is not None:
        # No bias on the context projection: the primary branch already has
        # one and two would be redundant.
        contribution = matvec(params['w2'], context, [0.0] * len(hidden))
        hidden = [a + b for a, b in zip(hidden, contribution)]
    hidden = [elu(value) for value in hidden]

    gated = gated_linear_unit(
        hidden, params['w_value'], params['b_value'],
        params['w_gate'], params['b_gate'],
    )

    # The residual is what makes "switch yourself off" meaningful: with the
    # gate at zero, the whole component becomes LayerNorm(a).
    residual = [a + g for a, g in zip(primary, gated)]
    return layer_norm(residual, params['ln_gain'], params['ln_shift'])


def softmax(scores):
    peak = max(scores)
    exponentials = [math.exp(score - peak) for score in scores]
    total = sum(exponentials)
    return [value / total for value in exponentials]


def variable_selection(variable_embeddings, static_context, params):
    """A softmax over VARIABLES, computed fresh at every timestep.

    variable_embeddings: one vector per input variable, all the same width.
    Returns the weighted combination and the weights themselves - the weights
    are the model's interpretability story and have to come back out.
    """
    flattened = [value for embedding in variable_embeddings for value in embedding]

    # Selection is conditioned on the static context, so a large store and a
    # small store can weight promotions differently.
    scores = gated_residual_network(flattened, static_context, params['selection'])
    weights = softmax(scores[:len(variable_embeddings)])

    # Each variable also passes through its OWN GRN before being combined,
    # so selection and transformation are separate jobs.
    transformed = [
        gated_residual_network(embedding, None, params['per_variable'][index])
        for index, embedding in enumerate(variable_embeddings)
    ]

    width = len(transformed[0])
    combined = [0.0] * width
    for weight, vector in zip(weights, transformed):
        for d in range(width):
            combined[d] += weight * vector[d]

    return combined, weights


def quantile_loss(predictions, target, quantiles):
    """Pinball loss, summed over quantile levels.

    For q = 0.9, under-predicting costs 9x what over-predicting costs, so the
    minimizer sits at the 90th percentile. That asymmetry is the entire
    mechanism - no distribution is ever assumed.
    """
    total = 0.0
    for prediction, q in zip(predictions, quantiles):
        error = target - prediction
        total += max(q * error, (q - 1.0) * error)
    return total


def masked_quantile_loss(predictions, targets, observed, quantiles):
    """Real panels are ragged. An unmasked loss trains on the padding."""
    total = 0.0
    counted = 0
    for series in range(len(predictions)):
        for horizon in range(len(predictions[series])):
            if not observed[series][horizon]:
                continue
            total += quantile_loss(
                predictions[series][horizon], targets[series][horizon], quantiles
            )
            counted += 1
    if counted == 0:
        raise ValueError("every target in the batch is masked out")
    return total / counted


def repair_crossing(prediction_row):
    """Quantile levels are trained INDEPENDENTLY, so nothing stops the 90th
    percentile coming out below the 50th. Sorting is the crude repair and it
    is what most production systems actually do."""
    return sorted(prediction_row)

"""The same three components, with the input-type contract made a type.

The substantive change is not the arithmetic. It is that the static /
known-future / observed-past split - which is TFT's central idea and its most
common production failure - stops being a convention the pipeline is trusted
to honour and becomes a structure that checks itself.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from enum import Enum
from typing import NamedTuple, Sequence


class InputKind(Enum):
    """The classification the whole architecture is built on.

    Getting this wrong is the model's signature failure: label an observed
    covariate as known-future and the backtest is invalidated completely,
    silently, and in the optimistic direction.
    """

    STATIC = 'static'
    KNOWN_FUTURE = 'known-future'
    OBSERVED_PAST = 'observed-past'


class LeakageError(ValueError):
    """Raised when a covariate is used beyond the point it is available.

    Its own type because the consequence is specific and severe: a model that
    has leaked does not fail, it reports excellent validation accuracy.
    """


class ShapeMismatch(ValueError):
    """Raised on a shape violation instead of broadcasting past it."""


@dataclass(frozen=True)
class InputSpec:
    """One covariate and where it may legitimately be read."""

    name: str
    kind: InputKind
    cardinality: int | None = None  # None means continuous

    @property
    def is_categorical(self) -> bool:
        return self.cardinality is not None


@dataclass(frozen=True)
class InputSchema:
    """The contract between the feature pipeline and the model.

    Frozen and validated, so a covariate cannot drift between categories
    between training and serving without the mismatch being visible.
    """

    specs: tuple[InputSpec, ...]

    def __post_init__(self) -> None:
        names = [spec.name for spec in self.specs]
        if len(names) != len(set(names)):
            raise ShapeMismatch('duplicate covariate names in the schema')
        if not self.by_kind(InputKind.OBSERVED_PAST):
            raise ShapeMismatch('a forecaster needs at least one observed history')

    def by_kind(self, kind: InputKind) -> tuple[InputSpec, ...]:
        return tuple(spec for spec in self.specs if spec.kind is kind)

    def encoder_variables(self) -> tuple[InputSpec, ...]:
        """History sees everything: the past is fully observed."""
        return tuple(
            spec for spec in self.specs if spec.kind is not InputKind.STATIC
        )

    def decoder_variables(self) -> tuple[InputSpec, ...]:
        """The future sees ONLY known-future covariates. This method is the
        whole leakage guarantee, expressed once rather than per call site."""
        return self.by_kind(InputKind.KNOWN_FUTURE)

    def fingerprint(self) -> tuple[tuple[str, str], ...]:
        """Comparable across training and serving, so a silently reclassified
        covariate is caught at load time rather than never."""
        return tuple((spec.name, spec.kind.value) for spec in self.specs)


def assert_schema_matches(trained: InputSchema, serving: InputSchema) -> None:
    """Guard clause for the failure that has no other symptom."""
    if trained.fingerprint() != serving.fingerprint():
        raise LeakageError(
            'serving schema differs from the trained schema; a covariate has '
            'changed kind and the model is no longer valid'
        )


@dataclass(frozen=True)
class QuantileSet:
    """Quantile levels, validated and kept sorted.

    Sorted because the crossing repair depends on it, and validated because a
    level outside (0, 1) produces a loss that is silently unbounded.
    """

    levels: tuple[float, ...] = (0.1, 0.5, 0.9)

    def __post_init__(self) -> None:
        if not self.levels:
            raise ShapeMismatch('at least one quantile level is required')
        if any(not 0.0 < level < 1.0 for level in self.levels):
            raise ShapeMismatch('quantile levels must lie strictly inside (0, 1)')
        if list(self.levels) != sorted(self.levels):
            raise ShapeMismatch('quantile levels must be given in ascending order')

    @property
    def median_index(self) -> int:
        """Where the point forecast lives, if one is needed at all."""
        return min(
            range(len(self.levels)),
            key=lambda i: abs(self.levels[i] - 0.5),
        )


class SelectionResult(NamedTuple):
    """The combined vector AND the weights.

    The weights are the interpretability story and must come back out rather
    than being discarded inside the layer - but they are importance to this
    model, not causal effect, and the name says combined first for a reason.
    """

    combined: list[float]
    weights: list[float]


@dataclass
class GrnParams:
    """One gated residual network's parameters."""

    w1: list[list[float]]
    b1: list[float]
    w_value: list[list[float]]
    b_value: list[float]
    w_gate: list[list[float]]
    b_gate: list[float]
    ln_gain: list[float]
    ln_shift: list[float]
    w_context: list[list[float]] | None = field(default=None)


def sigmoid(value: float) -> float:
    """Branch on sign so exp cannot overflow on a large negative input."""
    if value >= 0.0:
        return 1.0 / (1.0 + math.exp(-value))
    positive = math.exp(value)
    return positive / (1.0 + positive)


def elu(value: float) -> float:
    """ELU, not ReLU: a gated-off GRN should pass small negatives through
    smoothly rather than clipping them, keeping the residual path useful."""
    return value if value > 0.0 else math.expm1(value)


def matvec(weight: Sequence[Sequence[float]], vector: Sequence[float],
           bias: Sequence[float] | None = None) -> list[float]:
    if weight and len(weight[0]) != len(vector):
        raise ShapeMismatch(f'weight expects {len(weight[0])} inputs, got {len(vector)}')
    return [
        (bias[i] if bias is not None else 0.0) + math.fsum(
            w * v for w, v in zip(row, vector)
        )
        for i, row in enumerate(weight)
    ]


def layer_norm(vector: Sequence[float], gain: Sequence[float],
               shift: Sequence[float], epsilon: float = 1e-6) -> list[float]:
    mean = math.fsum(vector) / len(vector)
    variance = math.fsum((value - mean) ** 2 for value in vector) / len(vector)
    denominator = math.sqrt(variance + epsilon)
    return [
        (value - mean) / denominator * g + s
        for value, g, s in zip(vector, gain, shift)
    ]


def gated_residual_network(
    primary: Sequence[float],
    params: GrnParams,
    context: Sequence[float] | None = None,
) -> list[float]:
    """GRN(a, c) = LayerNorm(a + GLU(ELU(W1 a + W2 c))).

    With the gate at zero the whole component collapses to LayerNorm(a),
    which is what lets the network choose to be smaller. On the dataset sizes
    forecasting actually has, that is frequently the right choice.
    """
    if context is not None and params.w_context is None:
        raise ShapeMismatch('context supplied but this GRN has no context projection')

    hidden = matvec(params.w1, primary, params.b1)
    if context is not None and params.w_context is not None:
        # No bias here: the primary branch already carries one.
        for i, value in enumerate(matvec(params.w_context, context)):
            hidden[i] += value
    hidden = [elu(value) for value in hidden]

    values = matvec(params.w_value, hidden, params.b_value)
    gates = matvec(params.w_gate, hidden, params.b_gate)
    gated = [value * sigmoid(gate) for value, gate in zip(values, gates)]

    residual = [a + g for a, g in zip(primary, gated)]
    return layer_norm(residual, params.ln_gain, params.ln_shift)


def stable_softmax(scores: Sequence[float]) -> list[float]:
    peak = max(scores)
    exponentials = [math.exp(score - peak) for score in scores]
    total = math.fsum(exponentials)
    return [value / total for value in exponentials]


def variable_selection(
    embeddings: Sequence[Sequence[float]],
    selection_params: GrnParams,
    per_variable_params: Sequence[GrnParams],
    static_context: Sequence[float] | None = None,
) -> SelectionResult:
    """Softmax over VARIABLES, computed fresh at every timestep."""
    if len(embeddings) != len(per_variable_params):
        raise ShapeMismatch('one per-variable GRN is required for each variable')

    flattened = [value for embedding in embeddings for value in embedding]
    scores = gated_residual_network(flattened, selection_params, static_context)
    weights = stable_softmax(scores[:len(embeddings)])

    # Selection and transformation are separate jobs: each variable passes
    # through its own GRN before the weighted combination.
    transformed = [
        gated_residual_network(embedding, params)
        for embedding, params in zip(embeddings, per_variable_params)
    ]

    width = len(transformed[0])
    combined = [0.0] * width
    for weight, vector in zip(weights, transformed):
        for d in range(width):
            combined[d] += weight * vector[d]

    return SelectionResult(combined=combined, weights=weights)


def quantile_loss(
    predictions: Sequence[float],
    target: float,
    quantiles: QuantileSet,
) -> float:
    """Pinball loss summed over levels.

    For q = 0.9 under-prediction costs nine times over-prediction, so the
    minimizer is the 90th percentile. No distribution is assumed anywhere.
    """
    if len(predictions) != len(quantiles.levels):
        raise ShapeMismatch(
            f'{len(predictions)} predictions for {len(quantiles.levels)} levels'
        )
    total = 0.0
    for prediction, level in zip(predictions, quantiles.levels):
        error = target - prediction
        total += max(level * error, (level - 1.0) * error)
    return total


def repair_crossing(predictions: Sequence[float]) -> list[float]:
    """Levels are trained independently, so nothing prevents the 90th
    percentile falling below the 50th. Sorting is the crude repair and it is
    what most production systems do - but it should be COUNTED, because
    frequent crossing means the tail quantiles are not trustworthy."""
    return sorted(predictions)


def crossing_rate(batch: Sequence[Sequence[float]]) -> float:
    """Monitor rather than silently repair: this belongs on a dashboard."""
    if not batch:
        return 0.0
    crossed = sum(1 for row in batch if list(row) != sorted(row))
    return crossed / len(batch)

"""Batched over series, horizons and variables. Every GRN becomes a GEMM.

Three changes:

  1. A GRN is two matrix products and pointwise arithmetic. Batched across
     (series x timestep) it is one GEMM, and the value and gate branches
     share ONE weight with 2 * d columns so the pair costs one product.
  2. The per-variable GRNs, which the loop version runs V times, become a
     single batched matmul over a (V, d_in, d_out) weight stack - einsum,
     not a loop, because V is typically a dozen and the loop overhead
     dominates the arithmetic at that size.
  3. Quantile loss vectorizes exactly: the pinball max is a where over a
     sign, so every level and horizon is scored in one pass.
"""

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray

FLOAT = np.float32


def elu_(values: NDArray[np.float32]) -> NDArray[np.float32]:
    """In-place ELU. expm1 rather than exp minus one: for small negative
    inputs the naive form loses every significant digit."""
    negative = values < 0.0
    values[negative] = np.expm1(values[negative])
    return values


def gated_residual_network(
    primary: NDArray[np.float32],
    w1: NDArray[np.float32],
    b1: NDArray[np.float32],
    w_gated: NDArray[np.float32],
    b_gated: NDArray[np.float32],
    gain: NDArray[np.float32],
    shift: NDArray[np.float32],
    context: NDArray[np.float32] | None = None,
    w_context: NDArray[np.float32] | None = None,
) -> NDArray[np.float32]:
    """primary is (..., d). Everything leading batches.

    w_gated is (d_hidden, 2 * d): value and gate side by side, so the GLU
    costs one GEMM rather than two over the same input.
    """
    hidden = primary @ w1
    hidden += b1
    if context is not None and w_context is not None:
        # Static context broadcasts over the time axis rather than being
        # tiled - the whole reason static covariates are kept separate.
        hidden += context @ w_context

    elu_(hidden)

    projected = hidden @ w_gated
    projected += b_gated
    width = projected.shape[-1] // 2
    value = projected[..., :width]
    gate = projected[..., width:]

    # Fused sigmoid-and-multiply, in place over the gate buffer, so the
    # gated result reuses memory the projection already owns.
    np.negative(gate, out=gate)
    np.exp(gate, out=gate)
    gate += 1.0
    np.divide(value, gate, out=value)
    value += primary

    # Layer norm with the two-pass variance: the one-pass E[x^2] - E[x]^2
    # form cancels catastrophically in float32 once the mean is large
    # relative to the spread, which is where residual streams end up.
    mean = value.mean(axis=-1, keepdims=True, dtype=np.float32)
    value -= mean
    variance = np.mean(value * value, axis=-1, keepdims=True, dtype=np.float32)
    value /= np.sqrt(variance + FLOAT(1e-6))
    value *= gain
    value += shift
    return value


def variable_selection(
    embeddings: NDArray[np.float32],
    selection_weights: tuple[NDArray[np.float32], ...],
    per_variable_w1: NDArray[np.float32],
    per_variable_b1: NDArray[np.float32],
    per_variable_gated: NDArray[np.float32],
    per_variable_gain: NDArray[np.float32],
    per_variable_shift: NDArray[np.float32],
    static_context: NDArray[np.float32] | None = None,
) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
    """embeddings is (batch, time, num_variables, d).

    Returns the combined vector and the selection weights, which have to come
    back out - they are the interpretability story, and they are importance
    to this model rather than causal effect.
    """
    batch, steps, num_variables, width = embeddings.shape

    # Selection scores from the flattened per-variable embeddings.
    flattened = embeddings.reshape(batch, steps, num_variables * width)
    scores = gated_residual_network(
        flattened, *selection_weights, context=static_context,
    )[..., :num_variables]

    scores -= scores.max(axis=-1, keepdims=True)
    np.exp(scores, out=scores)
    scores /= scores.sum(axis=-1, keepdims=True)

    # Every per-variable GRN at once: a (V, d, d_hidden) weight stack
    # contracted against a (B, T, V, d) input. The loop version runs V small
    # GEMMs; this runs one batched contraction, and with V around a dozen
    # the loop overhead was the dominant cost.
    hidden = np.einsum('btvd,vdh->btvh', embeddings, per_variable_w1, optimize=True)
    hidden += per_variable_b1
    elu_(hidden)

    projected = np.einsum('btvh,vhg->btvg', hidden, per_variable_gated, optimize=True)
    half = projected.shape[-1] // 2
    value, gate = projected[..., :half], projected[..., half:]
    np.negative(gate, out=gate)
    np.exp(gate, out=gate)
    gate += 1.0
    np.divide(value, gate, out=value)
    value += embeddings

    mean = value.mean(axis=-1, keepdims=True, dtype=np.float32)
    value -= mean
    variance = np.mean(value * value, axis=-1, keepdims=True, dtype=np.float32)
    value /= np.sqrt(variance + FLOAT(1e-6))
    value *= per_variable_gain
    value += per_variable_shift

    # The weighted combination is a contraction over the variable axis, not
    # a loop with an accumulator.
    combined = np.einsum('btv,btvd->btd', scores, value, optimize=True)
    return combined, scores


def quantile_loss(
    predictions: NDArray[np.float32],
    targets: NDArray[np.float32],
    levels: NDArray[np.float32],
    observed: NDArray[np.bool_],
) -> np.float64:
    """predictions (batch, horizon, num_quantiles), targets (batch, horizon).

    The pinball max reduces to a single where on the sign of the error, so
    every level and every horizon is scored in one pass with no branching.
    """
    counted = observed.sum(dtype=np.int64)
    if counted == 0:
        raise ValueError('every target in the batch is masked out')

    error = targets[..., None] - predictions
    # max(q*e, (q-1)*e) is q*e when e >= 0 and (q-1)*e otherwise. Written as
    # a select rather than two products and a maximum, which halves the
    # arithmetic and materializes one fewer full-size temporary.
    loss = np.where(error >= 0.0, levels * error, (levels - 1.0) * error)
    loss = loss.sum(axis=-1)
    loss *= observed
    return loss.sum(dtype=np.float64) / counted


def repair_crossing(predictions: NDArray[np.float32]) -> NDArray[np.float32]:
    """Sort along the quantile axis. Levels are trained independently, so
    nothing prevents the 90th percentile falling below the 50th."""
    return np.sort(predictions, axis=-1)


def crossing_rate(predictions: NDArray[np.float32]) -> float:
    """Monitor rather than silently repair: frequent crossing means the tail
    quantiles have stopped being trustworthy, and sorting hides that."""
    ordered = np.diff(predictions, axis=-1) >= 0.0
    return float(1.0 - ordered.all(axis=-1).mean())