1"""Shared matplotlib styling and the figures used by train_nn.py, run_demo.py and run_experiments.py."""2from __future__ import annotations3 4import matplotlib5 6matplotlib.use("Agg")7import matplotlib.pyplot as plt # noqa: E4028import numpy as np # noqa: E4029 10# validated categorical slots (first three pass all-pairs CVD checks) + neutral inks11SERIES = ["#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#6250d6"]12VAN_COLORS = SERIES[:3]13SURFACE, INK, INK2, MUTED, GRID = "#fcfcfb", "#0b0b0b", "#52514e", "#8a8984", "#e4e3df"14LATE = "#e34948"15 16 17def style():18 plt.rcParams.update({19 "figure.facecolor": SURFACE, "axes.facecolor": SURFACE, "savefig.facecolor": SURFACE,20 "axes.edgecolor": MUTED, "axes.labelcolor": INK2, "xtick.color": INK2, "ytick.color": INK2,21 "text.color": INK, "axes.titlecolor": INK, "axes.titleweight": "bold", "axes.titlesize": 12,22 "axes.spines.top": False, "axes.spines.right": False, "axes.grid": True,23 "grid.color": GRID, "grid.linewidth": 0.8, "axes.axisbelow": True,24 "lines.linewidth": 2.0, "font.size": 10, "legend.frameon": False, "figure.dpi": 110,25 })26 27 28style()29 30 31def save(fig, path):32 fig.tight_layout()33 fig.savefig(path, dpi=130, bbox_inches="tight")34 plt.close(fig)35 return path36 37 38def hm(t):39 t = int(round(t))40 return f"{t // 60:02d}:{t % 60:02d}"41 42 43# --------------------------------------------------------------------------- NN figures44def learning_curve(history, path):45 h = np.asarray(history)46 fig, ax = plt.subplots(figsize=(6, 3.4))47 ax.plot(h[:, 0], h[:, 2], color=SERIES[0], marker="o", ms=4, label="validation MAE (min)")48 best = int(np.argmin(h[:, 2]))49 ax.scatter([h[best, 0]], [h[best, 2]], s=80, facecolor="none", edgecolor=INK, zorder=5)50 ax.annotate(f"best epoch {int(h[best, 0])}: {h[best, 2]:.2f} min", (h[best, 0], h[best, 2]),51 xytext=(8, 10), textcoords="offset points", color=INK2)52 ax.set_xlabel("epoch"); ax.set_ylabel("validation MAE (minutes)")53 ax.set_title("NN training: early stopping on validation MAE")54 return save(fig, path)55 56 57def mae_by_hour(test, preds: dict, path):58 fig, ax = plt.subplots(figsize=(7.5, 3.6))59 hours = np.arange(24)60 for k, (name, p) in enumerate(preds.items()):61 err = np.abs(np.asarray(p) - test.duration_min.values)62 m = [err[test.hour.values == h].mean() if (test.hour.values == h).any() else np.nan for h in hours]63 ax.plot(hours, m, color=SERIES[k], marker="o", ms=3.5, label=name)64 cnt = np.bincount(test.hour.values, minlength=24)65 ax2 = ax.inset_axes([0.0, -0.42, 1.0, 0.22], sharex=ax)66 ax2.bar(hours, cnt, color=MUTED, width=0.8)67 ax2.set_ylabel("test trips", fontsize=8); ax2.tick_params(labelsize=8); ax2.grid(False)68 ax.set_xlim(-0.5, 23.5); ax.set_xticks(range(0, 24, 3))69 ax.set_ylabel("MAE (minutes)"); ax.set_title("Test MAE by pickup hour")70 ax.legend(loc="upper left", ncol=len(preds))71 return save(fig, path)72 73 74def pred_vs_actual(y, p, path, n=4000, seed=0):75 rng = np.random.default_rng(seed)76 idx = rng.choice(len(y), size=min(n, len(y)), replace=False)77 fig, ax = plt.subplots(figsize=(4.6, 4.4))78 ax.scatter(np.asarray(y)[idx], np.asarray(p)[idx], s=8, alpha=0.25, color=SERIES[0], edgecolor="none")79 ax.plot([0, 60], [0, 60], color=INK2, lw=1, ls="--")80 ax.set_xlim(0, 60); ax.set_ylim(0, 60)81 ax.set_xlabel("actual duration (min)"); ax.set_ylabel("NN predicted (min)")82 ax.set_title("NN on held-out week (4k sample)")83 return save(fig, path)84 85 86def matrix_heatmaps(M, zones, path, hours=(8, 18)):87 labels = [z[:24] for z in zones.sort_values("idx").zone]88 fig, axes = plt.subplots(1, len(hours) + 1, figsize=(16, 5.6),89 gridspec_kw={"width_ratios": [1] * len(hours) + [0.9], "wspace": 0.12})90 vmax = np.percentile(M[list(hours)], 98)91 for k, (ax, h) in enumerate(zip(axes, hours)):92 im = ax.imshow(M[h], cmap="Blues", vmin=0, vmax=vmax)93 ax.set_title(f"NN predicted taxi minutes, {h:02d}:00"); ax.grid(False)94 ax.set_xticks(range(len(labels))); ax.set_yticks(range(len(labels)))95 ax.set_xticklabels(labels, rotation=90, fontsize=6.5)96 ax.set_yticklabels(labels if k == 0 else [], fontsize=6.5)97 ax.set_xlabel("drop-off zone")98 if k == 0:99 ax.set_ylabel("pickup zone")100 fig.colorbar(im, ax=axes[:len(hours)], shrink=0.7, label="minutes")101 ax = axes[-1]102 off = ~np.eye(M.shape[1], dtype=bool)103 ax.plot(range(24), [M[h][off].mean() for h in range(24)], color=SERIES[0], marker="o", ms=3)104 ax.set_title("Mean over zone pairs by hour"); ax.set_xlabel("departure hour"); ax.set_ylabel("minutes")105 ax.set_xticks(range(0, 24, 3))106 fig.savefig(path, dpi=130, bbox_inches="tight"); plt.close(fig)107 return path