CERC/matsim-proto
19
0
Fork 0
Code Issues 1 Pull Requests Packages Projects Releases Wiki Activity

Add viewport to visualizer

Browse Source
This commit is contained in:
Ruben Sanchez 2024-02-16 11:47:55 -05:00
parent 50d7fbaca2
commit 48b6afe31d
1 changed files with 27 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

64
matsim_visualizer.py
View File

@ -10,11 +10,12 @@ import matplotlib.colors as colors
class MatsimVisualizer():
def __init__(self, network_file_path, events_file_path, output_file_path):
self._nodes = None
self._links = None
def __init__(self, network_file_path, events_file_path, output_file_path, viewport):
self._nodes = {}
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,