#!/usr/bin/env python3 """ F1 Race Data Fetcher Usage: python fetch_data.py [year] [grand_prix] [session] Example: python fetch_data.py 2024 Bahrain R python fetch_data.py 2023 Monaco Q """ import sys import json import math from pathlib import Path import fastf1 import pandas as pd import numpy as np def hex_color(team_color): """Normalise team color to #RRGGBB, fall back to white.""" if not team_color or (isinstance(team_color, float) and math.isnan(team_color)): return "#ffffff" s = str(team_color).strip() return s if s.startswith("#") else f"#{s}" def safe_seconds(td): """Convert timedelta / NaT to float seconds, or None.""" try: import pandas as _pd if td is None or td is _pd.NaT: return None v = float(td.total_seconds()) return None if math.isnan(v) else v except Exception: return None def safe_str(val, default="UNKNOWN"): """Return a clean uppercase string, treating NaN/None/empty as default.""" import pandas as _pd try: if val is None or val is _pd.NaT: return default if isinstance(val, float) and math.isnan(val): return default s = str(val).strip().upper() return s if s and s != 'NAN' and s != 'NONE' else default except Exception: return default def fetch_race_data( year: int = 2024, grand_prix: str = "Bahrain", session_type: str = "R", output_path: str = "race_data.json", progress_cb=None, ): if progress_cb is None: progress_cb = print ff1_cache = Path(output_path).parent / "cache" ff1_cache.mkdir(parents=True, exist_ok=True) fastf1.Cache.enable_cache(str(ff1_cache)) progress_cb(f"Loading {year} {grand_prix} {session_type} …") session = fastf1.get_session(year, grand_prix, session_type) session.load(telemetry=True, weather=False, messages=False) driver_data: dict = {} track_points: list = [] pit_lane: list | None = None # ── Load pre-built track layout (and pit lane) if available ─────────────── # build_tracks.py pre-generates layouts for all 24 circuits keyed by # slug(EventName), avoiding broken telemetry-derived outlines entirely. # Format is now: {slug: {"track": [[x,y],...], "pit": [[x,y],...] or null}} # Old list-only format is still handled for backwards compatibility. def _slug(s: str) -> str: return s.lower().replace(" ", "_").replace("/", "_").replace("'", "") layouts_file = Path(output_path).parent / "track_layouts.json" if layouts_file.exists(): try: with open(layouts_file, encoding="utf-8") as _lf: _layouts = json.load(_lf) _key = _slug(grand_prix) if _key in _layouts: _entry = _layouts[_key] if isinstance(_entry, dict): track_points = _entry.get("track") or [] pit_lane = _entry.get("pit") or None else: # Old format: entry is just the track points list track_points = _entry progress_cb(f" Track layout: {len(track_points)} pts (pre-built), " f"pit lane: {'yes' if pit_lane else 'none'}") except Exception: pass # ───────────────────────────────────────────────────────────────────────── for driver_num in session.drivers: try: drv = session.get_driver(driver_num) laps = session.laps.pick_driver(driver_num) if laps.empty: continue # Combined telemetry across all laps tel = laps.get_telemetry() if tel.empty or "X" not in tel.columns: continue tel = tel.dropna(subset=["X", "Y", "SessionTime"]) if tel.empty: continue # Build track outline from the single fastest clean lap's raw GPS trace. # One lap sorted by Distance gives the actual circuit shape. # Averaging across multiple laps distorts the track because cars take # different lines — the averaged X,Y at each distance bin is off-circuit. if not track_points: try: clean = laps.dropna(subset=["LapTime"]) if not clean.empty: min_lt = clean["LapTime"].min() # Use laps within 107% of fastest to exclude SC/formation laps clean = clean[clean["LapTime"] <= min_lt * 1.07] # Pick the fastest lap ref_lap = ( clean.loc[clean["LapTime"].idxmin()] if not clean.empty else laps.iloc[1] ) lap_tel = ref_lap.get_telemetry() lap_tel = lap_tel.dropna(subset=["X", "Y", "Distance"]) lap_tel = lap_tel.sort_values("Distance").reset_index(drop=True) if not lap_tel.empty: # Downsample to ~800 pts max; keep every point on short laps step = max(1, len(lap_tel) // 800) track_points = [ [round(float(r.X), 1), round(float(r.Y), 1)] for _, r in lap_tel.iloc[::step].iterrows() ] # Trim first/last ~2% to avoid start/finish GPS scatter trim = max(1, len(track_points) // 50) track_points = track_points[trim:-trim] except Exception: pass # Downsample to ~2500 points per driver step = max(1, len(tel) // 2500) tel_s = tel.iloc[::step] # ── Lap start times ──────────────────────────────────────────── lap_starts = [] for _, lap in laps.iterrows(): t = safe_seconds(lap.get("LapStartTime")) if t is not None: lap_starts.append(round(t, 2)) # ── Tyre compounds (one entry per lap where compound is known) ─ tyres = [] for _, lap in laps.iterrows(): t = safe_seconds(lap.get("LapStartTime")) if t is None: continue compound = safe_str(lap.get("Compound"), "UNKNOWN") raw_life = lap.get("TyreLife") try: tyre_life = int(float(raw_life)) if raw_life is not None else 0 except (ValueError, TypeError): tyre_life = 0 tyres.append([round(t, 2), compound, tyre_life]) # ── Pit stops ────────────────────────────────────────────────── # FastF1 layout: PitInTime on the inlap, PitOutTime on the outlap. # Build a lookup of session-time → PitOutTime from out-laps. pits = [] lap_list = list(laps.iterrows()) for idx, (_, lap) in enumerate(lap_list): pi = safe_seconds(lap.get("PitInTime")) if pi is None: continue # PitOutTime lives on the NEXT lap row (the out-lap) po = None if idx + 1 < len(lap_list): po = safe_seconds(lap_list[idx + 1][1].get("PitOutTime")) if po is None: po = pi + 28 pits.append([round(pi, 2), round(po, 2)]) # ── Telemetry positions ──────────────────────────────────────── positions = [] for _, row in tel_s.iterrows(): try: positions.append([ round(row["SessionTime"].total_seconds(), 2), round(float(row["X"]), 1), round(float(row["Y"]), 1), round(float(row.get("Speed", 0) or 0), 1), round(float(row.get("Distance", 0) or 0), 1), ]) except Exception: pass if not positions: continue name = drv.get("Abbreviation", str(driver_num)) driver_data[driver_num] = { "name": name, "team": drv.get("TeamName", ""), "color": hex_color(drv.get("TeamColor")), "laps": sorted(lap_starts), "tyres": tyres, "pits": pits, "positions": positions, } progress_cb(f" {name}: {len(positions)} pts, {len(pits)} pit(s)") except Exception as exc: progress_cb(f" Skipping {driver_num}: {exc}") if not driver_data: progress_cb("No driver data found — nothing exported.") return all_t = [p[0] for d in driver_data.values() for p in d["positions"]] # ── Track status (Safety Car / VSC / Yellow / Red) ──────────────────── track_status_list = [] try: ts_data = session.track_status if ts_data is not None and not ts_data.empty: for _, row in ts_data.iterrows(): t = safe_seconds(row.get("Time")) if t is None: continue status = str(row.get("Status", "")).strip() msg = str(row.get("Message", "")).strip() track_status_list.append([round(t, 2), status, msg]) except Exception: pass # ── Total laps for the session ──────────────────────────────────────── total_laps = 0 try: all_laps = session.laps for num in session.drivers: laps_drv = all_laps.pick_driver(num) n = int(laps_drv["LapNumber"].max()) if not laps_drv.empty else 0 if n > total_laps: total_laps = n except Exception: pass output = { "grand_prix": grand_prix, "year": year, "session": session_type, "track": track_points, "pit_lane": pit_lane, "drivers": driver_data, "t_start": round(min(all_t), 2), "t_end": round(max(all_t), 2), "track_status": track_status_list, "total_laps": total_laps, } Path(output_path).parent.mkdir(parents=True, exist_ok=True) with open(output_path, "w") as fh: json.dump(output, fh, separators=(",", ":")) size_mb = Path(output_path).stat().st_size / 1e6 progress_cb(f"\nExported {len(driver_data)} drivers → {Path(output_path).name} ({size_mb:.1f} MB)") if __name__ == "__main__": _year = int(sys.argv[1]) if len(sys.argv) > 1 else 2024 _gp = sys.argv[2] if len(sys.argv) > 2 else "Bahrain" _sess = sys.argv[3] if len(sys.argv) > 3 else "R" fetch_race_data(_year, _gp, _sess)