RNN, LSTM and time series

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader
from sklearn.linear_model import LinearRegression
from sklearn.neighbors import KNeighborsRegressor
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error
from sklearn.model_selection import train_test_split

torch.set_num_threads(4)

r = np.random.default_rng(7)
MONTHS = 400
month = np.arange(MONTHS)
series_values = (420 + 0.9 * month + 38 * np.sin(2 * np.pi * month / 12) + r.normal(0, 14, MONTHS)).round(2)
series = pd.Series(series_values, index=pd.date_range("1993-01-01", periods=MONTHS, freq="MS"))

print(f"Series length : {MONTHS} months")
print(f"Range         : {series_values.min():.1f} .. {series_values.max():.1f}")
print(f"First year    : {series_values[:12].mean():.1f}")
print(f"Last year     : {series_values[-12:].mean():.1f}")
print()
print("It contains three things: a TREND (growth), SEASONALITY (a 12-month wave), NOISE.")

figure, (a1, a2) = plt.subplots(2, 1, figsize=(12, 6.5))
a1.plot(series.index, series.values, lw=1.1)
a1.set_title("The full series: trend + seasonality")
a1.set_ylabel("price")
a2.plot(series.index[-48:], series.values[-48:], marker="o", ms=3.5, lw=1.2)
a2.set_title("The last 48 months: the seasonality up close")
a2.set_ylabel("price")
plt.tight_layout()
fig("qator.png", caption="A time series: a 12-month seasonal wave on top of a growth trend.")

L = 12
X = np.stack([series_values[i:i + L] for i in range(MONTHS - L)])
y = series_values[L:]

print(f"Window length L = {L}")
print(f"X shape: {X.shape}   y shape: {y.shape}")
print()
print("An important observation: neighbouring windows are very similar:")
print(f"  window 0: {X[0][:5]} ...")
print(f"  window 1: {X[1][:5]} ...")
print(f"  values in common: {len(set(X[0]) & set(X[1]))} / {L}")

cut = int(len(X) * 0.8)
Xtr, Xte, ytr, yte = X[:cut], X[cut:], y[:cut], y[cut:]
Xa, Xb, ya, yb = train_test_split(X, y, test_size=0.2, random_state=42)

print(f"Chronological: train {len(Xtr)} months -> test {len(Xte)} months (the following period)")
print(f"Random       : train {len(Xa)} months -> test {len(Xb)} months (mixed across time)")
print()
print(f"Train target range (chronological): {ytr.min():.1f} .. {ytr.max():.1f}")
print(f"Test  target range (chronological): {yte.min():.1f} .. {yte.max():.1f}")
print()
print(f"{float((yte > ytr.max()).mean()):.1%} of the test set is ABOVE the LARGEST value in train.")
print("So the model has to EXTRAPOLATE: go into a range it has never seen.")

models = [
    ("linear regression", lambda: LinearRegression()),
    ("KNN (k=3)", lambda: KNeighborsRegressor(3)),
    ("random forest", lambda: RandomForestRegressor(200, random_state=0)),
]
table = []
for name, make in models:
    chrono = mean_absolute_error(yte, make().fit(Xtr, ytr).predict(Xte))
    rand = mean_absolute_error(yb, make().fit(Xa, ya).predict(Xb))
    table.append({"model": name, "chronological MAE": round(chrono, 2),
                  "random MAE": round(rand, 2), "gap": round(rand - chrono, 2)})
show(pd.DataFrame(table))

m_knn = KNeighborsRegressor(3).fit(Xtr, ytr)
m_lin = LinearRegression().fit(Xtr, ytr)
figure, ax = plt.subplots(figsize=(12, 4.2))
ind = np.arange(len(yte))
ax.plot(ind, yte, "k-", lw=2, label="actual")
ax.plot(ind, m_knn.predict(Xte), "--", lw=1.6, label="KNN prediction")
ax.plot(ind, m_lin.predict(Xte), "--", lw=1.6, label="linear regression prediction")
ax.axhline(ytr.max(), color="red", ls=":", lw=1.4, label=f"the maximum in train ({ytr.max():.0f})")
ax.set_title("KNN CANNOT go above the red line")
ax.set_xlabel("test month"); ax.set_ylabel("price"); ax.legend(fontsize=9)
fig("ekstrapolyatsiya.png", caption="The KNN prediction hits the ceiling of the training maximum. Linear regression continues the trend.")

baselines = {
    "naive (the last value)": Xte[:, -1],
    "seasonal naive (12 months earlier)": Xte[:, 0],
}
result = [{"method": k, "MAE": round(mean_absolute_error(yte, v), 2)} for k, v in baselines.items()]
result.append({"method": "linear regression (12 lags)",
               "MAE": round(mean_absolute_error(yte, m_lin.predict(Xte)), 2)})
show(pd.DataFrame(result))
print("If a complex model doesn't give a lower MAE than THESE numbers, it's unnecessary.")

diagramma(
    {
        "x1\nmonth 1": (0.2, 2.5, 1.4, 0.7, "manba"),
        "x2\nmonth 2": (2.1, 2.5, 1.4, 0.7, "manba"),
        "x3\nmonth 3": (4.0, 2.5, 1.4, 0.7, "manba"),
        "h1": (0.2, 1.3, 1.4, 0.7, "qayta"),
        "h2": (2.1, 1.3, 1.4, 0.7, "qayta"),
        "h3": (4.0, 1.3, 1.4, 0.7, "qayta"),
        "answer": (4.0, 0.1, 1.4, 0.7, "xizmat"),
    },
    [("x1\nmonth 1", "h1"), ("x2\nmonth 2", "h2"), ("x3\nmonth 3", "h3"),
     ("h1", "h2", "state"), ("h2", "h3", "state"), ("h3", "answer")],
    nom="rnn_sxema.png",
    izoh="An RNN walks step by step. The hidden state is passed to the right, "
         "and the last state gives the answer. The first month's influence has to travel down this chain.",
)

h = 32
rnn = nn.RNN(1, h, batch_first=True)
lstm = nn.LSTM(1, h, batch_first=True)
print(f"nn.RNN(1, {h})  parameters: {sum(p.numel() for p in rnn.parameters())}")
print(f"nn.LSTM(1, {h}) parameters: {sum(p.numel() for p in lstm.parameters())}  (4 times as many: 4 gates)")
print()
print("The LSTM's 4 gates: forget, input, output, and the cell update.")
print("The theory: the gates learn 'what to remember'.")

def memory_task(n, T, seed=0):
    rr = np.random.default_rng(seed)
    Xs = rr.normal(0, 1, (n, T, 1)).astype(np.float32)
    return torch.tensor(Xs), torch.tensor((Xs[:, 0, 0] > 0).astype(np.int64))

class SeqModel(nn.Module):
    def __init__(self, kind, h=32):
        super().__init__()
        self.rnn = nn.RNN(1, h, batch_first=True) if kind == "rnn" else nn.LSTM(1, h, batch_first=True)
        self.fc = nn.Linear(h, 2)
    def forward(self, x):
        out, _ = self.rnn(x)
        return self.fc(out[:, -1])

def seq_train(kind, T, seed, epochs=60, n=2000):
    torch.manual_seed(seed)
    Xtr_, ytr_ = memory_task(n, T, 1)
    Xte_, yte_ = memory_task(600, T, 2)
    m = SeqModel(kind)
    opt = torch.optim.Adam(m.parameters(), lr=0.005)
    crit = nn.CrossEntropyLoss()
    dl = DataLoader(TensorDataset(Xtr_, ytr_), batch_size=64, shuffle=True,
                    generator=torch.Generator().manual_seed(seed))
    for _ in range(epochs):
        m.train()
        for xb, yb_ in dl:
            opt.zero_grad(); crit(m(xb), yb_).backward(); opt.step()
    m.eval()
    with torch.no_grad():
        return float((m(Xte_).argmax(1) == yte_).float().mean())

print("A random answer = 0.50 accuracy. Each configuration with 3 SEEDS.")
rows = []
for T in [5, 20, 60]:
    for kind in ["rnn", "lstm"]:
        acc = [seq_train(kind, T, s) for s in (0, 1, 2)]
        rows.append({"length T": T, "model": kind,
                     "seeds": str([round(a, 3) for a in acc]),
                     "mean": round(float(np.mean(acc)), 4),
                     "worst": round(min(acc), 4)})
show(pd.DataFrame(rows))

print("The gradient's size at each input step (T=60, before training):")
print()
for kind in ["rnn", "lstm"]:
    torch.manual_seed(0)
    Xg, yg = memory_task(256, 60, 1)
    Xg.requires_grad_(True)
    m = SeqModel(kind)
    nn.CrossEntropyLoss()(m(Xg), yg).backward()
    g = Xg.grad.abs().mean(dim=(0, 2))
    print(f"  {kind.upper():>5}: step 0 -> {g[0]:.3e} | step 30 -> {g[30]:.3e} | "
          f"step 59 -> {g[59]:.3e}")
    print(f"         the ratio step 59 / step 0: {g[59] / g[0]:.2e}")

def lstm_forecast(seed, epochs, difference=False):
    torch.manual_seed(seed)
    if difference:
        Xa, Xb = np.diff(Xtr, axis=1), np.diff(Xte, axis=1)
        ya, base = ytr - Xtr[:, -1], Xte[:, -1]
    else:
        Xa, Xb, ya, base = Xtr, Xte, ytr, 0.0
    mu_, sd_ = Xa.mean(), Xa.std()
    ymu, ysd = ya.mean(), ya.std()
    core = nn.LSTM(1, 32, batch_first=True)
    head = nn.Linear(32, 1)
    opt = torch.optim.Adam(nn.ModuleList([core, head]).parameters(), lr=0.01)
    xt = torch.tensor(((Xa - mu_) / sd_).astype(np.float32)).unsqueeze(-1)
    yt = torch.tensor(((ya - ymu) / ysd).astype(np.float32)).unsqueeze(-1)
    dl = DataLoader(TensorDataset(xt, yt), batch_size=32, shuffle=True,
                    generator=torch.Generator().manual_seed(seed))
    for _ in range(epochs):
        for xb, yb_ in dl:
            opt.zero_grad()
            o, _ = core(xb)
            nn.MSELoss()(head(o[:, -1]), yb_).backward()
            opt.step()
    with torch.no_grad():
        o, _ = core(torch.tensor(((Xb - mu_) / sd_).astype(np.float32)).unsqueeze(-1))
        return head(o[:, -1]).squeeze(-1).numpy() * ysd + ymu + base

experiments = [
    ("LSTM, raw series, 120 epochs", dict(epochs=120)),
    ("LSTM, raw series, 400 epochs", dict(epochs=400)),
    ("LSTM, DIFFERENCED, 400 epochs", dict(epochs=400, difference=True)),
]
final = [
    {"method": "naive (the last value)", "MAE": round(mean_absolute_error(yte, Xte[:, -1]), 2),
     "seeds": "-"},
    {"method": "linear regression (12 lags)",
     "MAE": round(mean_absolute_error(yte, m_lin.predict(Xte)), 2), "seeds": "-"},
]
for name, kw in experiments:
    xs_ = [mean_absolute_error(yte, lstm_forecast(s, **kw)) for s in (0, 1, 2)]
    final.append({"method": name, "MAE": round(float(np.mean(xs_)), 2),
                  "seeds": str([round(v, 1) for v in xs_])})
show(pd.DataFrame(final))

from sklearn.model_selection import TimeSeriesSplit

tss = TimeSeriesSplit(n_splits=5)
print("TimeSeriesSplit: the train window grows, and the test always comes after it:")
print()
print(f"{'fold':>6} {'train':>14} {'test':>14}")
print("-" * 38)
for i, (tr, te) in enumerate(tss.split(X), 1):
    print(f"{i:>6} {f'0..{tr[-1]}':>14} {f'{te[0]}..{te[-1]}':>14}")
print()
xs = [mean_absolute_error(y[te], LinearRegression().fit(X[tr], y[tr]).predict(X[te]))
      for tr, te in tss.split(X)]
print(f"MAE in each fold : {[round(v, 2) for v in xs]}")
print(f"Mean             : {np.mean(xs):.2f}  (+/- {np.std(xs):.2f})")

show(pd.DataFrame([
    {"Concept": "The lag window", "In one sentence": "The way to turn a time series into a table"},
    {"Concept": "The chronological split", "In one sentence": "Mandatory: a random split makes the result look good"},
    {"Concept": "Extrapolation", "In one sentence": "Trees and KNN can't go above the training maximum"},
    {"Concept": "Differencing", "In one sentence": "Removing the trend: done BEFORE choosing a model"},
    {"Concept": "TimeSeriesSplit", "In one sentence": "Replaces KFold for a time series"},
    {"Concept": "The time-series baseline", "In one sentence": "naive and seasonal naive: the first rivals"},
    {"Concept": "RNN", "In one sentence": "Walks step by step, keeps a hidden state"},
    {"Concept": "LSTM gates", "In one sentence": "4 times the parameters; long memory in theory"},
    {"Concept": "Seed variance", "In one sentence": "Measured: one seed isn't enough, record the worst"},
    {"Concept": "The vanishing gradient", "In one sentence": "Measured: almost no signal reaches the first step"},
]))