File size: 3,596 Bytes
a64bfbe c328a93 a64bfbe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | """
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
# Encoder
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)
# Decoder
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)
# Normalisation (set from data)
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]
|