finalmodel.py

#!/usr/bin/env python3
"""
QHD-GRU: Quantum-Inspired Hilbert Delta-GRU
Google Cluster 2011 (task_usage)

✔ Streaming (RAM < 1GB)
✔ CUDA
✔ Train / Val / Test split
✔ Multi-horizon (5min → 1hour)
✔ TEST MSE, MAE, MAPE (paper-ready)

Author: Abhishek Singh
"""

# ================= IMPORTS =================
import os, glob
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import IterableDataset, DataLoader
from scipy.signal import hilbert

# ================= CONFIG =================
TASK_DIR = "/home/abhishek/google_cluster_2011/clusterdata-2011-2/task_usage"

MAX_FILES   = 8
CHUNK_SIZE = 100_000
LOOKBACK   = 12          # 12 × 5min = 1 hour
BATCH_SIZE = 64
EPOCHS     = 10
MAX_STEPS  = 200
DEVICE     = "cuda"

HORIZONS = {
    "5min": 1,
    "10min": 2,
    "15min": 3,
    "30min": 6,
    "1hour": 12
}

# ================= CUDA CHECK =================
assert torch.cuda.is_available()
print("Device:", torch.cuda.get_device_name(0))

FILES = sorted(glob.glob(os.path.join(TASK_DIR, "*.csv.gz")))[:MAX_FILES]
print("Using files:", len(FILES))

# =========================================================
# DATASET
# =========================================================
class GoogleHilbertDeltaDataset(IterableDataset):
    def __init__(self, horizon, mode="train"):
        self.h = horizon
        self.mode = mode

    def __iter__(self):
        for f in FILES:
            for chunk in pd.read_csv(
                f, header=None, compression="gzip", chunksize=CHUNK_SIZE
            ):
                chunk = chunk[[0, 4, 5]]
                chunk.columns = ["time", "machine", "cpu"]
                chunk.dropna(inplace=True)

                for _, g in chunk.groupby("machine"):
                    x = g.sort_values("time")["cpu"].values.astype(np.float32)
                    if len(x) < LOOKBACK + self.h + 10:
                        continue

                    # ---- Train / Val / Test split ----
                    n = len(x)
                    i1, i2 = int(0.6*n), int(0.8*n)
                    if self.mode == "train":
                        xs = x[:i1]
                    elif self.mode == "val":
                        xs = x[i1:i2]
                    else:
                        xs = x[i2:]

                    if len(xs) < LOOKBACK + self.h + 1:
                        continue

                    # EMA smoothing
                    xs = pd.Series(xs).ewm(span=10).mean().values

                    # Delta CPU
                    dx = np.diff(xs, prepend=xs[0])
                    dx = np.clip(dx, -0.2, 0.2)

                    # Hilbert transform
                    analytic = hilbert(dx)
                    amp = np.abs(analytic)
                    phase = np.unwrap(np.angle(analytic))

                    # Stabilization
                    amp = np.log1p(amp)
                    amp = (amp - amp.mean()) / (amp.std() + 1e-6)
                    amp = np.tanh(amp)
                    phase = np.tanh(phase / np.pi)

                    for i in range(len(dx) - LOOKBACK - self.h):
                        X = np.stack(
                            [amp[i:i+LOOKBACK], phase[i:i+LOOKBACK]],
                            axis=1
                        ).astype(np.float32)

                        y = dx[i+LOOKBACK:i+LOOKBACK+self.h].astype(np.float32)
                        last = np.float32(xs[i+LOOKBACK-1])

                        yield (
                            torch.from_numpy(X),
                            torch.from_numpy(y),
                            torch.tensor(last)
                        )

# =========================================================
# MODEL
# =========================================================
class QuantumInspiredLayer(nn.Module):
    def __init__(self, in_f, out_f):
        super().__init__()
        self.theta = nn.Parameter(torch.randn(in_f))
        self.fc = nn.Linear(in_f * 2, out_f)

    def forward(self, x):
        s = torch.sin(x * self.theta)
        c = torch.cos(x * self.theta)
        return self.fc(torch.cat([s, c], dim=-1))

class QHD_GRU(nn.Module):
    def __init__(self, horizon):
        super().__init__()
        self.q = QuantumInspiredLayer(2, 16)
        self.gru = nn.GRU(16, 48, batch_first=True)
        self.fc = nn.Linear(48, horizon)

    def forward(self, x):
        x = self.q(x)
        h, _ = self.gru(x)
        return self.fc(h[:, -1])

# =========================================================
# METRICS
# =========================================================
def compute_metrics(pred, true):
    mse  = torch.mean((pred - true) ** 2).item()
    mae  = torch.mean(torch.abs(pred - true)).item()
    mape = torch.mean(torch.abs((true - pred) / (true + 1e-6))).item() * 100
    return mse, mae, mape

# =========================================================
# TRAIN / VAL / TEST
# =========================================================
def run_horizon(h, label):
    print(f"\n===== {label} =====")

    train_dl = DataLoader(
        GoogleHilbertDeltaDataset(h, "train"),
        batch_size=BATCH_SIZE,
        num_workers=0
    )
    val_dl = DataLoader(
        GoogleHilbertDeltaDataset(h, "val"),
        batch_size=BATCH_SIZE,
        num_workers=0
    )
    test_dl = DataLoader(
        GoogleHilbertDeltaDataset(h, "test"),
        batch_size=BATCH_SIZE,
        num_workers=0
    )

    model = QHD_GRU(h).to(DEVICE)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)

    mse_loss = nn.MSELoss()
    huber = nn.HuberLoss(delta=0.05)

    # ---------------- TRAIN ----------------
    for ep in range(EPOCHS):
        model.train()
        steps = 0
        tr_mse = tr_mae = tr_mape = 0.0

        for x, y, last in train_dl:
            x, y, last = x.to(DEVICE), y.to(DEVICE), last.to(DEVICE)
            opt.zero_grad()

            dp = model(x)
            cpu_p = torch.clamp(
                torch.cumsum(dp, dim=1) + last.unsqueeze(1), 0.0, 1.0
            )
            cpu_t = torch.cumsum(y, dim=1) + last.unsqueeze(1)

            loss = mse_loss(cpu_p, cpu_t) + 0.4 * huber(dp, y)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()

            mse, mae, mape = compute_metrics(cpu_p, cpu_t)
            tr_mse += mse; tr_mae += mae; tr_mape += mape
            steps += 1
            if steps >= MAX_STEPS:
                break

        tr_mse /= steps; tr_mae /= steps; tr_mape /= steps

        # ---------------- VALIDATION ----------------
        model.eval()
        vsteps = 0
        val_mse = val_mae = val_mape = 0.0
        with torch.no_grad():
            for x, y, last in val_dl:
                x, y, last = x.to(DEVICE), y.to(DEVICE), last.to(DEVICE)
                dp = model(x)

                cpu_p = torch.clamp(
                    torch.cumsum(dp, dim=1) + last.unsqueeze(1), 0.0, 1.0
                )
                cpu_t = torch.cumsum(y, dim=1) + last.unsqueeze(1)

                mse, mae, mape = compute_metrics(cpu_p, cpu_t)
                val_mse += mse; val_mae += mae; val_mape += mape
                vsteps += 1
                if vsteps >= 100:
                    break

        val_mse /= vsteps; val_mae /= vsteps; val_mape /= vsteps

        print(
            f"Epoch {ep+1:02d} | "
            f"Train MSE={tr_mse:.6f} MAE={tr_mae:.4f} MAPE={tr_mape:.2f}% | "
            f"Val MSE={val_mse:.6f} MAE={val_mae:.4f} MAPE={val_mape:.2f}%"
        )

    # ---------------- TEST ----------------
    model.eval()
    tsteps = 0
    test_mse = test_mae = test_mape = 0.0
    with torch.no_grad():
        for x, y, last in test_dl:
            x, y, last = x.to(DEVICE), y.to(DEVICE), last.to(DEVICE)
            dp = model(x)

            cpu_p = torch.clamp(
                torch.cumsum(dp, dim=1) + last.unsqueeze(1), 0.0, 1.0
            )
            cpu_t = torch.cumsum(y, dim=1) + last.unsqueeze(1)

            mse, mae, mape = compute_metrics(cpu_p, cpu_t)
            test_mse += mse; test_mae += mae; test_mape += mape
            tsteps += 1
            if tsteps >= 200:
                break

    test_mse /= tsteps; test_mae /= tsteps; test_mape /= tsteps

    print(
        f"✅ FINAL TEST ({label}) | "
        f"MSE={test_mse:.6f} MAE={test_mae:.4f} MAPE={test_mape:.2f}%"
    )

    del model
    torch.cuda.empty_cache()

# =========================================================
# MAIN
# =========================================================
if __name__ == "__main__":
    for lbl, h in HORIZONS.items():
        run_horizon(h, lbl)