TARA.py

import math
from dataclasses import dataclass
from typing import Optional, Tuple

import numpy as np
from numpy.lib.stride_tricks import sliding_window_view

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


# -----------------------------
# utilities
# -----------------------------

def zscore(x: np.ndarray, eps: float = 1e-6) -> Tuple[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
    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 choose_heads(d_model: int, target_head_dim: int = 32, max_heads: int = 8, require_even_head_dim: bool = True) -> int:
    n_heads = max(1, min(max_heads, d_model // target_head_dim))
    while n_heads > 1:
        if d_model % n_heads != 0:
            n_heads -= 1
            continue
        head_dim = d_model // n_heads
        if require_even_head_dim and (head_dim % 2 != 0):
            n_heads -= 1
            continue
        break
    return n_heads


def make_nextstep_contexts(X: np.ndarray, ell: int) -> Tuple[np.ndarray, np.ndarray]:
    """
    1-step prediction from length-ell contexts:
      input  : ctx[i] = X[i : i+ell]    shape (ell, D)
      target : y[i]   = X[i+ell]        shape (D,)
    """
    X = np.asarray(X, dtype=np.float32)
    if X.ndim == 1:
        X = X[:, None]
    T, D = X.shape
    if T <= ell:
        raise ValueError(f"T={T} must be > ell={ell}")
    ctx = sliding_window_view(X, window_shape=ell, axis=0)[:-1]  # (N, ell, D)
    y = X[ell:]                                                  # (N, D)
    return np.array(ctx, dtype=np.float32), np.array(y, dtype=np.float32)


def build_rotary_inv_freq(head_dim: int, base: float = 10000.0) -> jnp.ndarray:
    if head_dim % 2 != 0:
        raise ValueError("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: jnp.ndarray, inv_freq: jnp.ndarray) -> jnp.ndarray:
    """
    x: (batch, heads, seq_len, head_dim)
    """
    b, h, t, d = x.shape
    half = d // 2
    pos = jnp.arange(t, dtype=jnp.float32)              # (t,)
    freqs = jnp.einsum("t,f->tf", pos, inv_freq)        # (t, half)
    cos = jnp.cos(freqs)[None, None, :, :]              # (1,1,t,half)
    sin = jnp.sin(freqs)[None, None, :, :]

    x1 = x[..., :half]
    x2 = x[..., half:]
    x1r = x1 * cos - x2 * sin
    x2r = x1 * sin + x2 * cos
    return jnp.concatenate([x1r, x2r], axis=-1)


# -----------------------------
# modules
# -----------------------------

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

    @nn.compact
    def __call__(self, x: jnp.ndarray, *, deterministic: bool) -> jnp.ndarray:
        hidden = int(self.d_model * self.mlp_ratio)
        x = nn.Dense(hidden)(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 MultiHeadCausalAttention(nn.Module):
    """
    If diffusion=False: standard causal attention logits = (QK^T)/sqrt(dh)
    If diffusion=True : "MHD" diffusion-map logits (DAC-style):
        D^2_ij = (g_j - G_ij) + (g_i - G_ji), with G_ij=<q_i,k_j>, g_j=G_jj
        P^+ = softmax_j(-beta D^2)  (row-stochastic diffusion kernel)
      Using row-softmax shift invariance, the -beta*g_i term can be dropped,
      giving logits ~ beta*(G_ij + G_ji) - beta*g_j.
    """
    d_model: int
    n_heads: int
    dropout: float = 0.0
    use_rope: bool = True

    diffusion: bool = False
    diffusion_beta: float = 1.0

    @nn.compact
    def __call__(self, x: jnp.ndarray, *, deterministic: bool) -> jnp.ndarray:
        b, t, d_model = x.shape
        if d_model != self.d_model:
            raise ValueError(f"Expected d_model={self.d_model}, got {d_model}")
        if self.d_model % self.n_heads != 0:
            raise ValueError("d_model must be divisible by n_heads")

        head_dim = self.d_model // self.n_heads
        if self.use_rope and (head_dim % 2 != 0):
            raise ValueError("head_dim must be even when use_rope=True")

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

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

        scale = 1.0 / math.sqrt(head_dim)

        # logits
        if not self.diffusion:
            # standard causal attention
            logits = jnp.einsum("bhid,bhjd->bhij", q, k) * scale          # (b,h,t,t)
        else:
            # DAC-style diffusion-map row kernel P^+ = softmax_j(-beta D^2_ij)
            # G_ij = <q_i, k_j> (scaled), G_ji = <k_i, q_j> (scaled)
            G   = jnp.einsum("bhid,bhjd->bhij", q, k) * scale             # (b,h,t,t)
            GT  = jnp.einsum("bhid,bhjd->bhij", k, q) * scale             # (b,h,t,t)
            # g_j = G_jj = <q_j, k_j> (scaled)
            g = jnp.einsum("bhid,bhid->bhi", q, k) * scale                # (b,h,t)
            g_key = g[:, :, None, :]                                      # broadcast over i

            beta = float(self.diffusion_beta)
            logits = beta * (G + GT - g_key)                              # (b,h,t,t)

        # causal mask (forecasting)
        mask = jnp.tril(jnp.ones((t, t), dtype=jnp.bool_))[None, None, :, :]
        logits = jnp.where(mask, logits, -1e9)

        attn = jax.nn.softmax(logits, axis=-1)
        attn = nn.Dropout(rate=self.dropout)(attn, deterministic=deterministic)

        out = jnp.einsum("bhij,bhjd->bhid", attn, v)                      # (b,h,t,hd)
        out = out.transpose(0, 2, 1, 3).reshape(b, t, self.d_model)       # (b,t,d)
        out = nn.Dense(self.d_model, use_bias=False, name="out")(out)
        out = nn.Dropout(rate=self.dropout)(out, deterministic=deterministic)
        return out


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

    use_ffn: bool = False
    mlp_ratio: float = 4.0

    diffusion: bool = False
    diffusion_beta: float = 1.0

    @nn.compact
    def __call__(self, x: jnp.ndarray, *, deterministic: bool) -> jnp.ndarray:
        h = nn.LayerNorm()(x)
        h = MultiHeadCausalAttention(
            d_model=self.d_model,
            n_heads=self.n_heads,
            dropout=self.dropout,
            use_rope=self.use_rope,
            diffusion=self.diffusion,
            diffusion_beta=self.diffusion_beta,
        )(h, deterministic=deterministic)
        x = x + h

        if self.use_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 TARAModel(nn.Module):
    """
    Hankel/Takens AR model.

    Input:  (batch, ell, D)
    Tokenize: Conv1D with kernel_size=L, VALID => tokens (batch, K, d_model) where K=ell-L+1
    Attention: causal over K tokens (window-start index T)
    Output: next-step prediction (batch, D) from last token
    """
    input_dim: int
    ell: int
    L: int
    d_model: int
    n_heads: int
    depth: int
    dropout: float = 0.0
    use_learned_pos: bool = True
    use_rope: bool = True

    use_ffn: bool = False
    mlp_ratio: float = 4.0

    diffusion: bool = False
    diffusion_beta: float = 1.0

    @nn.compact
    def __call__(self, x: jnp.ndarray, *, deterministic: bool = True) -> jnp.ndarray:
        b, t, d_in = x.shape
        if d_in != self.input_dim:
            raise ValueError(f"Expected input_dim={self.input_dim}, got {d_in}")
        if t != self.ell:
            raise ValueError(f"Expected ell={self.ell}, got t={t}")

        # Takens/Hankel tokenization without explicit Hankel tensor:
        # h_T = sum_{c,X} x_{T+c,X} * F_{c,X,M}
        h = nn.Conv(
            features=self.d_model,
            kernel_size=(self.L,),
            strides=(1,),
            padding="VALID",
            use_bias=True,
            name="token_conv",
        )(x)  # (b, K, d_model)
        K = h.shape[1]

        if self.use_learned_pos:
            pos_emb = self.param(
                "pos_emb",
                nn.initializers.normal(stddev=0.02),
                (K, self.d_model),
            )
            h = h + pos_emb[None, :, :]

        for i in range(self.depth):
            h = TARABlock(
                d_model=self.d_model,
                n_heads=self.n_heads,
                dropout=self.dropout,
                use_rope=self.use_rope,
                use_ffn=self.use_ffn,
                mlp_ratio=self.mlp_ratio,
                diffusion=self.diffusion,
                diffusion_beta=self.diffusion_beta,
                name=f"block_{i}",
            )(h, deterministic=deterministic)

        h = nn.LayerNorm(name="final_ln")(h)
        y = nn.Dense(self.input_dim, name="out_proj")(h[:, -1, :])  # (b,D)
        return y


# -----------------------------
# training wrapper
# -----------------------------

@dataclass
class TARAConfig:
    ell: int
    L: int
    d_model: int = 64
    depth: int = 4
    n_heads: Optional[int] = None

    use_ffn: bool = False
    mlp_ratio: float = 4.0

    diffusion: bool = False
    diffusion_beta: float = 1.0

    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
    patience: int = 5


class TARA:
    def __init__(self, R_tX: np.ndarray, cfg: TARAConfig):
        R = np.asarray(R_tX, dtype=np.float32)
        if R.ndim == 1:
            R = R[:, None]
        T, D = R.shape
        if T <= cfg.ell:
            raise ValueError(f"T={T} must be > ell={cfg.ell}")
        if cfg.ell < cfg.L:
            raise ValueError(f"Need ell >= L, got ell={cfg.ell}, L={cfg.L}")

        self.cfg = cfg
        self.D = int(D)

        # normalize
        Rn, (mu, sd) = zscore(R)
        self._mu, self._sd = mu, sd

        # dataset: contexts only (no explicit Hankel)
        X_ctx, Y = make_nextstep_contexts(Rn, ell=cfg.ell)  # (N,ell,D), (N,D)

        # Ensure X_ctx is (N, ell, D), not (N, D, ell)
        if X_ctx.ndim == 3 and X_ctx.shape[1] == self.D and X_ctx.shape[2] == cfg.ell:
            X_ctx = np.transpose(X_ctx, (0, 2, 1))

        
        N = X_ctx.shape[0]
        n_val = max(1, int(cfg.val_split * N))
        n_tr = N - n_val

        Xtr, Ytr = X_ctx[:n_tr], Y[:n_tr]
        Xva, Yva = X_ctx[n_tr:], Y[n_tr:]

        d_model = int(max(16, cfg.d_model))
        n_heads = cfg.n_heads
        if n_heads is None:
            n_heads = choose_heads(d_model, require_even_head_dim=cfg.use_rope)
        n_heads = int(n_heads)

        print(
            f"[TARA] N={N} (train={n_tr}, val={n_val}) | "
            f"ell={cfg.ell} L={cfg.L} => K={cfg.ell-cfg.L+1} | D={self.D} | "
            f"d_model={d_model} heads={n_heads} depth={cfg.depth} | "
            f"use_ffn={cfg.use_ffn} diffusion={cfg.diffusion}"
        )

        self.model = TARAModel(
            input_dim=self.D,
            ell=int(cfg.ell),
            L=int(cfg.L),
            d_model=d_model,
            n_heads=n_heads,
            depth=int(cfg.depth),
            dropout=float(cfg.dropout),
            use_learned_pos=bool(cfg.use_learned_pos),
            use_rope=bool(cfg.use_rope),
            use_ffn=bool(cfg.use_ffn),
            mlp_ratio=float(cfg.mlp_ratio),
            diffusion=bool(cfg.diffusion),
            diffusion_beta=float(cfg.diffusion_beta),
        )

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

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

        state = create_state(cfg.init_lr)

        # JAX arrays
        Xtr_j, Ytr_j = jnp.asarray(Xtr), jnp.asarray(Ytr)
        Xva_j, Yva_j = jnp.asarray(Xva), jnp.asarray(Yva)

        batch_size = int(cfg.batch_size)
        num_batches = lambda n: (n + batch_size - 1) // batch_size

        @jax.jit
        def train_step(st, xb, yb, rng_key):
            dropout_rng, rng_key = jax.random.split(rng_key)

            def loss_fn(p):
                pred = st.apply_fn(
                    {"params": p},
                    xb,
                    deterministic=False,
                    rngs={"dropout": dropout_rng},
                )
                return jnp.mean((pred - yb) ** 2)

            loss, grads = jax.value_and_grad(loss_fn)(st.params)
            st = st.apply_gradients(grads=grads)
            return st, loss, rng_key

        @jax.jit
        def eval_step(st, xb, yb):
            pred = st.apply_fn({"params": st.params}, xb, deterministic=True)
            return jnp.mean((pred - yb) ** 2)

        best_val = float("inf")
        best_params = state.params
        curr_lr = float(cfg.init_lr)
        epochs_no_gain = 0
        rng_np = np.random.default_rng(cfg.seed)
        idx = np.arange(n_tr, dtype=np.int32)

        for epoch in range(int(cfg.max_epochs)):
            rng_np.shuffle(idx)

            # train
            tr_losses = []
            for bi in range(num_batches(n_tr)):
                s = bi * batch_size
                e = min((bi + 1) * batch_size, n_tr)
                xb = Xtr_j[idx[s:e]]
                yb = Ytr_j[idx[s:e]]
                state, l, rng = train_step(state, xb, yb, rng)
                tr_losses.append(float(l))
            tr_loss = float(np.mean(tr_losses))

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

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

            if va_loss + 1e-10 < best_val:
                best_val = va_loss
                best_params = state.params
                epochs_no_gain = 0
            else:
                epochs_no_gain += 1

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

        self.state = state.replace(params=best_params)
        self.params = self.state.params
        self.param_count = sum(p.size for p in jax.tree_util.tree_leaves(self.params))
        print(f"[TARA] params: {self.param_count:,}")

    # -----------------------------
    # inference
    # -----------------------------

    def _prep_context(self, F_tX: np.ndarray) -> np.ndarray:
        X = np.asarray(F_tX, 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]}")

        ell = int(self.cfg.ell)
        if X.shape[0] < ell:
            pad_len = ell - X.shape[0]
            pad = np.repeat(X[:1], repeats=pad_len, axis=0)
            X = np.concatenate([pad, X], axis=0)
        else:
            X = X[-ell:, :]

        Xn = (X - self._mu) / self._sd
        return Xn.astype(np.float32)

    def predict_one(self, context: np.ndarray) -> np.ndarray:
        ctx = self._prep_context(context)  # (ell,D)
        yhat = self.model.apply({"params": self.params}, jnp.asarray(ctx[None, ...]), deterministic=True)
        yhat = np.asarray(yhat[0], dtype=np.float32)  # (D,) normalized
        return (yhat * self._sd.reshape(-1) + self._mu.reshape(-1)).astype(np.float32)

    def __call__(self, context: np.ndarray, *, steps: int = 1) -> np.ndarray:
        H = int(steps)
        if H <= 0:
            return np.zeros((0, self.D), dtype=np.float32)

        ctx = np.asarray(context, dtype=np.float32)
        if ctx.ndim == 1:
            ctx = ctx[:, None]

        out = np.zeros((H, self.D), dtype=np.float32)
        for h in range(H):
            y = self.predict_one(ctx)
            out[h] = y
            ctx = np.vstack([ctx, y[None, :]])
        return out