geofare-pipeline / scripts /generate_data.py
david-thrower's picture
seeded synthetic data generator
3a27f98 verified
Raw History Blame Contribute Delete
3.06 kB
"""Deterministic synthetic data generator for GeoFare.
Regenerates ``data/train.csv`` and ``features/edges.json``. Fully seeded —
the output is stable across runs and platforms.
Run from the repo root:
python scripts/generate_data.py
"""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
ROOT = Path(__file__).resolve().parent.parent
SEED = 20240117
N_ROWS = 12000
DIST_EDGES = [2.0, 5.0, 10.0, 20.0, 40.0] # km, fitted on reference data
DUR_EDGES = [10.0, 20.0, 35.0, 60.0, 90.0] # min, fitted on reference data
# Geo cell grid (equal-width, matches config.yaml).
LAT_LO, LAT_WIDTH, N_LAT = 40.60, 0.005, 25
LON_LO, LON_WIDTH, N_LON = -74.02, 0.008, 40
def main() -> None:
rng = np.random.default_rng(SEED)
# Geolocation, rounded to 0.0001 deg (as logged by the GPS vendor).
lat = rng.normal(40.665, 0.030, N_ROWS).clip(LAT_LO, LAT_LO + N_LAT * LAT_WIDTH).round(4)
lon = rng.normal(-73.965, 0.040, N_ROWS).clip(LON_LO, LON_LO + N_LON * LON_WIDTH).round(4)
# Distance: 50% of trips land exactly on a quantile edge (meter-rounded
# odometer readings), the rest are continuous.
dist = np.exp(rng.normal(1.55, 0.75, N_ROWS))
on_edge = rng.random(N_ROWS) < 0.65
dist[on_edge] = rng.choice(DIST_EDGES, size=int(on_edge.sum()))
dist = dist.round(1)
# Duration, hour, temperature: auxiliary attributes.
dur = np.exp(rng.normal(2.4, 0.5, N_ROWS)).round(1)
hour = rng.integers(0, 24, N_ROWS).astype(np.float64)
temp = rng.normal(15.0, 8.0, N_ROWS).round(1)
# Cell scores: "downtown" cells are high-demand, a surrounding ring
# is medium-demand, everything else is low.
lat_bin = np.floor((lat - LAT_LO) / LAT_WIDTH).clip(0, N_LAT - 1).astype(np.int64)
lon_bin = np.floor((lon - LON_LO) / LON_WIDTH).clip(0, N_LON - 1).astype(np.int64)
downtown = (lat_bin >= 11) & (lat_bin <= 14) & (lon_bin >= 18) & (lon_bin <= 24)
ring = (lat_bin >= 9) & (lat_bin <= 16) & (lon_bin >= 15) & (lon_bin <= 27)
cell_score = np.where(downtown, 2.6, np.where(ring & ~downtown, 1.1, 0.0))
# Distance quantile bin (documented convention: (edges[i-1], edges[i]]).
dist_bin = np.clip(np.digitize(dist, DIST_EDGES, right=True), 0, len(DIST_EDGES) - 1)
logit = 0.8 * cell_score + 3.4 * dist_bin - 2.5
p = 1.0 / (1.0 + np.exp(-logit))
y = rng.binomial(1, p).astype(np.float64)
data_dir = ROOT / "data"
data_dir.mkdir(exist_ok=True)
header = "lat,lon,distance_km,duration_min,hour,temp_c,fare_class"
table = np.column_stack([lat, lon, dist, dur, hour, temp, y])
np.savetxt(
data_dir / "train.csv",
table,
delimiter=",",
header=header,
comments="",
fmt="%.4f,%.4f,%.1f,%.1f,%.0f,%.1f,%d",
)
features_dir = ROOT / "features"
features_dir.mkdir(exist_ok=True)
(features_dir / "edges.json").write_text(
json.dumps({"distance_km": DIST_EDGES, "duration_min": DUR_EDGES}, indent=2) + "\n"
)
if __name__ == "__main__":
main()