code/train.py

"""
Data loading + training for Campus Weather VAE.
"""
import os, glob, time, json
import numpy as np
import pandas as pd
import torch
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from model import WeatherVAE, get_config

VAR_COLS = ['WindSpeed Ave (m/s)', 'WindDir Ave (degrees)', 'AirTemp Ave (C)',
            'RelHum Ave (%)', 'AtmPress Ave (hPa)', 'GlobalRad Ave (W/m2)']
VAR_NAMES = ['WindSpeed', 'WindDir', 'AirTemp', 'RelHum', 'AtmPress', 'GlobalRad']
VAR_UNITS = ['m/s', '°', '°C', '%', 'hPa', 'W/m²']


def load_nus40(imputed_dir: str, station_ids=None):
    """
    Load NUS-40 data. 
    station_ids: list of ints (1-40) to load, or None for all.
    Returns: data (T,N,V), coords (N,2), datetimes
    """
    files = sorted(glob.glob(os.path.join(imputed_dir, '*.csv')))
    if station_ids is not None:
        files = [f for f in files if int(f.split('WS')[1][:2]) in station_ids]
    
    dfs = [pd.read_csv(f) for f in files]
    N = len(dfs)
    T = len(dfs[0])
    V = len(VAR_COLS)
    
    data = np.full((T, N, V), np.nan, dtype=np.float32)
    coords = np.zeros((N, 2), dtype=np.float32)
    for i, df in enumerate(dfs):
        data[:, i, :] = df[VAR_COLS].values.astype(np.float32)
        coords[i] = [df['Latitude'].iloc[0], df['Longitude'].iloc[0]]
    
    # Fill remaining NaN with per-variable mean
    for v in range(V):
        col = data[:, :, v]
        m = np.nanmean(col)
        col[np.isnan(col)] = m
        data[:, :, v] = col
    
    datetimes = pd.to_datetime(dfs[0]['Datetime'])
    print(f"Loaded {N} stations, T={T}, V={V} | {datetimes.iloc[0]} to {datetimes.iloc[-1]}")
    return data, coords, datetimes


class FlatDataset(Dataset):
    """Flattens (T, N, V) → individual observation vectors for VAE training."""
    def __init__(self, data):
        # data: (T, N, V) → (T*N, V)
        T, N, V = data.shape
        self.X = torch.from_numpy(data.reshape(T * N, V))
    def __len__(self): return len(self.X)
    def __getitem__(self, i): return self.X[i]


class WindowDataset(Dataset):
    """Sliding windows for temporal tasks."""
    def __init__(self, data, window=24, stride=6):
        self.data = data
        self.window = window
        self.indices = list(range(0, data.shape[0] - window + 1, stride))
    def __len__(self): return len(self.indices)
    def __getitem__(self, i):
        s = self.indices[i]
        return torch.from_numpy(self.data[s:s + self.window])


def train_model(data_dir, config_size='base', epochs=100, batch_size=256,
                lr=5e-4, device='cpu', save_dir='checkpoints', station_ids=None):
    """Full training pipeline."""
    os.makedirs(save_dir, exist_ok=True)
    
    data, coords, datetimes = load_nus40(data_dir, station_ids)
    T, N, V = data.shape
    t_tr, t_va = int(T * 0.7), int(T * 0.85)
    
    tr_ds = FlatDataset(data[:t_tr])
    va_ds = FlatDataset(data[t_tr:t_va])
    tr_ld = DataLoader(tr_ds, batch_size=batch_size, shuffle=True, drop_last=True)
    va_ld = DataLoader(va_ds, batch_size=batch_size)
    
    cfg = get_config(config_size)
    model = WeatherVAE(**cfg).to(device)
    mean = torch.tensor(data[:t_tr].mean(axis=(0, 1)), dtype=torch.float32).to(device)
    std = torch.tensor(data[:t_tr].std(axis=(0, 1)), dtype=torch.float32).to(device)
    model.set_normalisation(mean, std)
    
    opt = optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    warmup = 5
    sched = optim.lr_scheduler.LambdaLR(opt, lambda ep: 
        min((ep + 1) / warmup, 0.5 * (1 + np.cos(np.pi * max(0, ep - warmup) / max(1, epochs - warmup)))))
    
    params = sum(p.numel() for p in model.parameters())
    print(f"Model: {params:,} params | d_latent={cfg['d_latent']} | device={device}")
    
    best_val = float('inf')
    history = []
    
    for ep in range(epochs):
        t0 = time.time()
        model.train()
        tr_r, tr_k, n = 0, 0, 0
        for batch in tr_ld:
            out = model(batch.to(device))
            out['loss'].backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step(); opt.zero_grad()
            tr_r += out['loss_recon']; tr_k += out['loss_kl']; n += 1
        
        model.eval()
        va_r, va_k, vn = 0, 0, 0
        with torch.no_grad():
            for batch in va_ld:
                out = model(batch.to(device))
                va_r += out['loss_recon']; va_k += out['loss_kl']; vn += 1
        
        sched.step()
        val_loss = va_r / vn + cfg['beta'] * va_k / vn
        history.append({'train_recon': tr_r/n, 'train_kl': tr_k/n, 'val_recon': va_r/vn, 'val_kl': va_k/vn})
        
        if val_loss < best_val:
            best_val = val_loss
            torch.save({'model': model.state_dict(), 'config': cfg,
                        'mean': mean.cpu(), 'std': std.cpu(),
                        'coords': coords, 'station_ids': station_ids,
                        'history': history, 'epoch': ep}, f'{save_dir}/best.pt')
        
        if (ep + 1) % 10 == 0 or ep == 0:
            print(f"Ep {ep+1:3d}/{epochs} ({time.time()-t0:.1f}s) | "
                  f"Train R={tr_r/n:.4f} KL={tr_k/n:.3f} | Val R={va_r/vn:.4f} KL={va_k/vn:.3f}")
    
    print(f"Done. Best val loss: {best_val:.4f}")
    
    # Extract and save embeddings for entire dataset
    model.eval()
    all_emb = []
    with torch.no_grad():
        for t in range(0, T, 1000):
            chunk = torch.from_numpy(data[t:t+1000].reshape(-1, V)).to(device)
            emb = model.get_embedding(chunk).cpu().numpy()
            all_emb.append(emb)
    embeddings = np.concatenate(all_emb).reshape(T, N, -1)
    np.savez_compressed(f'{save_dir}/embeddings.npz', 
                        embeddings=embeddings, data=data, coords=coords,
                        datetimes=datetimes.values.astype(str))
    print(f"Embeddings saved: {embeddings.shape}")
    
    return model, data, coords, datetimes, embeddings


def load_trained(save_dir, device='cpu'):
    """Load trained model and embeddings."""
    ckpt = torch.load(f'{save_dir}/best.pt', map_location=device, weights_only=False)
    model = WeatherVAE(**ckpt['config']).to(device)
    model.load_state_dict(ckpt['model'])
    model.set_normalisation(ckpt['mean'].to(device), ckpt['std'].to(device))
    model.eval()
    
    npz = np.load(f'{save_dir}/embeddings.npz', allow_pickle=True)
    return model, npz['data'], npz['coords'], npz['embeddings'], ckpt


if __name__ == '__main__':
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument('--data', default='/app/data_tmp/imputed')
    p.add_argument('--epochs', type=int, default=100)
    p.add_argument('--config', default='base')
    p.add_argument('--device', default='cpu')
    p.add_argument('--save', default='/app/campus_weather/results/checkpoints')
    args = p.parse_args()
    train_model(args.data, args.config, args.epochs, device=args.device, save_dir=args.save)