1"""Hand-written numpy MLP that predicts taxi trip duration (minutes).2 3Inputs (all available at planning time, no leakage of the actual trip):4 pickup zone one-hot (Z) -> the first layer's weights on it are a learned zone embedding5 drop-off zone one-hot (Z) -> second embedding6 hour one-hot (24) + sin/cos of the hour (smoothness hint)7 weekday one-hot (5)8 zone-pair mean distance (train split only), standardized, plus its log9 10Network: D -> 64 -> 32 -> 1, ReLU, Huber loss on the standardized target,11Adam, mini-batches, L2 weight decay, early stopping on validation MAE.12"""13from __future__ import annotations14 15import time16from dataclasses import dataclass, field17from pathlib import Path18 19import numpy as np20 21 22# --------------------------------------------------------------------------- features23class Featurizer:24 def __init__(self, n_zones: int, pair_dist: np.ndarray):25 self.Z = n_zones26 self.pair_dist = np.asarray(pair_dist, dtype=float)27 d = self.pair_dist.ravel()28 self.d_mu, self.d_sd = float(d.mean()), float(d.std() + 1e-9)29 ld = np.log1p(d)30 self.ld_mu, self.ld_sd = float(ld.mean()), float(ld.std() + 1e-9)31 32 @property33 def dim(self) -> int:34 return 2 * self.Z + 24 + 2 + 5 + 235 36 def transform(self, pu, do, hour, weekday) -> np.ndarray:37 pu, do = np.asarray(pu, int), np.asarray(do, int)38 hour, weekday = np.asarray(hour, int), np.asarray(weekday, int)39 n, Z = len(pu), self.Z40 X = np.zeros((n, self.dim), dtype=np.float32)41 r = np.arange(n)42 X[r, pu] = 1.043 X[r, Z + do] = 1.044 X[r, 2 * Z + hour] = 1.045 o = 2 * Z + 2446 X[:, o] = np.sin(2 * np.pi * hour / 24)47 X[:, o + 1] = np.cos(2 * np.pi * hour / 24)48 X[r, o + 2 + weekday] = 1.049 d = self.pair_dist[pu, do]50 X[:, o + 7] = (d - self.d_mu) / self.d_sd51 X[:, o + 8] = (np.log1p(d) - self.ld_mu) / self.ld_sd52 return X53 54 def state(self) -> dict:55 return {"Z": self.Z, "pair_dist": self.pair_dist}56 57 @classmethod58 def from_state(cls, s) -> "Featurizer":59 return cls(int(s["Z"]), s["pair_dist"])60 61 62# --------------------------------------------------------------------------- network63def huber(r: np.ndarray, delta: float):64 a = np.abs(r)65 loss = np.where(a <= delta, 0.5 * r * r, delta * (a - 0.5 * delta))66 grad = np.where(a <= delta, r, delta * np.sign(r))67 return loss.mean(), grad / len(r)68 69 70class MLP:71 def __init__(self, sizes, seed: int = 0):72 rng = np.random.default_rng(seed)73 self.sizes = list(sizes)74 self.W, self.b = [], []75 for a, b in zip(sizes[:-1], sizes[1:]):76 self.W.append(rng.normal(0, np.sqrt(2.0 / a), size=(a, b)).astype(np.float64))77 self.b.append(np.zeros(b))78 79 def params(self):80 return self.W + self.b81 82 def forward(self, X):83 acts = [X]84 h = X85 for k, (W, b) in enumerate(zip(self.W, self.b)):86 z = h @ W + b87 h = np.maximum(z, 0.0) if k < len(self.W) - 1 else z88 acts.append(h)89 return h[:, 0], acts90 91 def backward(self, acts, dout):92 """dout: dL/d(output) shape (n,). Returns grads aligned with params()."""93 gW = [None] * len(self.W)94 gb = [None] * len(self.b)95 g = dout[:, None]96 for k in reversed(range(len(self.W))):97 gW[k] = acts[k].T @ g98 gb[k] = g.sum(axis=0)99 if k > 0:100 g = (g @ self.W[k].T) * (acts[k] > 0)101 return gW + gb102 103 104class Adam:105 def __init__(self, params, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8, weight_decay=0.0):106 self.p = params107 self.lr, self.b1, self.b2, self.eps, self.wd = lr, b1, b2, eps, weight_decay108 self.m = [np.zeros_like(x) for x in params]109 self.v = [np.zeros_like(x) for x in params]110 self.t = 0111 112 def step(self, grads, n_weight_mats: int):113 self.t += 1114 for i, (p, g) in enumerate(zip(self.p, grads)):115 if self.wd and i < n_weight_mats: # decay weights, not biases116 g = g + self.wd * p117 self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * g118 self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * g * g119 mh = self.m[i] / (1 - self.b1 ** self.t)120 vh = self.v[i] / (1 - self.b2 ** self.t)121 p -= self.lr * mh / (np.sqrt(vh) + self.eps)122 123 124# --------------------------------------------------------------------------- model125@dataclass126class TrainConfig:127 hidden: tuple = (64, 32)128 lr: float = 1e-3129 batch: int = 512130 epochs: int = 40131 patience: int = 5132 weight_decay: float = 1e-5133 huber_delta: float = 1.0134 seed: int = 42135 136 137@dataclass138class TravelTimeNN:139 feat: Featurizer140 net: MLP141 y_mu: float142 y_sd: float143 history: list = field(default_factory=list)144 145 def predict(self, pu, do, hour, weekday) -> np.ndarray:146 X = self.feat.transform(pu, do, hour, weekday)147 out, _ = self.net.forward(X)148 return np.clip(out * self.y_sd + self.y_mu, 1.0, 120.0)149 150 # ---- persistence151 def save(self, path: Path):152 arrs = {f"W{i}": w for i, w in enumerate(self.net.W)}153 arrs.update({f"b{i}": b for i, b in enumerate(self.net.b)})154 np.savez_compressed(path, sizes=np.array(self.net.sizes), y_mu=self.y_mu, y_sd=self.y_sd,155 Z=self.feat.Z, pair_dist=self.feat.pair_dist,156 history=np.array(self.history, dtype=float), **arrs)157 158 @classmethod159 def load(cls, path: Path) -> "TravelTimeNN":160 f = np.load(path)161 sizes = [int(s) for s in f["sizes"]]162 net = MLP(sizes)163 net.W = [f[f"W{i}"] for i in range(len(sizes) - 1)]164 net.b = [f[f"b{i}"] for i in range(len(sizes) - 1)]165 feat = Featurizer(int(f["Z"]), f["pair_dist"])166 return cls(feat, net, float(f["y_mu"]), float(f["y_sd"]), f["history"].tolist())167 168 169def train(trips_train, trips_val, n_zones: int, pair_dist: np.ndarray,170 cfg: TrainConfig = TrainConfig(), verbose: bool = True) -> TravelTimeNN:171 feat = Featurizer(n_zones, pair_dist)172 cols = ["pu", "do", "hour", "weekday"]173 Xtr = feat.transform(*[trips_train[c].values for c in cols]).astype(np.float64)174 Xva = feat.transform(*[trips_val[c].values for c in cols]).astype(np.float64)175 ytr = trips_train.duration_min.values.astype(float)176 yva = trips_val.duration_min.values.astype(float)177 y_mu, y_sd = float(ytr.mean()), float(ytr.std())178 ttr = (ytr - y_mu) / y_sd179 180 net = MLP([feat.dim, *cfg.hidden, 1], seed=cfg.seed)181 opt = Adam(net.params(), lr=cfg.lr, weight_decay=cfg.weight_decay)182 rng = np.random.default_rng(cfg.seed)183 model = TravelTimeNN(feat, net, y_mu, y_sd)184 best, best_state, bad = np.inf, None, 0185 n = len(Xtr)186 for ep in range(1, cfg.epochs + 1):187 t0 = time.time()188 perm = rng.permutation(n)189 tot = 0.0190 for s in range(0, n, cfg.batch):191 idx = perm[s:s + cfg.batch]192 out, acts = net.forward(Xtr[idx])193 loss, dout = huber(out - ttr[idx], cfg.huber_delta)194 tot += loss * len(idx)195 opt.step(net.backward(acts, dout), len(net.W))196 pva = np.clip(net.forward(Xva)[0] * y_sd + y_mu, 1, 120)197 va_mae = float(np.abs(pva - yva).mean())198 model.history.append((ep, tot / n, va_mae))199 if verbose:200 print(f" epoch {ep:2d} train huber {tot/n:.4f} val MAE {va_mae:.3f} min ({time.time()-t0:.1f}s)")201 if va_mae < best - 1e-4:202 best, bad = va_mae, 0203 best_state = ([w.copy() for w in net.W], [b.copy() for b in net.b])204 else:205 bad += 1206 if bad >= cfg.patience:207 if verbose:208 print(f" early stop at epoch {ep} (best val MAE {best:.3f})")209 break210 net.W, net.b = best_state211 # Adam holds references to the old arrays; the model itself now uses the best weights212 return model213 214 215def numerical_grad_check(seed: int = 0, eps: float = 1e-6) -> float:216 """Max relative error between backprop and finite-difference gradients (for tests)."""217 rng = np.random.default_rng(seed)218 net = MLP([7, 5, 4, 1], seed=seed)219 X = rng.normal(size=(9, 7))220 y = rng.normal(size=9)221 222 def L():223 return huber(net.forward(X)[0] - y, 1.0)[0]224 225 out, acts = net.forward(X)226 _, dout = huber(out - y, 1.0)227 grads = net.backward(acts, dout)228 worst = 0.0229 for p, g in zip(net.params(), grads):230 it = np.nditer(p, flags=["multi_index"])231 for _ in it:232 i = it.multi_index233 old = p[i]234 p[i] = old + eps; lp = L()235 p[i] = old - eps; lm = L()236 p[i] = old237 num = (lp - lm) / (2 * eps)238 worst = max(worst, abs(num - g[i]) / max(1e-8, abs(num) + abs(g[i])))239 return worst