TART.py

import math
from typing import Optional

import numpy as np
import jax
import jax.numpy as jnp
import flax.linen as nn
import optax
from flax.training import train_state

def zscore(x: np.ndarray, eps: float = 1e-6):
    x = np.asarray(x, dtype=np.float32)
    mu = x.mean(axis=0, keepdims=True)
    sd = x.std(axis=0, keepdims=True) + eps
    return (x - mu) / sd, (mu, sd)


def make_nextstep_windows(X: np.ndarray, L: int):
    """
    Simple 1-step-ahead dataset:
      input  X_in[i]  = X[i : i+L]      (L, D)
      target Y_out[i] = X[i+L]          (D,)
    """
    X = np.asarray(X, dtype=np.float32)
    T, D = X.shape
    if T <= L:
        raise ValueError(f"T={T} must be > L={L}")
    N = T - L
    X_in = np.stack([X[i:i+L] for i in range(N)], axis=0)   # (N, L, D)
    Y_out = np.stack([X[i+L] for i in range(N)], axis=0)    # (N, D)
    return X_in, Y_out


def choose_heads(d_model: int, target_head_dim: int = 32, max_heads: int = 8) -> int:
    """Pick a reasonable number of heads given d_model."""
    n_heads = max(1, min(max_heads, d_model // target_head_dim))
    # ensure divisibility
    while d_model % n_heads != 0 and n_heads > 1:
        n_heads -= 1
    return n_heads

def build_rotary_inv_freq(head_dim: int, base: float = 10000.0):
    assert head_dim % 2 == 0, "head_dim must be even for RoPE."
    return 1.0 / (base ** (jnp.arange(0, head_dim, 2, dtype=jnp.float32) / head_dim))


def apply_rope(x, inv_freq):
    """
    Apply RoPE to q/k.

    x:        (batch, heads, seq_len, head_dim)
    inv_freq: (head_dim/2,)
    """
    b, h, t, d = x.shape
    half = d // 2
    # positions
    positions = jnp.arange(t, dtype=jnp.float32)     # (t,)
    freqs = jnp.einsum("i,j->ij", positions, inv_freq)  # (t, half)
    cos = jnp.cos(freqs)[None, None, :, :]  # (1,1,t,half)
    sin = jnp.sin(freqs)[None, None, :, :]

    x1 = x[..., :half]   # (b,h,t,half)
    x2 = x[..., half:]   # (b,h,t,half)

    x1_rot = x1 * cos - x2 * sin
    x2_rot = x1 * sin + x2 * cos
    x_rot = jnp.concatenate([x1_rot, x2_rot], axis=-1)  # (b,h,t,d)
    return x_rot

class MultiHeadSelfAttention(nn.Module):
    d_model: int
    n_heads: int
    dropout: float = 0.0
    use_rope: bool = True

    @nn.compact
    def __call__(self, x, deterministic: bool):
        """
        x: (batch, seq_len, d_model)
        returns: (batch, seq_len, d_model)
        """
        b, t, d_model = x.shape
        assert d_model == self.d_model
        assert self.d_model % self.n_heads == 0, "d_model must be divisible by n_heads"
        head_dim = self.d_model // self.n_heads
        assert head_dim % 2 == 0, "head_dim must be even for RoPE"

        # project to qkv
        qkv = nn.Dense(3 * self.d_model, use_bias=False, name="qkv")(x)  # (b, t, 3*d_model)
        qkv = qkv.reshape(b, t, 3, self.n_heads, head_dim)
        qkv = qkv.transpose(2, 0, 3, 1, 4)  # (3, b, h, t, head_dim)
        q, k, v = qkv[0], qkv[1], qkv[2]    # each: (b, h, t, head_dim)

        if self.use_rope:
            inv_freq = build_rotary_inv_freq(head_dim)
            q = apply_rope(q, inv_freq)
            k = apply_rope(k, inv_freq)

        # scaled dot-product attention with causal mask
        scale = 1.0 / math.sqrt(head_dim)
        attn_logits = jnp.einsum("bhqd, bhkd -> bhqk", q, k) * scale  # (b,h,t,t)

        # causal mask
        mask = jnp.tril(jnp.ones((t, t), dtype=jnp.bool_))
        attn_logits = jnp.where(mask, attn_logits, -1e9)

        attn_weights = nn.softmax(attn_logits, axis=-1)
        attn_weights = nn.Dropout(rate=self.dropout)(attn_weights, deterministic=deterministic)

        attn_output = jnp.einsum("bhqk, bhkd -> bhqd", attn_weights, v)  # (b,h,t,head_dim)
        attn_output = attn_output.transpose(0, 2, 1, 3).reshape(b, t, self.d_model)  # (b,t,d_model)
        out = nn.Dense(self.d_model, name="out_proj")(attn_output)
        out = nn.Dropout(rate=self.dropout)(out, deterministic=deterministic)
        return out

class FeedForward(nn.Module):
    d_model: int
    mlp_ratio: float = 4.0
    dropout: float = 0.0

    @nn.compact
    def __call__(self, x, deterministic: bool):
        hidden_dim = int(self.d_model * self.mlp_ratio)
        x = nn.Dense(hidden_dim)(x)
        x = nn.gelu(x)
        x = nn.Dropout(rate=self.dropout)(x, deterministic=deterministic)
        x = nn.Dense(self.d_model)(x)
        x = nn.Dropout(rate=self.dropout)(x, deterministic=deterministic)
        return x


class TransformerBlock(nn.Module):
    d_model: int
    n_heads: int
    dropout: float = 0.0
    mlp_ratio: float = 4.0
    use_rope: bool = True

    @nn.compact
    def __call__(self, x, deterministic: bool):
        # Self-attention
        h = nn.LayerNorm()(x)
        h = MultiHeadSelfAttention(
            d_model=self.d_model,
            n_heads=self.n_heads,
            dropout=self.dropout,
            use_rope=self.use_rope,
        )(h, deterministic=deterministic)
        x = x + h

        # FFN
        h2 = nn.LayerNorm()(x)
        h2 = FeedForward(
            d_model=self.d_model,
            mlp_ratio=self.mlp_ratio,
            dropout=self.dropout,
        )(h2, deterministic=deterministic)
        x = x + h2
        return x

class TARTModel(nn.Module):
    """
    Core autoregressive Transformer that maps
      (batch, L, D_in) -> (batch, L, D_in)
    and we train it to predict the *next* token at each step.
    For forecasting, we use the last position as the next-step prediction.
    """
    input_dim: int
    d_model: int
    n_heads: int
    depth: int
    max_len: int
    dropout: float = 0.0
    use_learned_pos: bool = True
    use_rope: bool = True

    @nn.compact
    def __call__(self, x, deterministic: bool = True):
        """
        x: (batch, seq_len, input_dim)
        returns: (batch, seq_len, input_dim)
        """
        b, t, d_in = x.shape
        assert d_in == self.input_dim
        if t > self.max_len:
            raise ValueError(f"seq_len={t} exceeds max_len={self.max_len}")

        # Project input to d_model
        h = nn.Dense(self.d_model, name="token_embed")(x)

        # Learned positional embedding
        if self.use_learned_pos:
            pos_emb = self.param(
                "pos_emb",
                nn.initializers.normal(stddev=0.02),
                (self.max_len, self.d_model),
            )  # (max_len, d_model)
            positions = jnp.arange(t)[None, :]          # (1, t)
            h = h + pos_emb[positions, :]               # broadcast to (b,t,d_model)

        # Stacked transformer blocks
        for i in range(self.depth):
            h = TransformerBlock(
                d_model=self.d_model,
                n_heads=self.n_heads,
                dropout=self.dropout,
                mlp_ratio=4.0,
                use_rope=self.use_rope,
                name=f"block_{i}",
            )(h, deterministic=deterministic)

        h = nn.LayerNorm(name="final_ln")(h)
        out = nn.Dense(self.input_dim, name="out_proj")(h)  # (b,t,D_in)
        return out

class TART:
    """
    Time-series AutoReg Transformer (TART)

    Usage:
      forecaster = TART(R_tX, L=128, d_model=64, depth=4, ...)
      preds = forecaster(F_cX, steps=200)

    - R_tX: training time-series, shape (T, D)
    - L:    context length (sequence length seen by the Transformer)
    """

    def __init__(self,
                 R_tX: np.ndarray,
                 L: int,
                 d_model: int = 64,
                 depth: int = 4,
                 n_heads: Optional[int] = None,
                 use_learned_pos: bool = True,
                 use_rope: bool = True,
                 dropout: float = 0.0,
                 seed: int = 0,
                 val_split: float = 0.1,
                 batch_size: int = 64,
                 max_epochs: int = 50,
                 init_lr: float = 3e-4,
                 min_lr: float = 1e-5,
                 lr_decay: float = 0.5,
                 tol_rel_improve: float = 1e-3,
                 patience: int = 5):

        R_tX = np.asarray(R_tX, dtype=np.float32)
        if R_tX.ndim == 1:
            R_tX = R_tX[:, None]
        T, D = R_tX.shape
        if T <= L:
            raise ValueError(f"T={T} must be > L={L}")

        self.L = int(L)
        self.D = int(D)

        # choose d_model / n_heads
        d_model = int(max(16, d_model))
        if n_heads is None:
            n_heads = choose_heads(d_model)
        self.d_model = d_model
        self.n_heads = int(n_heads)
        self.depth = int(depth)

        # normalize training data
        F_norm, (mu, sd) = zscore(R_tX)
        self._mu = mu
        self._sd = sd

        # build 1-step-ahead dataset
        X_in, Y_out = make_nextstep_windows(F_norm, L=self.L)  # (N,L,D), (N,D)
        N = X_in.shape[0]
        n_val = max(1, int(val_split * N))
        n_tr  = N - n_val
        Xtr, Ytr = X_in[:n_tr], Y_out[:n_tr]
        Xva, Yva = X_in[n_tr:], Y_out[n_tr:]

        print(f"[TART] Train N={N} (train={n_tr}, val={n_val}) | "
              f"L={self.L} D={self.D} | d_model={d_model} heads={n_heads} depth={depth}")

        # build model
        max_len = self.L          # always feed sequences of length L
        model = TARTModel(
            input_dim=self.D,
            d_model=self.d_model,
            n_heads=self.n_heads,
            depth=self.depth,
            max_len=max_len,
            dropout=dropout,
            use_learned_pos=use_learned_pos,
            use_rope=use_rope,
        )
        self.model = model
        self.max_len = max_len

        rng = jax.random.PRNGKey(seed)
        dummy_x = jnp.zeros((1, self.L, self.D), dtype=jnp.float32)
        params = model.init(rng, dummy_x, deterministic=True)["params"]

        def create_state(lr):
            tx = optax.adamw(learning_rate=lr, weight_decay=0.0)
            return train_state.TrainState.create(apply_fn=model.apply,
                                                 params=params,
                                                 tx=tx)

        state = create_state(init_lr)
        self._rng = rng

        @jax.jit
        def train_step(state, x_batch, y_batch, rng):
            """One training step (MSE on next-step prediction of last token)."""
            dropout_rng, new_rng = jax.random.split(rng)

            def loss_fn(p):
                preds = state.apply_fn({"params": p},
                                       x_batch,
                                       deterministic=False,
                                       rngs={"dropout": dropout_rng})
                pred_last = preds[:, -1, :]            # (B,D)
                loss = jnp.mean((pred_last - y_batch) ** 2)
                return loss

            loss, grads = jax.value_and_grad(loss_fn)(state.params)
            new_state = state.apply_gradients(grads=grads)
            return new_state, loss, new_rng

        @jax.jit
        def eval_step(state, x_batch, y_batch):
            preds = state.apply_fn({"params": state.params},
                                   x_batch,
                                   deterministic=True)
            pred_last = preds[:, -1, :]
            loss = jnp.mean((pred_last - y_batch) ** 2)
            return loss

        # move training data to JAX arrays once
        Xtr_j = jnp.asarray(Xtr)
        Ytr_j = jnp.asarray(Ytr)
        Xva_j = jnp.asarray(Xva)
        Yva_j = jnp.asarray(Yva)

        best_params = state.params
        best_val = float("inf")
        curr_lr = init_lr
        epochs_no_gain = 0

        num_batches = lambda N: int(np.ceil(N / batch_size))

        for epoch in range(1, max_epochs + 1):
            # shuffle training indices
            idx = np.arange(n_tr)
            np.random.default_rng(epoch + seed).shuffle(idx)

            # training loop
            train_losses = []
            rng = self._rng
            for bi in range(num_batches(n_tr)):
                s = bi * batch_size
                e = min((bi + 1) * batch_size, n_tr)
                batch_idx = idx[s:e]
                xb = Xtr_j[batch_idx]
                yb = Ytr_j[batch_idx]
                state, loss, rng = train_step(state, xb, yb, rng)
                train_losses.append(float(loss))

            self._rng = rng
            tr_loss = float(np.mean(train_losses))

            # validation
            val_losses = []
            for bi in range(num_batches(n_val)):
                s = bi * batch_size
                e = min((bi + 1) * batch_size, n_val)
                xb = Xva_j[s:e]
                yb = Yva_j[s:e]
                val_losses.append(float(eval_step(state, xb, yb)))
            va_loss = float(np.mean(val_losses))

            print(f"[TART] epoch {epoch:03d} | train {tr_loss:.6e} | val {va_loss:.6e} | lr {curr_lr:.2e}")

            # track best
            if va_loss + 1e-8 < best_val:
                rel_gain = (best_val - va_loss) / max(best_val, 1e-8) if best_val < float("inf") else 1.0
                best_val = va_loss
                best_params = state.params
                epochs_no_gain = 0
            else:
                rel_gain = 0.0
                epochs_no_gain += 1

            # LR schedule on plateau
            if epochs_no_gain >= patience:
                if curr_lr > min_lr * (1.0 + 1e-9):
                    curr_lr = max(min_lr, curr_lr * lr_decay)
                    print(f"[TART] plateau → lowering LR to {curr_lr:.2e}")
                    # rebuild optimizer with new LR, keep params
                    params = state.params
                    state = train_state.TrainState.create(
                        apply_fn=state.apply_fn,
                        params=params,
                        tx=optax.adamw(learning_rate=curr_lr, weight_decay=0.0),
                    )
                    epochs_no_gain = 0
                else:
                    print(f"[TART] early stop: lr at min and no improvement (best val {best_val:.6e})")
                    break

        # store best params
        self.state = state.replace(params=best_params)
        self.params = self.state.params

        # parameter count (for paper)
        self.param_count = sum(p.size for p in jax.tree_util.tree_leaves(self.params))
        print(f"[TART] params: {self.param_count:,}")

    def __call__(self, F_cX: np.ndarray, steps: int) -> np.ndarray:
        """
        Forecast `steps` points given any context slice F_cX that ENDS at
        the forecast start. Uses last L points (with left-padding if needed).
        """
        X = np.asarray(F_cX, dtype=np.float32)
        if X.ndim == 1:
            X = X[:, None]
        if X.shape[1] != self.D:
            raise ValueError(f"Expected D={self.D}, got {X.shape[1]}")

        # take last L rows (pad on the left if needed)
        if X.shape[0] < self.L:
            pad_len = self.L - X.shape[0]
            pad = np.repeat(X[:1], repeats=pad_len, axis=0)
            ctx = np.concatenate([pad, X], axis=0)
        else:
            ctx = X[-self.L:]

        # normalize with training stats
        ctx_n = (ctx - self._mu) / self._sd
        ctx_n = jnp.asarray(ctx_n[None, :, :])  # (1, L, D)

        preds_n = []
        params = self.params

        for _ in range(steps):
            # deterministic forward (no dropout)
            out = self.model.apply({"params": params}, ctx_n, deterministic=True)  # (1,L,D)
            next_n = out[:, -1, :]          # (1,D)
            preds_n.append(np.array(next_n[0]))

            # update context: drop oldest, append new
            ctx_n = jnp.concatenate([ctx_n[:, 1:, :], next_n[:, None, :]], axis=1)

        preds_n = np.stack(preds_n, axis=0)      # (steps, D) normalized
        preds = preds_n * self._sd + self._mu    # unnormalize
        return preds