Add viewport to visualizer

matsim_visualizer.py

@@ -10,11 +10,12 @@ import matplotlib.colors as colors
 
 
 class MatsimVisualizer():
-  def __init__(self, network_file_path, events_file_path, output_file_path):
+  def __init__(self, network_file_path, events_file_path, output_file_path, viewport):
-    self._nodes = None
+    self._nodes = {}
-    self._links = None
+    self._links = []
     self._pos = None
     self.norm = None
+    self._viewport = viewport
 
     self._output_file_path = output_file_path
     self._network_file_path = network_file_path
@@ -32,16 +33,24 @@ class MatsimVisualizer():
     with gzip.open(self._network_file_path, 'rb') as file:
       network_doc = etree.parse(file)
 
-    self._nodes = {node.get('id'): (float(node.get('x')), float(node.get('y')))
+    for node in network_doc.xpath('/network/nodes/node'):
-                   for node in network_doc.xpath('/network/nodes/node')}
+      x = float(node.get('x'))
+      y = float(node.get('y'))
+      if self._viewport[0] < x < self._viewport[2] and self._viewport[1] < y < self._viewport[3]:
+        self._nodes[node.get('id')] = (x, y)
+
+    for link in network_doc.xpath('/network/links/link'):
+      start_node = link.get('from')
+      end_node = link.get('to')
+      id = link.get('id')
+      if start_node in self._nodes and end_node in self._nodes:
+        self._links.append({
+          'id': id,
+          'from': start_node,
+          'to': end_node,
+          'vehicles': 0,
+        })
 
-    self._links = [{
-      'id': link.get('id'),
-      'from': link.get('from'),
-      'to': link.get('to'),
-      'vehicles': 0,
-    } for link in network_doc.xpath('/network/links/link')]
-    seen_links = set()
     cumulative_traffic = defaultdict(int)
 
     # ====== LOAD EVENTS ====== #
@@ -49,48 +58,29 @@ class MatsimVisualizer():
       events_doc = etree.parse(file)
 
     self._tick = 0
-    last_time = None
     for event in events_doc.xpath('/events/event'):
       link_id = event.get('link')
       event_type = event.get('type')
       time = float(event.get('time'))
-      vehicle_id = event.get('vehicle')
       ticked = False
 
       if link_id is not None and event_type is not None and time is not None:
         if event_type in ['entered link', 'vehicle enters traffic']:
           self._traffic_per_tick[self._tick][link_id] += 1
-          seen_links.add(link_id)
+          ticked = True
+
           cumulative_traffic[link_id] += 1
+          # We need to find the maximum value for traffic at any given point for the heatmap
           if self._max_traffic < cumulative_traffic[link_id]:
             self._max_traffic = cumulative_traffic[link_id]
-          ticked = True
         elif event_type in ['left link', 'vehicle leaves traffic']:
           self._traffic_per_tick[self._tick][link_id] -= 1
           cumulative_traffic[link_id] -= 1
           ticked = True
 
-        if ticked and last_time is not None and last_time != time:
+        if ticked:
           self._tick += 1
 
-        last_time = time
-
-    # ====== FILTER LINKS ====== #
-    unique_links = []
-    for link in self._links:
-      if link['id'] in seen_links:
-        unique_links.append(link)
-    self._links = unique_links
-
-    # ====== FILTER NODES ====== #
-    seen_nodes = set()
-    for link in self._links:
-      seen_nodes.add(link['from'])
-      seen_nodes.add(link['to'])
-
-    filtered_nodes = {node_id: self._nodes[node_id] for node_id in seen_nodes if node_id in self._nodes}
-    self._nodes = filtered_nodes
-
   def create_graph(self):
     for node_id, coords in self._nodes.items():
       self._G.add_node(node_id, pos=coords)
@@ -101,7 +91,7 @@ class MatsimVisualizer():
   def setup_color_mapping(self):
     self.norm = colors.Normalize(vmin=0, vmax=self._max_traffic)
 
-  def update(self, frame_number):
+  def update(self, frame):
     traffic_change = self._traffic_per_tick[self._tick]
 
     edge_colors = []
@@ -114,7 +104,7 @@ class MatsimVisualizer():
           break
 
       edge_colors.append(self._cmap(link['vehicles']))
-      edge_widths.append(1 + self.norm(link['vehicles'])*10)
+      edge_widths.append(1 + self.norm(link['vehicles']) * 10)
 
     plt.cla()
     nx.draw(self._G, self._pos, node_size=0, node_color='blue', width=edge_widths, edge_color=edge_colors,