Upload code/train.py with huggingface_hub

code/train.py

@@ -0,0 +1,177 @@
+"""
+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)