Better options to visu script

This commit is contained in:
Lucas Verney 2017-08-31 14:42:20 +02:00
parent a0b379bfcd
commit 203beebe37

38
visu.py
View File

@ -12,11 +12,12 @@ import logging
import os import os
import pickle import pickle
import sqlite3 import sqlite3
import sys
import matplotlib import matplotlib
matplotlib.use('AGG') # Use non-interactive backend matplotlib.use('AGG') # Use non-interactive backend
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
import progressbar import progressbar # progressbar2 pip module
import smopy import smopy
from scipy.spatial import Voronoi, voronoi_plot_2d from scipy.spatial import Voronoi, voronoi_plot_2d
@ -33,18 +34,28 @@ def get_hue(percentage):
return hue / 360.0 return hue / 360.0
if len(sys.argv) < 3:
sys.exit('Usage: %s db_file out_dir' % sys.argv[0])
db_file = sys.argv[1]
out_dir = sys.argv[2]
first_timestamp = None
if len(sys.argv) > 3:
first_timestamp = int(sys.argv[3])
progressbar.streams.wrap_stderr() progressbar.streams.wrap_stderr()
logging.basicConfig(level=logging.INFO) logging.basicConfig(level=logging.INFO)
# Ensure out folder exists # Ensure out folder exists
if not os.path.isdir('out'): if not os.path.isdir(out_dir):
logging.info('Creating out folder…') logging.info('Creating output folder %s', out_dir)
os.mkdir('out') os.makedirs(out_dir)
# Load all stations from the database # Load all stations from the database
logging.info('Loading all stations from the database…') logging.info('Loading all stations from the database…')
conn = sqlite3.connect("data.db") conn = sqlite3.connect(db_file)
c = conn.cursor() c = conn.cursor()
stations = c.execute( stations = c.execute(
"SELECT id, latitude, longitude, bike_stands, name FROM stations" "SELECT id, latitude, longitude, bike_stands, name FROM stations"
@ -100,9 +111,10 @@ for point_index, region_index in enumerate(vor.point_region):
"mpl_surface": None # Will store the drawn matplotlib surface (to update it easily) "mpl_surface": None # Will store the drawn matplotlib surface (to update it easily)
} }
# Dumping Voronoi diagram # Dumping Voronoi diagram
with open('out/voronoi.dat', 'wb') as fh: voronoi_file = os.path.join(out_dir, 'voronoi.dat')
with open(voronoi_file, 'wb') as fh:
pickle.dump(vor_regions, fh) pickle.dump(vor_regions, fh)
logging.info('Dumped Voronoi diagram to out/voronoi.dat') logging.info('Dumped Voronoi diagram to voronoi.dat')
# Plotting # Plotting
@ -135,6 +147,12 @@ for station_id, region in vor_regions.items():
# Get time steps # Get time steps
logging.info('Loading time steps from the database.') logging.info('Loading time steps from the database.')
if first_timestamp:
time_data = c.execute(
"SELECT DISTINCT updated FROM stationsstats WHERE updated > ? ORDER BY updated ASC",
(first_timestamp,)
)
else:
time_data = c.execute( time_data = c.execute(
"SELECT DISTINCT updated FROM stationsstats WHERE updated ORDER BY updated ASC" "SELECT DISTINCT updated FROM stationsstats WHERE updated ORDER BY updated ASC"
) )
@ -168,15 +186,15 @@ for t, in bar(time_data):
region = vor_regions[station_data[0]] region = vor_regions[station_data[0]]
region["mpl_surface"].set_color(matplotlib.colors.hsv_to_rgb([get_hue(percentage), 1.0, 1.0])) region["mpl_surface"].set_color(matplotlib.colors.hsv_to_rgb([get_hue(percentage), 1.0, 1.0]))
except KeyError: except KeyError:
logging.warn('Unknown Voronoi region for station %d.', station_data[0]) logging.debug('Unknown Voronoi region for station %d.', station_data[0])
# Output frame if necessary # Output frame if necessary
if t >= last_t + timesteps: if t >= last_t + timesteps:
ax.set_title(datetime.datetime.fromtimestamp(t // 1000).strftime('%d/%m/%Y %H:%M')) ax.set_title(datetime.datetime.fromtimestamp(t // 1000).strftime('%d/%m/%Y %H:%M'))
fig.tight_layout() fig.tight_layout()
fig.savefig('out/%d.png' % t) fig.savefig(os.path.join(out_dir, '%d.png' % t))
last_t = t last_t = t
# Output last frame # Output last frame
fig.tight_layout() fig.tight_layout()
fig.savefig('out/%d.png' % t) fig.savefig(os.path.join(out_dir, '%d.png' % t))