1"""Loading helpers shared by run_demo.py and run_experiments.py."""2from __future__ import annotations3 4import numpy as np5 6from . import config as C7from .data import load_processed8from .travel_time import TimeModel9 10MODEL_NAMES = {"nn": "NN (MLP)", "hist": "hist-avg", "dist": "distance-linear"}11 12 13def load_all():14 trips, zones, dist = load_processed()15 path = C.MODELS / "time_matrices.npz"16 if not path.exists():17 raise SystemExit("models/time_matrices.npz missing -- run `python train_nn.py` first")18 mats = dict(np.load(path))19 return trips, zones, dist, mats20 21 22def time_model(mats, kind: str, weekday: int, factor: float = C.VAN_FACTOR) -> TimeModel:23 return TimeModel(mats[kind][weekday], factor=factor, name=MODEL_NAMES.get(kind, kind))24 25 26def zone_weights(zones):27 return zones.sort_values("idx").n_dropoffs.values