| """ |
| Campus Weather VAE |
| ================== |
| VAE for learning weather embeddings from 40 NUS campus weather stations. |
| Compresses 6 meteorological variables into a 12-dimensional latent space. |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from typing import Dict |
|
|
|
|
| class WeatherVAE(nn.Module): |
| """ |
| Variational Autoencoder for multivariate weather station data. |
| |
| Input: 6 weather variables (WindSpeed, WindDir, AirTemp, RelHum, AtmPress, GlobalRad) |
| Output: Compact latent embedding z ∈ R^d_latent |
| |
| The latent space is regularised (KL divergence) to be smooth and interpolable, |
| enabling spatial interpolation, forecasting, clustering, and anomaly detection. |
| """ |
| |
| def __init__(self, n_vars: int = 6, d_hidden: int = 128, d_latent: int = 12, |
| n_layers: int = 3, beta: float = 0.001): |
| super().__init__() |
| self.n_vars = n_vars |
| self.d_latent = d_latent |
| self.beta = beta |
| |
| |
| enc = [nn.Linear(n_vars, d_hidden), nn.LayerNorm(d_hidden), nn.GELU()] |
| for _ in range(n_layers - 1): |
| enc += [nn.Linear(d_hidden, d_hidden), nn.LayerNorm(d_hidden), nn.GELU()] |
| self.encoder = nn.Sequential(*enc) |
| self.mu_head = nn.Linear(d_hidden, d_latent) |
| self.logvar_head = nn.Linear(d_hidden, d_latent) |
| |
| |
| dec = [nn.Linear(d_latent, d_hidden), nn.GELU()] |
| for _ in range(n_layers - 1): |
| dec += [nn.Linear(d_hidden, d_hidden), nn.GELU()] |
| dec.append(nn.Linear(d_hidden, n_vars)) |
| self.decoder = nn.Sequential(*dec) |
| |
| |
| self.register_buffer('x_mean', torch.zeros(n_vars)) |
| self.register_buffer('x_std', torch.ones(n_vars)) |
| |
| def set_normalisation(self, mean: torch.Tensor, std: torch.Tensor): |
| self.x_mean.copy_(mean) |
| self.x_std.copy_(std) |
| |
| def normalise(self, x): |
| return (x - self.x_mean) / (self.x_std + 1e-8) |
| |
| def denormalise(self, x): |
| return x * (self.x_std + 1e-8) + self.x_mean |
| |
| def encode(self, x_raw: torch.Tensor) -> Dict[str, torch.Tensor]: |
| h = self.encoder(self.normalise(x_raw)) |
| mu = self.mu_head(h) |
| logvar = self.logvar_head(h) |
| if self.training: |
| z = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu) |
| else: |
| z = mu |
| return {'z': z, 'mu': mu, 'logvar': logvar} |
| |
| def decode(self, z: torch.Tensor) -> torch.Tensor: |
| return self.denormalise(self.decoder(z)) |
| |
| def forward(self, x_raw: torch.Tensor) -> Dict[str, torch.Tensor]: |
| enc = self.encode(x_raw) |
| x_hat = self.decode(enc['z']) |
| |
| loss_recon = F.mse_loss(self.normalise(x_hat), self.normalise(x_raw)) |
| loss_kl = -0.5 * torch.mean(1 + enc['logvar'] - enc['mu']**2 - enc['logvar'].exp()) |
| loss = loss_recon + self.beta * loss_kl |
| |
| return {'loss': loss, 'loss_recon': loss_recon.item(), 'loss_kl': loss_kl.item(), |
| 'x_hat': x_hat, 'z': enc['z'], 'mu': enc['mu']} |
| |
| def get_embedding(self, x_raw: torch.Tensor) -> torch.Tensor: |
| """Deterministic embedding (mu).""" |
| return self.encode(x_raw)['mu'] |
|
|
|
|
| def get_config(size='base'): |
| return { |
| 'small': dict(n_vars=6, d_hidden=64, d_latent=8, n_layers=2, beta=0.001), |
| 'base': dict(n_vars=6, d_hidden=128, d_latent=12, n_layers=3, beta=0.001), |
| 'large': dict(n_vars=6, d_hidden=256, d_latent=24, n_layers=4, beta=0.001), |
| }[size] |
|
|
