Download harness/graph.py from continker/metrollm-bench-mac: direct link, hf CLI and curl.
- Browser
- Download file 22.1 kB
-
https://huggingface.co/continker/metrollm-bench-mac/resolve/main/harness/graph.py
- Command line
-
hf download hf://continker/metrollm-bench-mac/harness/graph.py
-
curl -L -o graph.py https://huggingface.co/continker/metrollm-bench-mac/resolve/main/harness/graph.py
22.1 kB
| """Station graph operations using NetworkX — expanded line-graph routing.""" | |
| import json | |
| from collections import defaultdict | |
| from pathlib import Path | |
| from dataclasses import dataclass, field | |
| import networkx as nx | |
| TRANSFER_PENALTY_MIN = 5.0 | |
| class RouteResult: | |
| path: list[str] # station IDs in order | |
| stations: list[dict] # full station info per stop (name, line, is_transfer, etc.) | |
| distance_miles: float | |
| estimated_minutes: float | |
| transfers: int | |
| line_sequence: list[str] # e.g. ["red", "blue"] if transferring | |
| class MetroGraph: | |
| def __init__(self, system_dir: Path): | |
| """Load graph.json, stations.json, lines.json from a system directory.""" | |
| self.system_dir = system_dir | |
| with open(system_dir / "stations.json") as f: | |
| stations_list = json.load(f) | |
| self.stations: dict[str, dict] = {s["id"]: s for s in stations_list} | |
| with open(system_dir / "lines.json") as f: | |
| self.lines: dict[str, dict] = {l["id"]: l for l in json.load(f)} | |
| with open(system_dir / "graph.json") as f: | |
| graph_data = json.load(f) | |
| self._edges_raw: list[dict] = graph_data["edges"] | |
| # station_id -> set of line_ids serving that station (derived from edges) | |
| self.station_lines: dict[str, set[str]] = defaultdict(set) | |
| for edge in self._edges_raw: | |
| self.station_lines[edge["from"]].add(edge["line"]) | |
| self.station_lines[edge["to"]].add(edge["line"]) | |
| # Simple graph for connectivity checks (is_valid_path, adjacent_stations) | |
| self.G: nx.Graph = nx.Graph() | |
| for sid, sdata in self.stations.items(): | |
| self.G.add_node(sid, **sdata) | |
| for edge in self._edges_raw: | |
| self.G.add_edge( | |
| edge["from"], | |
| edge["to"], | |
| distance_miles=edge["distance_miles"], | |
| travel_time_min=edge["travel_time_min"], | |
| line=edge["line"], | |
| type=edge["type"], | |
| ) | |
| # Expanded directed graph for routing | |
| self._expanded = self._build_expanded(self._edges_raw, set(self.stations)) | |
| def _build_expanded( | |
| self, | |
| edges: list[dict], | |
| station_ids: set[str], | |
| ) -> nx.DiGraph: | |
| """Build the expanded line graph for transfer-aware Dijkstra. | |
| Nodes: | |
| ("enter", station_id) — virtual entry point | |
| (station_id, line_id) — station on a specific line | |
| ("exit", station_id) — virtual exit point | |
| Edges: | |
| entry: ("enter", s) → (s, line) weight=0, distance=0 | |
| exit: (s, line) → ("exit", s) weight=0, distance=0 | |
| travel: (sA, line) → (sB, line) weight=travel_time, distance=d | |
| transfer: (s, lineA) → (s, lineB) weight=TRANSFER_PENALTY_MIN, distance=0 | |
| """ | |
| G = nx.DiGraph() | |
| # Collect which lines serve each station | |
| station_lines: dict[str, set[str]] = defaultdict(set) | |
| for edge in edges: | |
| s_from, s_to = edge["from"], edge["to"] | |
| line = edge["line"] | |
| dist = edge["distance_miles"] | |
| time = edge["travel_time_min"] | |
| station_lines[s_from].add(line) | |
| station_lines[s_to].add(line) | |
| # Travel edges (both directions since graph is undirected) | |
| G.add_edge( | |
| (s_from, line), (s_to, line), | |
| weight=time, distance_miles=dist, line=line, | |
| edge_type="travel", | |
| ) | |
| G.add_edge( | |
| (s_to, line), (s_from, line), | |
| weight=time, distance_miles=dist, line=line, | |
| edge_type="travel", | |
| ) | |
| # Entry, exit, and transfer edges | |
| for sid in station_ids: | |
| lines = station_lines.get(sid, set()) | |
| for line in lines: | |
| # Entry | |
| G.add_edge( | |
| ("enter", sid), (sid, line), | |
| weight=0, distance_miles=0, edge_type="entry", | |
| ) | |
| # Exit | |
| G.add_edge( | |
| (sid, line), ("exit", sid), | |
| weight=0, distance_miles=0, edge_type="exit", | |
| ) | |
| # Transfer edges between all line pairs at this station | |
| lines_list = sorted(lines) | |
| for i, lineA in enumerate(lines_list): | |
| for lineB in lines_list[i + 1:]: | |
| G.add_edge( | |
| (sid, lineA), (sid, lineB), | |
| weight=TRANSFER_PENALTY_MIN, distance_miles=0, | |
| edge_type="transfer", | |
| ) | |
| G.add_edge( | |
| (sid, lineB), (sid, lineA), | |
| weight=TRANSFER_PENALTY_MIN, distance_miles=0, | |
| edge_type="transfer", | |
| ) | |
| return G | |
| def lines_for_station(self, station_id: str) -> set[str]: | |
| """Return the set of line ids that serve a station.""" | |
| sid = self._resolve_station(station_id) | |
| return set(self.station_lines.get(sid, set())) | |
| def _line_subgraph(self, line_id: str) -> nx.Graph: | |
| if line_id not in self.lines: | |
| raise ValueError(f"Unknown line: {line_id}") | |
| sub = nx.Graph() | |
| for edge in self._edges_raw: | |
| if edge["line"] == line_id: | |
| sub.add_edge(edge["from"], edge["to"]) | |
| return sub | |
| def is_loop_line(self, line_id: str) -> bool: | |
| """True if the line has no terminals (every station has degree >= 2 on its own line).""" | |
| sub = self._line_subgraph(line_id) | |
| if sub.number_of_nodes() == 0: | |
| return False | |
| return all(deg >= 2 for _, deg in sub.degree()) | |
| def line_terminals(self, line_id: str) -> list[str]: | |
| """Stations with degree 1 on the line subgraph. Empty list for loop lines.""" | |
| sub = self._line_subgraph(line_id) | |
| return [n for n, deg in sub.degree() if deg == 1] | |
| def expand_line_closures( | |
| self, | |
| closures: list[dict], | |
| ) -> list[tuple[str, str]]: | |
| """Expand line-level closures into segment_closures. | |
| Each closure dict: {"line": str, "from_station"?: str, "to_station"?: str}. | |
| Omitting both endpoints closes the entire line. Partial closure requires | |
| both endpoints and raises ValueError on a loop line (ambiguous). | |
| """ | |
| segments: list[tuple[str, str]] = [] | |
| for c in closures: | |
| line_id = c.get("line") | |
| if not line_id or line_id not in self.lines: | |
| raise ValueError(f"Unknown line: {line_id}") | |
| from_s = c.get("from_station") | |
| to_s = c.get("to_station") | |
| ordered = list(self.lines[line_id].get("stations", [])) | |
| if not ordered: | |
| raise ValueError(f"Line '{line_id}' has no stations defined") | |
| if from_s is None and to_s is None: | |
| keep = set(ordered) | |
| elif from_s is None or to_s is None: | |
| raise ValueError( | |
| f"Partial closure on line '{line_id}' requires both from_station and to_station" | |
| ) | |
| else: | |
| if self.is_loop_line(line_id): | |
| raise ValueError( | |
| f"Partial closure on loop line '{line_id}' is ambiguous — use whole-line closure or specify segments" | |
| ) | |
| a = self._resolve_station(from_s) | |
| b = self._resolve_station(to_s) | |
| if a not in ordered or b not in ordered: | |
| raise ValueError( | |
| f"Endpoints '{from_s}'/'{to_s}' are not on line '{line_id}'" | |
| ) | |
| i, j = ordered.index(a), ordered.index(b) | |
| lo, hi = min(i, j), max(i, j) | |
| keep = set(ordered[lo:hi + 1]) | |
| for edge in self._edges_raw: | |
| if ( | |
| edge["line"] == line_id | |
| and edge["from"] in keep | |
| and edge["to"] in keep | |
| ): | |
| segments.append((edge["from"], edge["to"])) | |
| return segments | |
| def shortest_path(self, origin: str, destination: str) -> RouteResult: | |
| """Find shortest path by time (with transfer penalty). Returns RouteResult. | |
| Raises ValueError if either station cannot be resolved. | |
| Raises nx.NetworkXNoPath if no path exists between the two stations. | |
| Raises nx.NodeNotFound if a resolved ID is not present in the graph. | |
| """ | |
| origin_id = self._resolve_station(origin) | |
| dest_id = self._resolve_station(destination) | |
| if origin_id == dest_id: | |
| station = self.stations[origin_id] | |
| stop = { | |
| "station_id": origin_id, | |
| "station_name": station["name"], | |
| "line": None, | |
| "is_transfer": False, | |
| "transfer_to": None, | |
| } | |
| return RouteResult( | |
| path=[origin_id], | |
| stations=[stop], | |
| distance_miles=0.0, | |
| estimated_minutes=0.0, | |
| transfers=0, | |
| line_sequence=[], | |
| ) | |
| return self._route_on_expanded(origin_id, dest_id, self._expanded) | |
| def shortest_path_avoiding( | |
| self, | |
| origin: str, | |
| destination: str, | |
| blocked_edges: list[tuple[str, str]] | None = None, | |
| blocked_stations: list[str] | None = None, | |
| ) -> RouteResult: | |
| """Compute shortest path avoiding specified edges and stations. | |
| Used by case generator for computing post-disruption alternative routes. | |
| Rebuilds the expanded graph with disrupted edges/stations removed. | |
| Raises ValueError if origin or destination is blocked or cannot be resolved. | |
| Raises nx.NetworkXNoPath if no alternative path exists. | |
| """ | |
| origin_id = self._resolve_station(origin) | |
| dest_id = self._resolve_station(destination) | |
| blocked_station_set = set(blocked_stations) if blocked_stations else set() | |
| blocked_edge_set = set() | |
| if blocked_edges: | |
| for u, v in blocked_edges: | |
| blocked_edge_set.add((u, v)) | |
| blocked_edge_set.add((v, u)) | |
| if origin_id in blocked_station_set: | |
| raise ValueError( | |
| f"Origin station '{origin}' is blocked by disruption" | |
| ) | |
| if dest_id in blocked_station_set: | |
| raise ValueError( | |
| f"Destination station '{destination}' is blocked by disruption" | |
| ) | |
| # Filter edges and stations | |
| remaining_stations = set(self.stations) - blocked_station_set | |
| remaining_edges = [ | |
| e for e in self._edges_raw | |
| if e["from"] not in blocked_station_set | |
| and e["to"] not in blocked_station_set | |
| and (e["from"], e["to"]) not in blocked_edge_set | |
| ] | |
| expanded = self._build_expanded(remaining_edges, remaining_stations) | |
| try: | |
| return self._route_on_expanded(origin_id, dest_id, expanded) | |
| except (nx.NetworkXNoPath, nx.NodeNotFound): | |
| raise nx.NetworkXNoPath( | |
| f"No alternative path between '{origin}' and '{destination}' " | |
| "with current disruption" | |
| ) | |
| def shortest_path_with_restrictions( | |
| self, | |
| origin: str, | |
| destination: str, | |
| station_restrictions: list[dict] | None = None, | |
| segment_closures: list[tuple[str, str]] | None = None, | |
| ) -> RouteResult: | |
| """Compute shortest path with typed station restrictions. | |
| station_restrictions: list of {"station": name_or_id, "restriction": type} | |
| - "closed": no entry, exit, transfer, or pass-through | |
| - "skip": trains pass through but don't stop (no entry/exit/transfer) | |
| - "no_transfer": can board/alight but cannot change lines | |
| segment_closures: list of (stationA, stationB) pairs where track is closed. | |
| Raises ValueError if origin/destination is closed or skip. | |
| Raises nx.NetworkXNoPath if no path exists with restrictions. | |
| """ | |
| if not station_restrictions and not segment_closures: | |
| return self.shortest_path(origin, destination) | |
| origin_id = self._resolve_station(origin) | |
| dest_id = self._resolve_station(destination) | |
| # Build restrictions map: station_id → restriction type | |
| restrictions_map: dict[str, str] = {} | |
| for r in (station_restrictions or []): | |
| sid = self._resolve_station(r["station"]) | |
| restrictions_map[sid] = r["restriction"] | |
| # Validate origin/destination | |
| for label, sid, name in [("Origin", origin_id, origin), | |
| ("Destination", dest_id, destination)]: | |
| restriction = restrictions_map.get(sid) | |
| if restriction in ("closed", "skip"): | |
| raise ValueError( | |
| f"{label} station '{name}' is {restriction} by disruption" | |
| ) | |
| # Build segment closure set (both directions) | |
| closed_segments: set[tuple[str, str]] = set() | |
| for seg in (segment_closures or []): | |
| u = self._resolve_station(seg[0]) | |
| v = self._resolve_station(seg[1]) | |
| closed_segments.add((u, v)) | |
| closed_segments.add((v, u)) | |
| # Build expanded graph with restrictions | |
| closed_stations = {s for s, r in restrictions_map.items() if r == "closed"} | |
| skip_stations = {s for s, r in restrictions_map.items() if r == "skip"} | |
| no_transfer_stations = {s for s, r in restrictions_map.items() | |
| if r == "no_transfer"} | |
| G = nx.DiGraph() | |
| station_lines: dict[str, set[str]] = defaultdict(set) | |
| # Phase 1: travel edges | |
| for edge in self._edges_raw: | |
| s_from, s_to = edge["from"], edge["to"] | |
| line = edge["line"] | |
| dist = edge["distance_miles"] | |
| time = edge["travel_time_min"] | |
| # Skip segment closures | |
| if (s_from, s_to) in closed_segments: | |
| continue | |
| # Skip travel edges touching closed stations | |
| if s_from in closed_stations or s_to in closed_stations: | |
| continue | |
| station_lines[s_from].add(line) | |
| station_lines[s_to].add(line) | |
| G.add_edge( | |
| (s_from, line), (s_to, line), | |
| weight=time, distance_miles=dist, line=line, | |
| edge_type="travel", | |
| ) | |
| G.add_edge( | |
| (s_to, line), (s_from, line), | |
| weight=time, distance_miles=dist, line=line, | |
| edge_type="travel", | |
| ) | |
| # Phase 2: entry, exit, transfer edges | |
| no_entry_exit = closed_stations | skip_stations | |
| no_transfer = closed_stations | skip_stations | no_transfer_stations | |
| for sid in set(self.stations) - closed_stations: | |
| lines = station_lines.get(sid, set()) | |
| if sid not in no_entry_exit: | |
| for line in lines: | |
| G.add_edge( | |
| ("enter", sid), (sid, line), | |
| weight=0, distance_miles=0, edge_type="entry", | |
| ) | |
| G.add_edge( | |
| (sid, line), ("exit", sid), | |
| weight=0, distance_miles=0, edge_type="exit", | |
| ) | |
| if sid not in no_transfer: | |
| lines_list = sorted(lines) | |
| for i, lineA in enumerate(lines_list): | |
| for lineB in lines_list[i + 1:]: | |
| G.add_edge( | |
| (sid, lineA), (sid, lineB), | |
| weight=TRANSFER_PENALTY_MIN, distance_miles=0, | |
| edge_type="transfer", | |
| ) | |
| G.add_edge( | |
| (sid, lineB), (sid, lineA), | |
| weight=TRANSFER_PENALTY_MIN, distance_miles=0, | |
| edge_type="transfer", | |
| ) | |
| try: | |
| return self._route_on_expanded(origin_id, dest_id, G) | |
| except (nx.NetworkXNoPath, nx.NodeNotFound): | |
| raise nx.NetworkXNoPath( | |
| f"No path between '{origin}' and '{destination}' " | |
| "with current restrictions" | |
| ) | |
| def _route_on_expanded( | |
| self, origin_id: str, dest_id: str, expanded: nx.DiGraph | |
| ) -> RouteResult: | |
| """Run Dijkstra on the expanded graph and convert to RouteResult.""" | |
| enter_node = ("enter", origin_id) | |
| exit_node = ("exit", dest_id) | |
| if enter_node not in expanded: | |
| raise nx.NodeNotFound( | |
| f"Node '{origin_id}' is not in the expanded graph" | |
| ) | |
| if exit_node not in expanded: | |
| raise nx.NodeNotFound( | |
| f"Node '{dest_id}' is not in the expanded graph" | |
| ) | |
| try: | |
| exp_path = nx.shortest_path( | |
| expanded, enter_node, exit_node, weight="weight" | |
| ) | |
| except nx.NetworkXNoPath: | |
| raise nx.NetworkXNoPath( | |
| f"No path found between '{origin_id}' and '{dest_id}'" | |
| ) | |
| # Convert expanded path to station-level RouteResult | |
| path: list[str] = [] | |
| stops: list[dict] = [] | |
| line_sequence: list[str] = [] | |
| total_distance = 0.0 | |
| total_time = 0.0 | |
| transfers = 0 | |
| current_line: str | None = None | |
| for i in range(len(exp_path) - 1): | |
| node = exp_path[i] | |
| next_node = exp_path[i + 1] | |
| edge_data = expanded[node][next_node] | |
| edge_type = edge_data["edge_type"] | |
| if edge_type == "entry": | |
| # (enter, station) -> (station, line): add origin station | |
| station_id = node[1] | |
| line = next_node[1] | |
| current_line = line | |
| if line not in line_sequence: | |
| line_sequence.append(line) | |
| station = self.stations[station_id] | |
| path.append(station_id) | |
| stops.append({ | |
| "station_id": station_id, | |
| "station_name": station["name"], | |
| "line": current_line, | |
| "is_transfer": False, | |
| "transfer_to": None, | |
| }) | |
| elif edge_type == "travel": | |
| # (stationA, line) -> (stationB, line): add stationB | |
| station_id = next_node[0] | |
| total_distance += edge_data["distance_miles"] | |
| total_time += edge_data["weight"] | |
| station = self.stations[station_id] | |
| path.append(station_id) | |
| stops.append({ | |
| "station_id": station_id, | |
| "station_name": station["name"], | |
| "line": current_line, | |
| "is_transfer": False, | |
| "transfer_to": None, | |
| }) | |
| elif edge_type == "transfer": | |
| # (station, lineA) -> (station, lineB): transfer at station | |
| new_line = next_node[1] | |
| transfers += 1 | |
| total_time += edge_data["weight"] | |
| if new_line not in line_sequence: | |
| line_sequence.append(new_line) | |
| # Mark the last stop as a transfer point | |
| if stops: | |
| stops[-1]["is_transfer"] = True | |
| stops[-1]["transfer_to"] = new_line | |
| current_line = new_line | |
| # exit edges: no action needed | |
| return RouteResult( | |
| path=path, | |
| stations=stops, | |
| distance_miles=round(total_distance, 2), | |
| estimated_minutes=round(total_time, 1), | |
| transfers=transfers, | |
| line_sequence=line_sequence, | |
| ) | |
| def is_valid_path(self, path: list[str]) -> bool: | |
| """Check if all consecutive stations in path are adjacent in the graph.""" | |
| if len(path) == 0: | |
| return False | |
| for i in range(len(path) - 1): | |
| if not self.G.has_edge(path[i], path[i + 1]): | |
| return False | |
| return True | |
| def adjacent_stations(self, station_id: str) -> list[str]: | |
| """Return neighbor station IDs for a given station. | |
| Raises ValueError if the station cannot be resolved. | |
| """ | |
| sid = self._resolve_station(station_id) | |
| return list(self.G.neighbors(sid)) | |
| def station_info(self, station_id: str) -> dict | None: | |
| """Return full station data, or None if the station does not exist.""" | |
| try: | |
| sid = self._resolve_station(station_id) | |
| except ValueError: | |
| return None | |
| return self.stations.get(sid) | |
| def _resolve_station(self, name_or_id: str) -> str: | |
| """Resolve a station name or ID to its canonical ID. | |
| Accepts an exact station ID or a station name (case-insensitive). | |
| Also matches the base name without parenthetical suffixes, e.g. | |
| "Olympic Park" matches "Aolinpike Gongyuan (Olympic Park)". | |
| Raises ValueError if no match is found. | |
| """ | |
| if name_or_id in self.stations: | |
| return name_or_id | |
| name_lower = name_or_id.lower().strip() | |
| for sid, sdata in self.stations.items(): | |
| full = sdata["name"].lower() | |
| # Exact match | |
| if full == name_lower: | |
| return sid | |
| # Match base name (before parenthetical) | |
| if "(" in full: | |
| base = full.split("(")[0].strip() | |
| paren = full.split("(")[1].rstrip(")").strip() | |
| if name_lower == base or name_lower == paren: | |
| return sid | |
| raise ValueError(f"Unknown station: '{name_or_id}'") | |