Model code and training recipe

README.md

@@ -0,0 +1,120 @@
+---
+license: cc-by-4.0
+library_name: pytorch
+tags:
+  - time-series
+  - panel-data
+  - missing-data
+  - smoothing
+  - imputation
+  - air-quality
+  - pytorch
+datasets:
+  - prteek/sparse-panel-demo
+---
+
+# Prevailing level network for sparse, gappy panels
+
+A small PyTorch network that estimates a smooth **prevailing level** for every location and day of a panel,
+for two related quantities at once, when the observations are noisy daily means with long gaps and some series
+missing entirely. It was developed for the post
+[PM2.5, PM10 ~ station + time](https://prteek.github.io/regression_room/posts/beijing_air_quality/) and powers the
+[sparse panel demo Space](https://huggingface.co/spaces/prteek/sparse-panel-demo).
+
+**This repo holds the code and the training recipe, not trained weights.** The network learns an embedding for
+each location and a time basis over a fixed date range, so a fitted model only applies to the panel it was fitted
+on. To use it, fit it to your own panel: on the 12-station, four-year Beijing data a fit takes a few seconds on a
+CPU.
+
+## Files
+
+- `model.py`: the encoding, the network, the loss and `fit` / `predict`
+- `example.py`: fits the network to a small synthetic panel and predicts every cell (runs as is)
+
+## Model
+
+Each row is one location and day, with up to two targets $y_1, y_2$ (log values, standardised).
+
+$$
+\begin{aligned}
+x &= \big[\, e_{\text{area}} + e_{\text{station}},\ b(t) \,\big] \\
+h &= \operatorname{SiLU}\!\big(W_2 \operatorname{SiLU}(W_1 x)\big) \\
+\mu_1 &= w_1^\top h, \quad \mu_2 = w_2^\top h
+\end{aligned}
+$$
+
+- **Location**: an 8-dim area embedding plus a station offset that starts near zero and is penalised, so a
+  station with little data falls back to its area.
+- **Time**: fixed Gaussian bumps $b_k(t) = \exp\{-\tfrac12((t - c_k)/w)^2\}$ with centres every $w$ days, so the
+  level is a smooth function of time.
+- **Shared trunk, two heads**: two layers of 64 with SiLU, then one linear head per target. A day where only one
+  target was observed still trains everything the other depends on.
+
+**Loss.** An observed value is a mean of $n$ readings, so its variance around the level is $\sigma_j^2 / n + \tau_j^2$,
+with $\sigma_j$, $\tau_j$ learned per target. The Gaussian negative log-likelihood is averaged over observed cells
+only: a missing target is **masked**, not dropped. Added to it:
+
+- a **curvature penalty** $\lambda \cdot \overline{(\mu_{t+1} - 2\mu_t + \mu_{t-1})^2}$ over every location's
+  predicted series on the **whole** day grid, observed or not, which keeps the level smooth through gaps;
+- a small L2 penalty on the station offsets.
+
+## Training recipe
+
+Full-batch Adam with a cosine learning-rate schedule. The settings below were chosen by look-ahead
+cross-validation in the post:
+
+| Setting | Value |
+|---|---|
+| curvature penalty $\lambda$ (`rough`) | 1e5 |
+| bump spacing $w$ (`spacing`) | 30 days |
+| hidden width | 64 |
+| learning rate | 5e-3 |
+| weight decay | 1e-4 |
+| station offset L2 (`station_l2`) | 1e-3 |
+| schedule length | 1000 |
+| training steps | 100 |
+| seed | 0 |
+
+For a single-target model, pass `use=["PM2.5"]` (or the other target) to `fit`.
+
+## Input format
+
+Long format, one row per location × date × target that has data:
+
+| column | meaning |
+|---|---|
+| `station` | location id |
+| `area` | group the location belongs to (it falls back to this when its own data is thin) |
+| `date` | daily timestamp |
+| `pollutant` | target name, one of `model.TARGETS` |
+| `value` | the daily mean (> 0; modelled on the log scale) |
+| `n` | number of readings behind that mean (weights the observation) |
+
+The two target names are set in `model.TARGETS` (`["PM2.5", "PM10"]`); change them for other quantities.
+
+```python
+from model import Encoding, fit, predict
+
+enc = Encoding.fit(panel, end=panel["date"].max())
+model = fit(enc, enc.wide(panel), CFG, steps=100)
+level = predict(model, enc, enc.grid())   # every location x day, observed or not
+```
+
+## Limitations
+
+- **Not transferable.** Location embeddings and the time basis are specific to one panel and date range. Refit
+  for new locations or dates.
+- **Smoothness is a choice.** Cross-validation that scores against noisy readings rewards following short
+  episodes, so it tends to under-smooth; the post discusses this. Treat `rough` and `spacing` as decisions about
+  how smooth the level should be.
+- **Log-scale bias.** The model is fitted to logs of means over few readings, so its level sits slightly below a
+  moving average of full data (Jensen's inequality).
+- **A GAM is a strong alternative.** In the post, a mixed-effects GAM with hand-set smoothness does as well or
+  better.
+- **Memory.** Training is full batch and the penalty covers every location × day, so memory grows with panel size.
+
+## Data used in the post
+
+Beijing Multi-Site Air Quality, Song Chen (2017), UCI Machine Learning Repository,
+[doi:10.24432/C5RK5G](https://doi.org/10.24432/C5RK5G), CC BY 4.0. The demo's sparse readings and predictions
+are in the dataset [prteek/sparse-panel-demo](https://huggingface.co/datasets/prteek/sparse-panel-demo).

example.py

@@ -0,0 +1,39 @@
+"""Fit the prevailing-level network to a panel of your own and predict a smooth level for every cell.
+
+Input is long format, one row per station x date x pollutant that has data:
+  station, area, date, pollutant (one of model.TARGETS), value (> 0, a daily mean), n (readings behind it)
+Here a small synthetic panel stands in for real data.
+"""
+import numpy as np
+import pandas as pd
+
+from model import TARGETS, Encoding, fit, predict
+
+# Settings chosen by look-ahead cross-validation in the post
+CFG = dict(rough=1e5, spacing=30, lr=5e-3, hidden=64, weight_decay=1e-4, station_l2=1e-3, schedule=1000, seed=0)
+STEPS = 100
+
+
+def synthetic_panel(seed: int = 0) -> pd.DataFrame:
+    rng = np.random.default_rng(seed)
+    days = pd.date_range("2020-01-01", periods=365)
+    stations = {"s1": "urban", "s2": "urban", "s3": "suburban", "s4": "suburban"}
+    t = np.arange(len(days))
+    rows = []
+    for s, area in stations.items():
+        level = 4 + 0.5 * np.sin(2 * np.pi * t / 365) + (0.3 if area == "urban" else 0) + rng.normal(0, 0.1)
+        for j, pollutant in enumerate(TARGETS):
+            n = rng.poisson(1.2, len(days))  # hours observed per day; 0 means no reading
+            keep = n > 0
+            value = np.exp(level + 0.4 * j + rng.normal(0, 0.5 / np.sqrt(np.maximum(n, 1))))
+            rows.append(pd.DataFrame({"station": s, "area": area, "date": days[keep], "pollutant": pollutant,
+                                      "value": value[keep], "n": n[keep]}))
+    return pd.concat(rows, ignore_index=True)
+
+
+panel = synthetic_panel()
+enc = Encoding.fit(panel, end=panel["date"].max())
+model = fit(enc, enc.wide(panel), CFG, STEPS)            # joint model; use=["PM2.5"] for a single-target one
+level = predict(model, enc, enc.grid())                  # every station x day, observed or not
+print(level.head())
+print(f"{len(panel)} readings in, {len(level)} predicted cells out")

model.py

@@ -0,0 +1,149 @@
+"""Joint two-target network for a smooth 'prevailing' pollution level per station and day.
+
+Copied from the regression_room post (posts/beijing_air_quality/model.py):
+- 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  # single-target models ignore the other pollutant
+        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)  # small init so the area vectors are actually learned
+        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)