-
Notifications
You must be signed in to change notification settings - Fork 5
/
plots.py
155 lines (136 loc) · 5.62 KB
/
plots.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
"""Script to create Hub & Spoke plots.
Obtained and modified the code of:
https://github.com/ekhoda/plotly_hub_and_spoke
"""
import plotly.graph_objects as go
def plot_network(data, add_path=True, auto_open_map=False, title='', open_map_in_cell=False):
"""Plot a hub-and-spoke network."""
# Prepare the data
df_dict = create_location_dataframe(data)
locations = set_locations_map(df_dict)
lat_range, lon_range = get_lat_lon_range(df_dict)
scope = get_scope(lat_range, lon_range)
fig = go.Figure()
# Specify the layout attributes
title = title or f'{scope.upper()} Network'
layout = dict(title=title,
showlegend=True,
geo=dict(
scope=scope,
showland=True,
# # if you want to zoom only on the area with lat/lon rather than the whole scope
# lataxis=dict(range=lat_range),
# lonaxis=dict(range=lon_range),
landcolor='rgb(243, 243, 243)',
countrycolor='rgb(204, 204, 204)'))
fig.update_layout(layout)
# Add locations
for d in locations:
fig.add_trace(d)
# Add paths
if add_path:
for i, row in data.iterrows():
fig.add_trace(go.Scattergeo(
lon=[row['Origin Lon'], row['Dest Lon']],
lat=[row['Origin Lat'], row['Dest Lat']],
mode='lines',
line=dict(width=2, color='green'),
opacity=0.8,
showlegend=False,
))
if open_map_in_cell:
fig.show() # to show it in Jupyter notebook
# Save the map as an HTML in the current directory and whether to automatically open it or not
fig.write_html(file=f'{title}.html', auto_open=auto_open_map)
def create_location_dataframe(data):
"""Return a dataframe with the proper attributes needed for the plot."""
origin_df = data[['Origin', 'Origin Lat', 'Origin Lon']].drop_duplicates()
origin_df['text'] = origin_df['Origin']
origin_df['size'] = 20
origin_df['color'] = 'blue'
origin_df['shape'] = 'triangle-up'
origin_df['name'] = origin_df['Origin']
origin_df.rename(columns={'Origin Lat': 'lat', 'Origin Lon': 'lon'}, inplace=True)
dest_df = data[['Destination', 'Dest Lat', 'Dest Lon']].drop_duplicates()
dest_df['text'] = dest_df['Destination']
dest_df['size'] = 8
dest_df['color'] = 'orange'
dest_df['shape'] = 'circle'
dest_df['name'] = dest_df['Destination']
dest_df.rename(columns={'Dest Lat': 'lat', 'Dest Lon': 'lon'}, inplace=True)
# we'll concat origin_df with dest_df
origin_df.drop('Origin', axis=1, inplace=True)
dest_df.drop('Destination', axis=1, inplace=True)
df_dict = {'Plant': origin_df, 'Customer': dest_df}
return df_dict
def set_locations_map(df_dict):
"""Create a list of dictionaries that hold location attributes needed for the plot."""
locations = []
for name, df in df_dict.items():
try:
locations.append(
dict(type='scattergeo',
lon=df['lon'],
lat=df['lat'],
text=df['text'],
name=name,
marker=dict(size=df['size'],
symbol=df['shape'],
color=df['color'],
line=dict(width=3, color='rgba(68, 68, 68, 0)'))))
except KeyError:
pass
return locations
def get_lat_lon_range(df_dict):
"""Return the range of lat and lon in the data."""
min_lat = +90
max_lat = -90
min_lon = +180
max_lon = -180
for df in df_dict.values():
try:
df_lat_min = df['lat'].min()
df_lat_max = df['lat'].max()
df_lon_min = df['lon'].min()
df_lon_max = df['lon'].max()
if df_lat_min < min_lat:
min_lat = df_lat_min
if df_lat_max > max_lat:
max_lat = df_lat_max
if df_lon_min < min_lon:
min_lon = df_lon_min
if df_lon_max > max_lon:
max_lon = df_lon_max
except KeyError:
pass
return [min_lat, max_lat], [min_lon, max_lon]
def get_scope(lat_range, lon_range):
"""Assign the proper scope based on range of data's lat/lon.
The geographic scope in Plotly can be one of:
"world" | "usa" | "europe" | "asia" | "africa" | "north america" | "south america".
(source: https://plot.ly/python/reference/#layout-geo-scope)
We're interested in USA, North America, South America, and Europe.
"""
us_lat_rng = [24, 55]
us_lon_rng = [-127, -50]
na_lat_rng = [15, 85]
na_lon_rng = [-170, -50]
eu_lat_rng = [30, 80]
eu_lon_rng = [-20, 70]
sa_lat_rng = [-60, 12]
sa_lon_rng = [-81, -34]
if (lat_range[0] >= us_lat_rng[0] and lat_range[1] <= us_lat_rng[1]
and lon_range[0] >= us_lon_rng[0] and lon_range[1] <= us_lon_rng[1]):
scope = 'usa'
elif (lat_range[0] >= na_lat_rng[0] and lat_range[1] <= na_lat_rng[1]
and lon_range[0] >= na_lon_rng[0] and lon_range[1] <= na_lon_rng[1]):
scope = 'north america'
elif (lat_range[0] >= eu_lat_rng[0] and lat_range[1] <= eu_lat_rng[1]
and lon_range[0] >= eu_lon_rng[0] and lon_range[1] <= eu_lon_rng[1]):
scope = 'europe'
elif (lat_range[0] >= sa_lat_rng[0] and lat_range[1] <= sa_lat_rng[1]
and lon_range[0] >= sa_lon_rng[0] and lon_range[1] <= sa_lon_rng[1]):
scope = 'south america'
else: # can add asia and africa
scope = 'world' # default
return scope