Customized visualization class based on osmnx integrating smopy

routing_vis.py

import os
from urllib2 import URLError

import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import smopy
from matplotlib.collections import LineCollection


class GraphPlotter(object):
    def __init__(self, g, config):

        # data
        self._g = g
        self._points = np.array(nx.get_node_attributes(g, 'pos').values())

        # dimensions
        self._bbox = config['bbox']
        self._fig_width = config['fig_width']
        self._fig_height = config['fig_height']
        self._margin = config['margin']

        # node aesthetics:
        self._node_color = config['node_color']
        self._node_size = config['node_size']
        self._node_alpha = config['node_alpha']
        self._node_edge_color = config['node_edge_color']

        # edge aesthetics
        self._edge_widths = config['edge_widths']
        self._edge_color = config['edge_color']
        self._edge_alpha = config['edge_alpha']

        # route line aesthetics
        self._route_line_width = config['route_line_width']
        self._route_color = config['route_color']
        self._route_alpha = config['route_alpha']

        # route orig/dest aesthetics
        self._origin_point_color = config['origin_point_color']
        self._origin_node_color = config['origin_node_color']
        self._dest_point_color = config['dest_point_color']
        self._dest_node_color = config['dest_node_color']
        self._orig_dest_node_alpha = config['orig_dest_point_alpha']
        self._orig_dest_node_size = config['orig_dest_node_size']

        # save/show info
        self._axis = config['axis']
        self._image_folder = config['image_folder']

        # initialize the plot
        self._fig, self._ax = None, None
        self._init_plot()

    def _init_plot(self):
        if self._bbox is None:
            west = min(self._points[:, 0])
            east = max(self._points[:, 0])
            north = max(self._points[:, 1])
            south = min(self._points[:, 1])
            self._bbox = south, east, north, west
        else:
            south, east, north, west = self._bbox

        bbox_aspect_ratio = (north - south) / (east - west)

        if self._fig_width is None:
            self._fig_width = self._fig_height / bbox_aspect_ratio

        self._fig, self._ax = plt.subplots(figsize=(self._fig_width, self._fig_height))
        try:
            self.osm_map = smopy.Map((south, east, north, west), z=13)
        except URLError as err:
            print "URL error: {}, skipping map background...".format(err)
        self._ax = self.osm_map.show_mpl(self._ax)
        w, n = self.osm_map.to_pixels(north, west)
        e, s = self.osm_map.to_pixels(south, east)
        margin_ns = (n - s) * self._margin
        margin_ew = (e - w) * self._margin
        self._ax.set_ylim((s - margin_ns, n + margin_ns))
        self._ax.set_xlim((w - margin_ew, e + margin_ew))
        self._ax.get_xaxis().get_major_formatter().set_useOffset(False)
        self._ax.get_yaxis().get_major_formatter().set_useOffset(False)

        if not self._axis:
            self._ax.axis('off')

    def plot_graph(self, ec=None):

        node_Xs = self._points[:, 0]
        node_Ys = self._points[:, 1]

        lines = []

        for u, v in self._g.edges():
            x1 = self._g.node[u]['pos'][0]
            y1 = self._g.node[u]['pos'][1]
            x2 = self._g.node[v]['pos'][0]
            y2 = self._g.node[v]['pos'][1]
            p1 = self.osm_map.to_pixels(y1, x1)
            p2 = self.osm_map.to_pixels(y2, x2)
            line = [p1, p2]
            lines.append(line)

        lc = LineCollection(lines, colors=ec, linewidths=self._edge_widths, alpha=self._edge_alpha,
                            zorder=2)

        self._ax.add_collection(lc)
        self._ax.scatter(node_Xs, node_Ys, s=self._node_size, c=self._node_color, alpha=self._node_alpha,
                         edgecolor=self._node_edge_color,
                         zorder=1)

    def plot_graph_route(self, route, route_color='r', origin_point=None, destination_point=None):

        origin_node = route[0]
        destination_node = route[-1]

        if origin_point is None or destination_point is None:
            dest_coords = self.osm_map.to_pixels(list(reversed(self._g.node[destination_node]['pos'])))
            origin_coords = self.osm_map.to_pixels(list(reversed(self._g.node[origin_node]['pos'])))


        else:
            origin_coords = self.osm_map.to_pixels(origin_point[0], origin_point[1])
            dest_coords = self.osm_map.to_pixels(destination_point[0], destination_point[1])
            self._origin_node_color = self._origin_point_color
            self._dest_node_color = self._dest_point_color

        self._ax.scatter(origin_coords[0], origin_coords[1], s=self._orig_dest_node_size,
                         c=self._origin_node_color, alpha=self._orig_dest_node_alpha,
                         edgecolor=self._node_edge_color, zorder=4)

        self._ax.scatter(dest_coords[0], dest_coords[1], s=self._orig_dest_node_size,
                         c=self._dest_node_color, alpha=self._orig_dest_node_alpha,
                         edgecolor=self._node_edge_color, zorder=4)

        edge_nodes = list(zip(route[:-1], route[1:]))

        lines = []

        for u, v in edge_nodes:
            x1 = self._g.node[u]['pos'][0]
            y1 = self._g.node[u]['pos'][1]
            x2 = self._g.node[v]['pos'][0]
            y2 = self._g.node[v]['pos'][1]
            p1 = self.osm_map.to_pixels(y1, x1)
            p2 = self.osm_map.to_pixels(y2, x2)
            line = [p1, p2]

            lines.append(line)
        lc = LineCollection(lines, colors=route_color, linewidths=self._route_line_width, alpha=self._route_alpha,
                            zorder=3)
        self._ax.add_collection(lc)

    def save_and_show(self, show=True, save=True, filename='temp', file_format='png', dpi=150):
        if save:
            if not os.path.exists(self._image_folder):
                os.makedirs(self._image_folder)
            path_filename = '{}/{}.{}'.format(self._image_folder, filename, file_format)

            if file_format == 'svg':
                # if the file_format is svg, prep the fig/ax a bit for saving
                self._ax.axis('off')
                self._ax.set_position([0, 0, 1, 1])
                self._ax.patch.set_alpha(0.)
                self._fig.patch.set_alpha(0.)
                self._fig.savefig(path_filename, bbox_inches=0, transparent=True, format=file_format)
            else:
                self._fig.savefig(path_filename, dpi=dpi, bbox_inches='tight', format=file_format)

        if show:
            plt.show()


def make_default_config():
    config = {'bbox': None, 'fig_width': None, 'fig_height': 6, 'margin': 0.02, 'edge_color': '#8F8F8F',
              'edge_alpha': 0.75,
              'node_color': 'none', 'node_size': 15, 'node_alpha': 1,
              'node_edge_color': 'none', 'edge_widths': 1.5, 'route_line_width': 4,
              'route_color': 'r', 'route_alpha': 0.5, 'dest_point_color': 'g', 'dest_node_color': 'g',
              'origin_point_color': 'm',
              'origin_node_color': 'm',
              'orig_dest_point_alpha': 0.5, 'orig_dest_node_size': 100,
              'axis': False, 'image_folder': 'images'}
    return config