import os import numpy as np import pandas as pd import plotly.graph_objects as go class Utils: def __init__(self, dataset_dir = "dataset/"): self.london_stations_df = pd.read_csv(os.path.join(dataset_dir,"stations.csv")) self.london_connections_df = pd.read_csv(os.path.join(dataset_dir,"connections.csv")) self.london_lines_df = pd.read_csv(os.path.join(dataset_dir,"lines.csv")) def get_all_station_names(self): return self.london_stations_df['name'].tolist() def remove_spl_chars(self, text): return text.lower().replace(" ", "_").replace("'", "").replace("(", "_").replace(")", "_").replace("&", "and").replace(".", "").replace(",", "") def parse_route_by_prolog_name(self, prolog_name): return self.london_stations_df[self.london_stations_df['prolog_name'] == prolog_name] def get_line_color_for_prolog_name(self, name1, name2): station1 = self.parse_route_by_prolog_name(name1) station2 = self.parse_route_by_prolog_name(name2) connection = self.london_connections_df[(self.london_connections_df['station1'] == station1['id'].values[0]) & ( self.london_connections_df['station2'] == station2['id'].values[0])] if len(connection) == 0: connection = self.london_connections_df[(self.london_connections_df['station1'] == station2['id'].values[0]) & ( self.london_connections_df['station2'] == station1['id'].values[0])] line = self.london_lines_df[self.london_lines_df['line'] == connection['line'].values[0]] color = "#"+line['colour'].values[0] return color def get_line_name_for_prolog_name(self, name1, name2): station1 = self.parse_route_by_prolog_name(name1) station2 = self.parse_route_by_prolog_name(name2) connection = self.london_connections_df[(self.london_connections_df['station1'] == station1['id'].values[0]) & ( self.london_connections_df['station2'] == station2['id'].values[0])] if len(connection) == 0: connection = self.london_connections_df[(self.london_connections_df['station1'] == station2['id'].values[0]) & ( self.london_connections_df['station2'] == station1['id'].values[0])] line = self.london_lines_df[self.london_lines_df['line'] == connection['line'].values[0]] return line['name'].values[0] def get_traces(self): lines = [] for station1, station2, line, time in self.london_connections_df.values.tolist(): station1_lat = self.london_stations_df[self.london_stations_df['id'] == station1]['latitude'].values[0] station1_long = self.london_stations_df[self.london_stations_df['id'] == station1]['longitude'].values[0] station2_lat = self.london_stations_df[self.london_stations_df['id'] == station2]['latitude'].values[0] station2_long = self.london_stations_df[self.london_stations_df['id'] == station2]['longitude'].values[0] name = self.london_lines_df[self.london_lines_df['line'] == line]['name'].values[0] color = self.london_lines_df[self.london_lines_df['line'] == line]['colour'].values[0] lines.append(go.Scattermapbox( mode="lines", lat=[station1_lat, station2_lat], lon=[station1_long, station2_long], line={'width': 4, 'color': "#"+color}, name=name)) return lines def filter_duplicate_traces(self, fig): names = set() fig.for_each_trace(lambda trace: trace.update(showlegend=False) if (trace.name in names) else names.add(trace.name)) return fig def get_map_center(self): return self.london_stations_df['latitude'].mean(), self.london_stations_df['longitude'].mean() def get_default_map_layout(self, center_lat, center_lon): return go.Layout( mapbox_style="carto-positron", mapbox=dict( bearing=0, pitch=0, center=dict(lat=center_lat, lon=center_lon), zoom=11.8 )) def get_markers(self): return go.Scattermapbox( lat=self.london_stations_df['latitude'].tolist(), lon=self.london_stations_df['longitude'].tolist(), text=self.london_stations_df['name'].tolist(), name="Stations", mode='markers+text', textposition="bottom center", marker=go.scattermapbox.Marker(size=10) ) def filter_marker_and_traces(self, route): route_station = [self.parse_route_by_prolog_name(x) for x in route] lats = [x['latitude'].values[0] for x in route_station] longs = [x['longitude'].values[0] for x in route_station] names = [x['name'].values[0] for x in route_station] traces = [go.Scattermapbox( text=[names[i], names[i+1]], mode="markers+text+lines", lat=[lats[i], lats[i+1]], lon=[longs[i], longs[i+1]], line={'width': 4, 'color': self.get_line_color_for_prolog_name( route[i], route[i+1])}, name=self.get_line_name_for_prolog_name(route[i], route[i+1]), marker=go.scattermapbox.Marker(size=10) ) for i in range(len(route_station)-1)] traces.append(go.Scattermapbox( mode="markers", lat=[lats[0], lats[-1]], lon=[longs[0], longs[-1]], marker=go.scattermapbox.Marker( size=[10,15], color=["blue", "red"]), name="Stations" )) center_lat = np.array(lats).mean() center_lon = np.array(longs).mean() output_text = " -> ".join(names) return traces, center_lat, center_lon, output_text