1"""Real data: download and preprocess NYC TLC yellow-taxi trips.2 3Pipeline (all deterministic, seeded):4 1. download one month of trips + zone lookup + zone shapefile5 2. check the zone-ID columns exist, parse timestamps6 3. keep trips inside the month, on weekdays, duration in [1, 120] min,7 positive distance and plausible speed (<= 60 mph)8 4. pick the 20 busiest zones (pickups + drop-offs, unknown zones 264/265 excluded)9 and keep trips with BOTH ends in those zones10 5. subsample ~300k trips with a fixed seed11 6. split by day: Jan 3-19 train, Jan 20-21 validation, Jan 24-31 test12 7. zone-pair mean distance from the TRAIN split only (fallback for unseen pairs:13 linear fit on zone-centroid distance, also fitted on train only)14"""15from __future__ import annotations16 17import json18import urllib.request19from pathlib import Path20 21import numpy as np22import pandas as pd23 24from . import config as C25 26NEEDED_COLS = ["tpep_pickup_datetime", "tpep_dropoff_datetime",27 "PULocationID", "DOLocationID", "trip_distance"]28 29 30# --------------------------------------------------------------------------- download31def download(force: bool = False) -> dict:32 C.DATA_RAW.mkdir(parents=True, exist_ok=True)33 files = {34 "trips": (C.TRIP_URL, C.DATA_RAW / f"yellow_tripdata_{C.MONTH}.parquet"),35 "lookup": (C.ZONE_LOOKUP_URL, C.DATA_RAW / "taxi_zone_lookup.csv"),36 "shapes": (C.ZONE_SHAPES_URL, C.DATA_RAW / "taxi_zones.zip"),37 }38 out = {}39 for key, (url, path) in files.items():40 if force or not path.exists() or path.stat().st_size == 0:41 print(f" downloading {url}")42 urllib.request.urlretrieve(url, path)43 else:44 print(f" cached {path.name} ({path.stat().st_size/1e6:.1f} MB)")45 out[key] = path46 return out47 48 49# --------------------------------------------------------------------------- zones50def _ring_centroid(pts: np.ndarray) -> tuple[float, float, float]:51 x, y = pts[:, 0], pts[:, 1]52 x1, y1 = np.roll(x, -1), np.roll(y, -1)53 cross = x * y1 - x1 * y54 a = cross.sum() / 2.055 if abs(a) < 1e-9:56 return float(x.mean()), float(y.mean()), 0.057 cx = ((x + x1) * cross).sum() / (6 * a)58 cy = ((y + y1) * cross).sum() / (6 * a)59 return float(cx), float(cy), float(abs(a))60 61 62def read_zone_shapes(zip_path: Path) -> dict[int, list[np.ndarray]]:63 """LocationID -> list of polygon rings in km (EPSG:2263 feet * 0.0003048)."""64 import shapefile # pyshp65 r = shapefile.Reader(str(zip_path))66 fields = [f[0] for f in r.fields[1:]]67 id_col = fields.index("LocationID")68 shapes: dict[int, list[np.ndarray]] = {}69 for sr in r.iterShapeRecords():70 lid = int(sr.record[id_col])71 pts = np.asarray(sr.shape.points, dtype=float) * 0.000304872 parts = list(sr.shape.parts) + [len(pts)]73 for a, b in zip(parts[:-1], parts[1:]):74 shapes.setdefault(lid, []).append(pts[a:b])75 return shapes76 77 78def zone_centroids(zip_path: Path) -> pd.DataFrame:79 rows = []80 for lid, rings in read_zone_shapes(zip_path).items():81 cs = [_ring_centroid(r) for r in rings]82 w = np.array([c[2] for c in cs]) + 1e-1283 cx = float(np.average([c[0] for c in cs], weights=w))84 cy = float(np.average([c[1] for c in cs], weights=w))85 rows.append((lid, cx, cy))86 return pd.DataFrame(rows, columns=["LocationID", "cx_km", "cy_km"])87 88 89# --------------------------------------------------------------------------- trips90def load_raw_trips(path: Path) -> pd.DataFrame:91 import pyarrow.parquet as pq92 schema = pq.read_schema(path)93 missing = [c for c in NEEDED_COLS if c not in schema.names]94 if missing:95 raise ValueError(f"zone/time columns missing in {path.name}: {missing}")96 return pq.read_table(path, columns=NEEDED_COLS).to_pandas()97 98 99def preprocess(raw: pd.DataFrame, lookup: pd.DataFrame, cents: pd.DataFrame,100 seed: int = C.SEED, n_zones: int = C.N_ZONES,101 target: int = C.TARGET_TRIPS):102 stats = {"raw_rows": int(len(raw))}103 df = raw.rename(columns={"tpep_pickup_datetime": "pickup", "tpep_dropoff_datetime": "dropoff",104 "PULocationID": "pu_id", "DOLocationID": "do_id",105 "trip_distance": "distance_mi"})106 month_start = pd.Timestamp(C.MONTH + "-01")107 month_end = month_start + pd.offsets.MonthBegin(1)108 df = df[(df.pickup >= month_start) & (df.pickup < month_end)]109 stats["in_month"] = int(len(df))110 df = df[df.pickup.dt.weekday < 5]111 stats["weekdays"] = int(len(df))112 df = df.assign(duration_min=(df.dropoff - df.pickup).dt.total_seconds() / 60.0)113 df = df[(df.duration_min >= C.MIN_DURATION) & (df.duration_min <= C.MAX_DURATION)]114 stats["duration_1_120"] = int(len(df))115 speed = df.distance_mi / (df.duration_min / 60.0)116 df = df[(df.distance_mi > 0) & (speed <= C.MAX_SPEED_MPH)]117 stats["distance_speed_ok"] = int(len(df))118 119 known = df[(df.pu_id < 264) & (df.do_id < 264)]120 activity = known.pu_id.value_counts().add(known.do_id.value_counts(), fill_value=0)121 top = [int(z) for z in activity.sort_values(ascending=False).index[:n_zones]]122 top = sorted(top)123 df = df[df.pu_id.isin(top) & df.do_id.isin(top)]124 stats["both_ends_top_zones"] = int(len(df))125 if len(df) > target:126 df = df.sample(n=target, random_state=seed)127 stats["sampled"] = int(len(df))128 129 idx = {z: i for i, z in enumerate(top)}130 day = df.pickup.dt.normalize()131 split = np.where(day >= pd.Timestamp(C.TEST_START), "test",132 np.where(day.dt.strftime("%Y-%m-%d").isin(C.VAL_DAYS), "val", "train"))133 trips = pd.DataFrame({134 "pickup": df.pickup.values,135 "date": day.dt.strftime("%Y-%m-%d").values,136 "pu_id": df.pu_id.values.astype(int),137 "do_id": df.do_id.values.astype(int),138 "pu": df.pu_id.map(idx).values.astype(int),139 "do": df.do_id.map(idx).values.astype(int),140 "hour": df.pickup.dt.hour.values.astype(int),141 "weekday": df.pickup.dt.weekday.values.astype(int),142 "distance_mi": df.distance_mi.values.astype(float),143 "duration_min": df.duration_min.values.astype(float),144 "split": split,145 }).sort_values("pickup").reset_index(drop=True)146 for s in ("train", "val", "test"):147 stats[f"n_{s}"] = int((trips.split == s).sum())148 stats[f"days_{s}"] = sorted(trips.loc[trips.split == s, "date"].unique().tolist())149 150 zones = (lookup[lookup.LocationID.isin(top)]151 .merge(cents, on="LocationID", how="left")152 .sort_values("LocationID").reset_index(drop=True))153 zones.insert(0, "idx", [idx[z] for z in zones.LocationID])154 zones["n_pickups"] = zones.idx.map(trips.pu.value_counts()).fillna(0).astype(int)155 zones["n_dropoffs"] = zones.idx.map(trips["do"].value_counts()).fillna(0).astype(int)156 zones = zones.rename(columns={"Borough": "borough", "Zone": "zone", "service_zone": "service_zone"})157 return trips, zones, stats158 159 160def pair_distance_matrix(trips: pd.DataFrame, zones: pd.DataFrame):161 """Mean trip distance per (pu, do) pair from TRAIN rows only.162 163 Pairs never seen in train get a linear fit on centroid distance (fitted on train pairs).164 Returns (dist[Z,Z] miles, count[Z,Z], fit coefficients).165 """166 Z = len(zones)167 tr = trips[trips.split == "train"]168 g = tr.groupby(["pu", "do"]).distance_mi.agg(["mean", "size"])169 dist = np.full((Z, Z), np.nan)170 cnt = np.zeros((Z, Z), dtype=int)171 for (i, j), r in g.iterrows():172 dist[i, j], cnt[i, j] = r["mean"], r["size"]173 xy = zones.sort_values("idx")[["cx_km", "cy_km"]].to_numpy()174 cd_mi = np.linalg.norm(xy[:, None, :] - xy[None, :, :], axis=2) / 1.609344175 seen = ~np.isnan(dist)176 b, a = np.polyfit(cd_mi[seen], dist[seen], 1)177 dist = np.where(seen, dist, a + b * cd_mi)178 return dist, cnt, {"intercept": float(a), "slope_per_centroid_mile": float(b),179 "unseen_pairs": int((~seen).sum())}180 181 182# --------------------------------------------------------------------------- io183def save_processed(trips, zones, dist, cnt, stats, out: Path = C.DATA_PROC):184 out.mkdir(parents=True, exist_ok=True)185 trips.to_parquet(out / "trips.parquet", index=False)186 zones.to_csv(out / "zones.csv", index=False)187 np.save(out / "pair_distance.npy", dist)188 np.save(out / "pair_count.npy", cnt)189 (out / "preprocess_stats.json").write_text(json.dumps(stats, indent=2))190 191 192def load_processed(base: Path = C.DATA_PROC):193 """Returns (trips, zones, pair_distance). Falls back to the committed test fixture."""194 if (base / "trips.parquet").exists():195 trips = pd.read_parquet(base / "trips.parquet")196 zones = pd.read_csv(base / "zones.csv")197 dist = np.load(base / "pair_distance.npy")198 return trips, zones, dist199 return load_fixture()200 201 202def write_fixture(trips, zones, dist, n: int = 3000, seed: int = C.SEED, out: Path = C.FIXTURES):203 out.mkdir(parents=True, exist_ok=True)204 parts = [trips[trips.split == s].sample(n=min(n // 3, (trips.split == s).sum()), random_state=seed)205 for s in ("train", "val", "test")]206 sample = pd.concat(parts).sort_values("pickup")207 sample.drop(columns=["pickup"]).to_csv(out / "sample_trips.csv", index=False)208 zones.to_csv(out / "zones.csv", index=False)209 np.savetxt(out / "pair_distance.csv", dist, delimiter=",", fmt="%.4f")210 211 212def load_fixture(base: Path = C.FIXTURES):213 trips = pd.read_csv(base / "sample_trips.csv")214 zones = pd.read_csv(base / "zones.csv")215 dist = np.loadtxt(base / "pair_distance.csv", delimiter=",")216 return trips, zones, dist