| """Joint two-target network for a smooth 'prevailing' pollution level per station and day. |
| |
| The model of the paper "PM2.5, PM10 ~ Station + Time, Analysis" (www.alphaxiv.org/abs/2610.pm25-pm10-joint-smoothing): |
| - one row per station x day with up to two targets (log PM2.5, log PM10); a missing target is masked, not dropped; |
| - location = area embedding + small station offset (station falls back to its area when data is thin); |
| - time enters only through fixed Gaussian bumps, so the level is a smooth function of time; |
| - loss = Gaussian NLL of an observed daily mean of n hours, variance sigma^2 / n + tau^2; |
| - a roughness penalty on the second difference of every station's predicted series over the full day grid. |
| """ |
| from dataclasses import dataclass |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| from torch import nn |
|
|
| TARGETS = ["PM2.5", "PM10"] |
| KEY = ["station", "date"] |
|
|
|
|
| @dataclass |
| class Encoding: |
| stations: list |
| areas: list |
| station_area: dict |
| start: pd.Timestamp |
| n_days: int |
| y_mean: np.ndarray |
| y_std: np.ndarray |
|
|
| @classmethod |
| def fit(cls, df: pd.DataFrame, end: pd.Timestamp) -> "Encoding": |
| start = df["date"].min() |
| logv = np.log(df["value"]) |
| return cls(stations=sorted(df["station"].unique()), areas=sorted(df["area"].unique()), |
| station_area=df.drop_duplicates("station").set_index("station")["area"].to_dict(), |
| start=start, n_days=(end - start).days + 1, |
| y_mean=np.array([logv[df["pollutant"] == t].mean() for t in TARGETS]), |
| y_std=np.array([logv[df["pollutant"] == t].std() for t in TARGETS])) |
|
|
| def wide(self, df: pd.DataFrame) -> pd.DataFrame: |
| w = df.pivot_table(index=KEY, columns="pollutant", values=["value", "n"], aggfunc="first") |
| w.columns = [f"{v}_{p}" for v, p in w.columns] |
| for t in TARGETS: |
| for v in ["value", "n"]: |
| if f"{v}_{t}" not in w: |
| w[f"{v}_{t}"] = np.nan |
| return w.reset_index() |
|
|
| def inputs(self, keys: pd.DataFrame) -> dict: |
| s_idx = {s: i for i, s in enumerate(self.stations)} |
| a_idx = {a: i for i, a in enumerate(self.areas)} |
| return {"station": torch.tensor(keys["station"].map(s_idx).to_numpy()), |
| "area": torch.tensor(keys["station"].map(self.station_area).map(a_idx).to_numpy()), |
| "day": torch.tensor((keys["date"] - self.start).dt.days.to_numpy())} |
|
|
| def targets(self, w: pd.DataFrame, use=TARGETS): |
| val = w[[f"value_{t}" for t in TARGETS]].to_numpy(float) |
| n = w[[f"n_{t}" for t in TARGETS]].to_numpy(float) |
| mask = ~np.isnan(val) |
| mask[:, [t not in use for t in TARGETS]] = False |
| y = (np.log(np.where(mask, val, 1.0)) - self.y_mean) / self.y_std |
| return (torch.tensor(np.where(mask, y, 0.0), dtype=torch.float32), torch.tensor(mask), |
| torch.tensor(np.where(mask, n, 1.0), dtype=torch.float32)) |
|
|
| def to_value(self, mu: torch.Tensor) -> np.ndarray: |
| return np.exp(mu.detach().numpy() * self.y_std + self.y_mean) |
|
|
| def grid(self) -> pd.DataFrame: |
| days = pd.date_range(self.start, periods=self.n_days) |
| return pd.DataFrame({"station": self.stations}).merge(pd.DataFrame({"date": days}), how="cross") |
|
|
|
|
| class TimeBasis(nn.Module): |
| def __init__(self, n_days: int, spacing: int): |
| super().__init__() |
| self.register_buffer("centres", torch.arange(-spacing, n_days + 2 * spacing, spacing, dtype=torch.float32)) |
| self.width = float(spacing) |
|
|
| def forward(self, day): |
| return torch.exp(-0.5 * ((day.float()[:, None] - self.centres[None, :]) / self.width) ** 2) |
|
|
|
|
| class PrevailingNet(nn.Module): |
| def __init__(self, enc: Encoding, dim_loc: int = 8, hidden: int = 64, spacing: int = 30): |
| super().__init__() |
| self.time = TimeBasis(enc.n_days, spacing) |
| self.area = nn.Embedding(len(enc.areas), dim_loc) |
| nn.init.normal_(self.area.weight, std=0.1) |
| self.station = nn.Embedding(len(enc.stations), dim_loc) |
| nn.init.normal_(self.station.weight, std=0.01) |
| self.trunk = nn.Sequential(nn.Linear(dim_loc + len(self.time.centres), hidden), nn.SiLU(), |
| nn.Linear(hidden, hidden), nn.SiLU()) |
| self.heads = nn.ModuleList([nn.Linear(hidden, 1) for _ in TARGETS]) |
| self.log_sigma = nn.Parameter(torch.zeros(len(TARGETS))) |
| self.log_tau = nn.Parameter(torch.full((len(TARGETS),), -1.0)) |
| self.register_buffer("area_of", torch.tensor([enc.areas.index(enc.station_area[s]) for s in enc.stations])) |
| self.n_days = enc.n_days |
|
|
| def forward(self, x): |
| loc = self.area(x["area"]) + self.station(x["station"]) |
| h = self.trunk(torch.cat([loc, self.time(x["day"])], dim=1)) |
| return torch.cat([head(h) for head in self.heads], dim=1) |
|
|
| def all_series(self): |
| """Predictions for every station over the full day grid, shape (stations, days, targets).""" |
| s = torch.arange(len(self.area_of)) |
| x = {"station": s.repeat_interleave(self.n_days), "area": self.area_of[s].repeat_interleave(self.n_days), |
| "day": torch.arange(self.n_days).repeat(len(s))} |
| return self(x).view(len(s), self.n_days, len(TARGETS)) |
|
|
|
|
| def nll(mu, model, y, mask, n): |
| var = torch.exp(2 * model.log_sigma) / n + torch.exp(2 * model.log_tau) |
| out = 0.5 * (torch.log(var) + (y - mu) ** 2 / var) |
| return (out * mask).sum() / mask.sum() |
|
|
|
|
| def roughness(model): |
| mu = model.all_series() |
| return (mu[:, 2:] - 2 * mu[:, 1:-1] + mu[:, :-2]).pow(2).mean() |
|
|
|
|
| def fit(enc, w, cfg, steps, use=TARGETS, on_checkpoint=None, checkpoints=()): |
| """Full-batch Adam on a cosine schedule spanning `cfg['schedule']` steps (same meaning in CV and final fit).""" |
| torch.manual_seed(cfg.get("seed", 0)) |
| model = PrevailingNet(enc, hidden=cfg["hidden"], spacing=cfg["spacing"]) |
| opt = torch.optim.Adam(model.parameters(), lr=cfg["lr"], weight_decay=cfg["weight_decay"]) |
| sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, cfg["schedule"]) |
| x, (y, mask, n) = enc.inputs(w), enc.targets(w, use) |
| for step in range(1, steps + 1): |
| model.train() |
| loss = (nll(model(x), model, y, mask, n) + cfg["rough"] * roughness(model) |
| + cfg["station_l2"] * model.station.weight.pow(2).mean()) |
| opt.zero_grad() |
| loss.backward() |
| opt.step() |
| sched.step() |
| if on_checkpoint and step in checkpoints: |
| model.eval() |
| on_checkpoint(step, model) |
| model.eval() |
| return model |
|
|
|
|
| def predict(model, enc, keys: pd.DataFrame) -> pd.DataFrame: |
| with torch.no_grad(): |
| v = enc.to_value(model(enc.inputs(keys))) |
| return pd.concat([keys[KEY].assign(pollutant=t, pred=v[:, i]) for i, t in enumerate(TARGETS)], ignore_index=True) |
|
|
