""" 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)