1"""Baselines the NN is compared against.2 3HistoricalAverage : mean duration per (pickup zone, drop-off zone, hour) on the train split,4 backing off to (pickup, drop-off) and then the global mean when a cell has5 fewer than `min_count` trips (rare pairs at night).6DistanceLinear : ordinary least squares duration = a + b * zone-pair mean distance7 (numpy least squares; scikit-learn was dropped to save disk space).8 No time-of-day information at all.9"""10from __future__ import annotations11 12import numpy as np13 14 15class HistoricalAverage:16 name = "hist-avg"17 18 def __init__(self, n_zones: int, min_count: int = 3):19 self.Z, self.min_count = n_zones, min_count20 21 def fit(self, trips):22 Z = self.Z23 y = trips.duration_min.values24 self.global_mean = float(y.mean())25 s3 = np.zeros((Z, Z, 24)); n3 = np.zeros((Z, Z, 24))26 np.add.at(s3, (trips.pu.values, trips["do"].values, trips.hour.values), y)27 np.add.at(n3, (trips.pu.values, trips["do"].values, trips.hour.values), 1)28 s2, n2 = s3.sum(2), n3.sum(2)29 pair = np.where(n2 >= self.min_count, s2 / np.maximum(n2, 1), self.global_mean)30 cell = np.where(n3 >= self.min_count, s3 / np.maximum(n3, 1), pair[:, :, None])31 self.table, self.counts = cell, n332 return self33 34 def predict(self, pu, do, hour, weekday=None):35 return self.table[np.asarray(pu, int), np.asarray(do, int), np.asarray(hour, int)]36 37 38class DistanceLinear:39 name = "distance-linear"40 41 def __init__(self, pair_dist: np.ndarray):42 self.pair_dist = np.asarray(pair_dist, float)43 44 def fit(self, trips):45 x = self.pair_dist[trips.pu.values, trips["do"].values]46 X = np.column_stack([np.ones_like(x), x])47 (self.intercept_, self.slope_), *_ = np.linalg.lstsq(X, trips.duration_min.values, rcond=None)48 return self49 50 def predict(self, pu, do, hour=None, weekday=None):51 x = self.pair_dist[np.asarray(pu, int), np.asarray(do, int)]52 return np.clip(self.intercept_ + self.slope_ * x, 1.0, 120.0)53 54 55class NNPredictor:56 """Adapter so the NN has the same predict() signature as the baselines."""57 name = "nn"58 59 def __init__(self, model):60 self.model = model61 62 def predict(self, pu, do, hour, weekday):63 return self.model.predict(pu, do, hour, weekday)