-
Notifications
You must be signed in to change notification settings - Fork 7
/
plot_mpl.py
103 lines (88 loc) · 4.11 KB
/
plot_mpl.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
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import numpy as np
import pandas as pd
from mpl_toolkits.basemap import Basemap
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize, LinearSegmentedColormap, PowerNorm
def plot_map(in_filename, color_mode='screen',
out_filename='flights_map_mpl.png', absolute=False):
"""Plots the given CSV data files use matplotlib basemap and saves it to
a PNG file.
Args:
in_filename: Filename of the CSV containing the data points.
out_filename: Output image filename
color_mode: Use 'screen' if you intend to use the visualisation for
on screen display. Use 'print' to save the visualisation
with printer-friendly colors.
absolute: set to True if you want coloring to depend on your dataset
parameter value (ie for comparison).
When set to false, each coordinate pair gets a different
color.
"""
if color_mode == 'screen':
bg_color = (0.0, 0.0, 0, 1.0)
coast_color = (204/255.0, 0, 153/255.0, 0.7)
color_list = [(0.0, 0.0, 0.0, 0.0),
(204/255.0, 0, 153/255.0, 0.6),
(255/255.0, 204/255.0, 230/255.0, 1.0)]
else:
bg_color = (1.0, 1.0, 1.0, 1.0)
coast_color = (10.0/255.0, 10.0/255.0, 10/255.0, 0.8)
color_list = [(1.0, 1.0, 1.0, 0.0),
(255/255.0, 204/255.0, 230/255.0, 1.0),
(204/255.0, 0, 153/255.0, 0.6)
]
# define the expected CSV columns
CSV_COLS = ('dep_lat', 'dep_lon', 'arr_lat', 'arr_lon',
'nb_flights', 'CO2')
routes = pd.read_csv(in_filename, names=CSV_COLS, na_values=['\\N'],
sep=';', skiprows=1)
num_routes = len(routes.index)
# normalize the dataset for color scale
norm = PowerNorm(0.3, routes['nb_flights'].min(),
routes['nb_flights'].max())
# norm = Normalize(routes['nb_flights'].min(), routes['nb_flights'].max())
# create a linear color scale with enough colors
if absolute:
n = routes['nb_flights'].max()
else:
n = num_routes
cmap = LinearSegmentedColormap.from_list('cmap_flights', color_list,
N=n)
# create the map and draw country boundaries
plt.figure(figsize=(27, 20))
m = Basemap(projection='mill', lon_0=0)
m.drawcoastlines(color=coast_color, linewidth=1.0)
m.fillcontinents(color=bg_color, lake_color=bg_color)
m.drawmapboundary(fill_color=bg_color)
# plot each route with its color depending on the number of flights
for i, route in enumerate(routes.sort_values(by='nb_flights',
ascending=True).iterrows()):
route = route[1]
if absolute:
color = cmap(norm(int(route['nb_flights'])))
else:
color = cmap(i * 1.0 / num_routes)
line, = m.drawgreatcircle(route['dep_lon'], route['dep_lat'],
route['arr_lon'], route['arr_lat'],
linewidth=0.5, color=color)
# if the path wraps the image, basemap plots a nasty line connecting
# the points at the opposite border of the map.
# we thus detect path that are bigger than 30km and split them
# by adding a NaN
path = line.get_path()
cut_point, = np.where(np.abs(np.diff(path.vertices[:, 0])) > 30000e3)
if len(cut_point) > 0:
cut_point = cut_point[0]
vertices = np.concatenate([path.vertices[:cut_point, :],
[[np.nan, np.nan]],
path.vertices[cut_point+1:, :]])
path.codes = None # treat vertices as a serie of line segments
path.vertices = vertices
# save the map
plt.savefig(out_filename, format='png', bbox_inches='tight')
if __name__ == '__main__':
# use 'screen' color mode for on-screen display. Use 'print' if you intend
# to print the map
plot_map('data.csv', 'screen', absolute=False)