actlib_dataflow_neuro/test/prsim_plot.py

#!/usr/bin/python3

import sys
import glob
import re
import matplotlib.pyplot as plt
import numpy as np

def main(argv):

	# Load file
	assert len(argv) >= 2, "No arguments given. -h for help"

	if argv[1] == "-h":
		print("""Specify the path to prsim.out either by giving the full path,
or the folder name like 'buf_15'.
Use -include='regex' to specify signals to include (or -in).
Use -exclude='regex' to specify signals to exclude (or -ex).""")
		return


	file_path = argv[1]
	if not ".out" in file_path:
		file_path = f"./unit_tests/{file_path}/run/prsim.out"
	assert len(glob.glob(file_path)) >= 1, "prsim.out file not found!"
	print(f"Loading {file_path}")
	f = open(file_path,'r').read()

	# Start regexxing
	entries = re.findall(r"\t *(\d+) ([^:]+) : (\d)(  \[by.+\])?[\n\r]", f)
	assert len(entries) >= 1, "Could not find signal info in prsim.out!"

	# Check if user gave an include filter
	include_given = False
	include_re = None
	for arg in argv:
		r = re.findall(r'(-include|-in)=(.+)', arg)
		if len(r) >= 1:
			include_given = True
			include_re = r[0][1]

	# Check if user gave an exclude filter
	exclude_given = False
	exclude_re = None
	for arg in argv:
		r = re.findall(r'(-exclude|-ex)=(.+)', arg)
		if len(r) >= 1:
			exclude_given = True
			exclude_re = r[0][1]

	assert not (exclude_given and include_given), "Can't give include and exclude re simultaneously."
	if include_given: print(f"Including signals that match regex {include_re}")
	if exclude_given: print(f"Excluding signals that match regex {exclude_re}")

	if include_given:
		entries = [e for e in entries if not re.search(include_re, e[1]) == None]

	if exclude_given:
		entries = [e for e in entries if re.search(exclude_re, e[1]) == None]

	assert len(entries) >= 1, "No valid entries in prsim.out!"

	# Get list of all sigs and times
	times = np.array([int(e[0]) for e in entries])
	unique_times = np.unique(times)
	num_times = unique_times.shape[0]

	sigs = np.array([e[1] for e in entries])
	unique_sigs = np.unique(sigs)
	num_sigs = unique_sigs.shape[0]

	print(f"Plotting signals:")
	print(unique_sigs)


	# Some functions to order everything nicely
	# Should probably put these outside but whatever.
	def time_to_index(time):
		'''
		Since times are random, need to convert them to an index.
		'''
		out = np.argwhere(unique_times == int(time))
		return out[0][0]

	def sig_to_index(sig):
		'''
		Handles signal name ordering.
		Assume ordered like unique_sigs for now
		'''
		out = np.argwhere(unique_sigs == sig)
		return out[0][0]
	
	# Create matrix of signals over time and populate	
	signals_matrix = np.zeros((num_sigs, num_times), dtype = int)

	for sig in unique_sigs:
		entries_filtered = [e for e in entries if e[1] == sig]
		# make sure sorted
		entries_filtered = sorted(entries_filtered, key = lambda e: int(e[0]))
		for e in entries_filtered:
			val = int(e[2])
			val = 2*val -1
			signals_matrix[sig_to_index(sig),time_to_index(e[0]):] = val

	# Plot

	colour_undefined = (255,0,0)
	colour_high = (252, 186, 3)
	colour_low = (20, 184, 186)

	fig = plt.figure(figsize = (num_sigs/3,num_times/3))

	image = np.zeros((num_sigs, num_times, 3), dtype = int)
	image[signals_matrix == 0] = colour_undefined
	image[signals_matrix == 1] = colour_high
	image[signals_matrix == -1] = colour_low

	plt.imshow(image)
	ax = fig.gca()
	ax.set_xlabel("Time")
	# ax.set_ylabel("Signal")
	ax.set_yticks([])

	for sig in unique_sigs:
		ax.text(-1, sig_to_index(sig), sig, ha = "right", va = "center", size = 10)

	for i in range(num_sigs-1):
		ax.axhline(i+0.5, c = "white", lw = 2)

	for i in range(num_times-1):
		ax.axvline(i+0.5, c = "white", lw = 2)
		
	ax.axis("off")

	# Draw arrows
	for e in entries:
		# check if has a causal signal
		by = re.findall(r"\[by (.+):=(\d)",e[3])
		if len(by) == 0: continue
		sig = e[1]
		time = e[0]
		t_index = time_to_index(time)
		by_sig = by[0][0]
		by_val = int(by[0][1])
		t0,t1 = (t_index, t_index)

		# The sig that caused the change might have been excluded from plotting
		s0 = sig_to_index(sig)
		if by_sig in unique_sigs:
			s1 = sig_to_index(by_sig)
		else:
			s1 = s0


		arrow_c = "black" if by_val == 1 else "grey"
		if by_val == 1:
			plt.arrow(t0, s1, 0, s0-s1 + 0.2*np.sign(s0-s1), head_width = 0.5, width = 0.2, ec = "none", lw = 0, fc = "black", length_includes_head = True)
		else:
			plt.arrow(t0, s1, 0, s0-s1 + 0.2*np.sign(s0-s1), head_width = 0, width = 0.2, ec = "none", lw = 0, fc = "black", length_includes_head = True)
			plt.scatter((t0),(s0), c = "black", s = 30)
	
	file_out_path = file_path.replace(".out",".pdf")
	plt.savefig(file_out_path, bbox_inches = "tight")


if __name__ == "__main__":
	# print(sys.argv[0:])
	main(sys.argv)
   # main(sys.argv[1:])