xref: /petsc/share/petsc/bin/dmnetwork_view.py (revision 97c047f8306e861d004fa98651e63d4b3bca0606)
1#!/usr/bin/env python3
2import math
3
4def clamp(x, maxval, minval):
5	if x > maxval:
6		return maxval
7	elif x < minval:
8		return minval
9	else:
10		return x
11
12def parseIndexSet(text):
13	if text is None:
14		return None
15	indexSet = []
16	ranges = str(text).split(',')
17	for rangeStr in ranges:
18		if '-' in rangeStr != -1:
19			limits = rangeStr.split('-')
20			for rank in range(int(limits[0]), int(limits[1])+1):
21				if not rank in indexSet:
22					indexSet.append(rank)
23		else:
24			rank = int(rangeStr)
25			if not rank in indexSet:
26				indexSet.append(rank)
27	indexSet.sort()
28	return indexSet
29
30global nextWindowOffset
31nextWindowOffset = 50
32
33def parseRGBColor(color, defaultValue = (0.2, 0.2, 0.2, 1)):
34	# A '#' denotes the use of HTML colors
35	if isinstance(color, str) and color[0] == '#':
36		rgb = tuple(int(color[i:i+2], 16) for i in (1, 3, 5))
37		return (rgb[0] / 255, rgb[1] / 255, rgb[2] / 255, 1)
38	else:
39		return defaultValue
40
41class ColorParser:
42	def __init__(self, datasets):
43		self.colormap = plt.get_cmap('brg')
44
45		hasNormRange = False
46		normMaxVal = None
47		normMinVal = None
48		# Preprocess the datasets
49		for data in datasets:
50			for i,row in data.iterrows():
51				# For each node entry try to parse the color
52				if row['Type'] == 'Node':
53					color = row['Color']
54					try:
55						# If the color gets parsed as a float correctly use it to find the normalization range
56						color = float(color)
57						if not math.isnan(color):
58							hasNormRange = True
59							if normMaxVal is None or color > normMaxVal:
60								normMaxVal = color
61							if normMinVal is None or color < normMinVal:
62								normMinVal = color
63					except:
64						pass
65
66		if hasNormRange:
67			self.hasColorBar = True
68
69			# Clamp the minimum between negative infinity and zero and the corresponding for maximum
70			if normMinVal >= 0:
71				normMinVal = None
72			if normMaxVal <= 0:
73				normMaxVal = None
74
75			class BiasNorm(matplotlib.colors.Normalize):
76				def __init__(self, vmin, vmax, minbias = 0, maxbias = 0):
77					matplotlib.colors.Normalize.__init__(self, vmin, vmax)
78					self.minbias = minbias
79					self.maxbias = maxbias
80
81				def __call__(self, value, clip=None):
82					value = super(BiasNorm, self).__call__(value, clip)
83					value *= (1 - self.maxbias) - self.minbias
84					return value + self.minbias
85
86			# Four cases; Both min and max defined, one but not the either, or all values are uniform
87			if normMinVal is not None and normMaxVal is not None:
88				self.norm = matplotlib.colors.TwoSlopeNorm(0, normMinVal, normMaxVal)
89			elif normMinVal is not None:
90				self.norm = BiasNorm(normMinVal, 0, maxbias=0.5)
91			elif normMaxVal is not None:
92				self.norm = BiasNorm(0, normMaxVal, minbias=0.5)
93			else:
94				self.norm = None
95
96			self.normMinVal = normMinVal
97			self.normMaxVal = normMaxVal
98		else:
99			self.hasColorBar = False
100
101	# Parses a color value into an RGBA tuple to use with matplotlib
102	def getRGB(self, color, defaultValue = (0.2, 0.2, 0.2, 1)):
103		# Try to do the basic color parsing first
104		explicitColor = parseRGBColor(color, None)
105		if explicitColor is not None:
106			return explicitColor
107
108		# Try to parse as a number
109		try:
110			color = float(color)
111
112			# Normalize to [0,1], either by converting from [-1,1] or using a generated normalization range
113			if self.norm is not None:
114				color = self.norm(color)
115			else:
116				color = (color + 1) * 0.5
117
118			# Get the color from the color map
119			return self.colormap(color)
120		except:
121			return defaultValue
122
123
124class DisplayOptions:
125	def __init__(self,args,datasets):
126		# Parse any set node or edge colors
127		self.nodeColor = None
128		self.edgeColor = None
129		self.nodeTitleColor = None
130		self.edgeTitleColor = None
131
132		if 'set_node_color' in args:
133			self.nodeColor = parseRGBColor(args.set_node_color, None)
134		if 'set_edge_color' in args:
135			self.edgeColor = parseRGBColor(args.set_edge_color, None)
136
137		if 'set_node_title_color' in args:
138			self.nodeTitleColor = parseRGBColor(args.set_node_title_color, (1, 1, 1, 1))
139		if 'set_edge_title_color' in args:
140			self.edgeTitleColor = parseRGBColor(args.set_edge_title_color)
141
142		self.noNodes = 'no_nodes' in args and args.no_nodes
143		self.setTitle = args.set_title if 'set_title' in args else None
144		self.noNodeLabels = 'no_node_labels' in args and args.no_node_labels
145		self.noEdgeLabels = 'no_edge_labels' in args and args.no_edge_labels
146
147		self.nodeColorParser = ColorParser(datasets)
148
149		self.viewportShowVertices = parseIndexSet(args.viewport_show_vertices)
150		self.viewport = None
151		self.viewportPadding = float(args.viewport_padding) if args.viewport_padding else None
152
153	def adjustViewport(self, node):
154		# Only adjust if we are focusing on a set of vertices
155		if self.viewportShowVertices is not None and int(node.id) in self.viewportShowVertices:
156			x = node.position[0]
157			y = node.position[1]
158			pad = self.viewportPadding
159			# If no viewport is defined yet, set it directly
160			if self.viewport is None:
161				self.viewport = (x, x, y, y)
162			# Else compute by the minimum and maximum bounds
163			else:
164				self.viewport = (
165					min(self.viewport[0], x),
166					max(self.viewport[1], x),
167					min(self.viewport[2], y),
168					max(self.viewport[3], y)
169				)
170
171	def finalizeViewport(self):
172		if self.viewport is None:
173			return None
174		w = self.viewport[1] - self.viewport[0]
175		h = self.viewport[3] - self.viewport[2]
176		pad = self.viewportPadding or (max(w, h) * 0.1)
177		return (
178			self.viewport[0] - pad,
179			self.viewport[1] + pad,
180			self.viewport[2] - pad,
181			self.viewport[3] + pad
182		)
183
184# Class for holding the properties of a node
185class Node:
186	def __init__(self, row, opts: DisplayOptions):
187		# Set our ID and rank
188		self.id = row['ID']
189		self.rank = int(row['Rank'])
190
191		# Set our position
192		x = float(row['X'])
193		if np.isnan(x):
194			x = 0
195		y = float(row['Y'])
196		if np.isnan(y):
197			y = 0
198		z = float(row['Z'])
199		if np.isnan(z):
200			z = 0
201		self.position = (x,y,z)
202
203		# Set name and color, defaulting to a None name if not specified
204		self.name = row['Name']
205
206		self.color = opts.nodeColor or opts.nodeColorParser.getRGB(row['Color'])
207
208		# Adjust the viewport for the display options
209		opts.adjustViewport(self)
210
211# Class for holding the properties of an edge
212class Edge:
213	def __init__(self, row, opts: DisplayOptions, nodes):
214		# Set our ID and rank
215		self.id = row['ID']
216		self.rank = int(row['Rank'])
217
218		# Determine our starting and ending nodes from the X and Y properties
219		start = row['X']
220		if not start in nodes:
221			raise KeyError("No such node \'" + str(start) + "\' for start of edge \'" + str(self.id) + '\'')
222		self.startNode = nodes[start]
223		end = row['Y']
224		if not end in nodes:
225			raise KeyError ("No such node \'" + str(end) + "\' for end of edge \'" + str(self.id) + '\'')
226		self.endNode = nodes[end]
227
228		# Set name and color, defaulting to a None name if not specified
229		self.name = row['Name']
230
231		self.color = opts.edgeColor or parseRGBColor(row['Color'], (0.5, 0.5, 0.5, 1))
232
233
234# Class for holding the data for a rank
235class Rank:
236	def __init__(self, index):
237		self.id = index
238		self.nodes = {}
239		self.edges = {}
240
241	def display(self, opts: DisplayOptions, title):
242		# Create Numpy arrays for node and edge positions and colors
243		nodePositions = np.zeros((len(self.nodes), 2))
244		nodeColors = np.zeros((len(self.nodes), 4))
245		edgeSegments = np.zeros((len(self.edges), 2, 2))
246		edgeColors = np.zeros((len(self.edges), 4))
247
248		# Copy node positions and colors to the arrays
249		i = 0
250		for node in self.nodes.values():
251			nodePositions[i] = node.position[0], node.position[1]
252			nodeColors[i] = node.color
253			i += 1
254
255		# Copy edge positions and colors to the arrays
256		i = 0
257		for edge in self.edges.values():
258			start = edge.startNode.position
259			end = edge.endNode.position
260			edgeSegments[i] = [
261				(start[0], start[1]),
262				(end[0], end[1])
263			]
264			edgeColors[i] = edge.color
265			i += 1
266
267		# Start the figure for this rank
268		fig = plt.figure("Rank " + str(self.id) if self.id >= 0 else "Global")
269		try:
270			global nextWindowOffset
271			offset = nextWindowOffset
272			nextWindowOffset += 50
273
274			window = fig.canvas.manager.window
275			backend = matplotlib.get_backend()
276			if backend == 'TkAgg':
277				window.wm_geometry("+%d+%d" % (offset, offset))
278			elif backend == 'WXAgg':
279				window.SetPosition(offset, offset)
280			else:
281				window.move(offset, offset)
282		except Exception:
283			pass
284		# Get axis for the plot
285		axis = fig.add_subplot()
286
287		# Set the title of the plot if specified
288		if opts.setTitle:
289			title = (opts.setTitle, (0, 0, 0, 1))
290		if title is None:
291			title = ("Network", (0, 0, 0, 1))
292		if self.id != -1:
293			title = (title[0] + " (Rank " + str(self.id) + ")", title[1])
294		axis.set_title(title[0], color=title[1])
295
296		# Add a line collection to the axis for the edges
297		axis.add_collection(LineCollection(
298			segments=edgeSegments,
299			colors=edgeColors,
300			linewidths=2
301		))
302
303		if not opts.noNodes:
304			# Add a circle collection to the axis for the nodes
305			axis.add_collection(CircleCollection(
306				sizes=np.ones(len(self.nodes)) * (20 ** 2),
307				offsets=nodePositions,
308				transOffset=axis.transData,
309				facecolors=nodeColors,
310				# Place above the lines
311				zorder=3
312			))
313
314		if not opts.noNodeLabels and not opts.noNodes:
315			# For each node, plot its name at the center of its point
316			for node in self.nodes.values():
317				if node.name is not None:
318					axis.text(
319						x=node.position[0], y=node.position[1],
320						s=node.name,
321						# Center text vertically and horizontally
322						va='center', ha='center',
323						# Make sure the text is clipped within the plot area
324						clip_on=True,
325						color=opts.nodeTitleColor
326					)
327
328		if not opts.noEdgeLabels:
329			# For each edge, plot its name at the center of the line segment
330			for edge in self.edges.values():
331				if edge.name is not None:
332					axis.text(
333						x=(edge.startNode.position[0]+edge.endNode.position[0])/2,
334						y=(edge.startNode.position[1]+edge.endNode.position[1])/2,
335						s=edge.name,
336						va='center', ha='center',
337						clip_on=True,
338						color=opts.edgeTitleColor
339					)
340
341		# Scale the plot to the content
342		axis.autoscale()
343
344		# Adjust the viewport if requested
345		if opts.viewportShowVertices is not None:
346			viewport = opts.finalizeViewport()
347			if viewport:
348				plt.xlim(viewport[0], viewport[1])
349				plt.ylim(viewport[2], viewport[3])
350
351		# Draw the colorbar if allowed by options
352		colors = opts.nodeColorParser
353		if colors.hasColorBar:
354			norm = colors.norm
355			ylow = colors.normMinVal or 0
356			yhigh = colors.normMaxVal or 0
357			# If the actual range is zero
358			if ylow == yhigh:
359				# Clamp to the correct normalized extreme based on positive, negative, or zero
360				if ylow > 0:
361					ylow = yhigh = 1
362				elif ylow < 0:
363					ylow = yhigh = 0
364				else:
365					ylow = yhigh = 0.5
366				# Shift the low and high limits by an epsilon to make sure the
367				epsilon = ylow * 0.00001
368				cbar = plt.colorbar(matplotlib.cm.ScalarMappable(norm=norm, cmap=colors.colormap), ax=axis)
369				cbar.ax.set_ylim([ylow - epsilon, yhigh + epsilon])
370			else:
371				cbar = plt.colorbar(matplotlib.cm.ScalarMappable(norm=norm, cmap=colors.colormap), ax=axis)
372				cbar.ax.set_ylim([ylow, yhigh])
373
374def main(args):
375	datasets = []
376
377	# Read each file passed in arguments
378	for filename in args.filenames:
379		try:
380			# Read the data from the supplied CSV file
381			data = pd.read_csv(filename, skipinitialspace=True, dtype=str)
382			# Append to the list of datasets
383			datasets.append(data)
384		except Exception as e:
385			print("Warning! Could not read file \"" + filename + "\": " + str(e))
386			traceback.print_exc(file=sys.stdout)
387			exit(-1)
388
389
390	# Parse display options from arguments and datasets
391	opts = DisplayOptions(args, datasets)
392
393
394	# Variable storing a title to use or None
395	title = None
396
397	# The set of ranks
398	ranks = { -1: Rank(-1) }
399	#The global rank assigned to index -1
400	globalRank = ranks[-1]
401
402	# Function to get the specified rank
403	def getRank(rank: int):
404		if rank in ranks:
405			return ranks[rank]
406		else:
407			r = Rank(rank)
408			ranks[rank] = r
409			return r
410
411	# Now process the actual data
412	for data in datasets:
413		# Iterate each row of data in the file
414		for i,row in data.iterrows():
415			# Switch based on the type of the entry
416			type = row['Type']
417			if type == 'Type':
418				# If we encounter 'Type' again it is a duplicate header and should be skipped
419				continue
420			elif type == 'Title':
421				# Set the title based on name and color
422				titleColor = parseRGBColor(row['Color'])
423				title = (row['Name'], titleColor)
424			elif type == 'Node':
425				# Register a new node
426				node = Node(row, opts)
427				globalRank.nodes[node.id] = node
428				r = getRank(node.rank)
429				if r is not None:
430					r.nodes[node.id] = node
431			elif type == 'Edge':
432				# Register a new edge
433				edge = Edge(row, opts, globalRank.nodes)
434				globalRank.edges[edge.id] = edge
435				r = getRank(node.rank)
436				if r is not None:
437					edge = Edge(row, opts, r.nodes)
438					r.edges[edge.id] = edge
439
440	# Show the plot
441	if not args.no_display:
442		# Generate figures using ranks
443		if not args.no_combined_plot:
444			globalRank.display(opts, title)
445		if args.draw_rank_range:
446			ranges = parseIndexSet(args.draw_rank_range)
447			for rank in ranges:
448				ranks[rank].display(opts, title)
449		elif args.draw_all_ranks:
450			for rank in ranks:
451				if rank != -1:
452					ranks[rank].display(opts, title)
453
454		# Delay based on options
455		if args.display_time is not None:
456			plt.show(block=False)
457			plt.pause(float(args.display_time))
458			# Try to bring the window to front since we are displaying for a limited time
459			try:
460				window = plt.get_current_fig_manager().window
461				window.activateWindow()
462				window.raise_()
463			except AttributeError:
464				pass
465		else:
466			plt.show()
467		plt.close()
468
469
470if __name__ == "__main__":
471	try:
472		from argparse import ArgumentParser
473		# Construct the argument parse and parse the program arguments
474		argparser = ArgumentParser(
475			prog='dmnetwork_view.py',
476			description="Displays a CSV file generated from a DMNetwork using matplotlib"
477		)
478		argparser.add_argument('filenames', nargs='+')
479		argparser.add_argument('-t', '--set-title', metavar='TITLE', action='store', help="Sets the title for the generated plot, overriding any title set in the source file")
480		argparser.add_argument('-nnl', '--no-node-labels', action='store_true', help="Disables labeling nodes in the generated plot")
481		argparser.add_argument('-nel', '--no-edge-labels', action='store_true', help="Disables labeling edges in the generated plot")
482		argparser.add_argument('-nc', '--set-node-color', metavar='COLOR', action='store', help="Sets the color for drawn nodes, overriding any per-node colors")
483		argparser.add_argument('-ec', '--set-edge-color', metavar='COLOR', action='store', help="Sets the color for drawn edges, overriding any per-edge colors")
484		argparser.add_argument('-ntc', '--set-node-title-color', metavar='COLOR', action='store', help="Sets the color for drawn node titles, overriding any per-node colors")
485		argparser.add_argument('-etc', '--set-edge-title-color', metavar='COLOR', action='store', help="Sets the color for drawn edge titles, overriding any per-edge colors")
486		argparser.add_argument('-nd', '--no-display', action='store_true', help="Disables displaying the figure, but will parse as normal")
487		argparser.add_argument('-tx', '--test-execute', action='store_true', help="Returns from the program immediately, used only to test run the script")
488		argparser.add_argument('-dt', '--display-time', metavar='SECONDS', action='store', help="Sets the time to display the figure in seconds before automatically closing the window")
489		argparser.add_argument('-dar', '--draw-all-ranks', action='store_true', help="Draws each rank's network in a separate figure")
490		argparser.add_argument('-ncp', '--no-combined-plot', action='store_true', help="Disables drawing the combined network figure")
491		argparser.add_argument('-drr', '--draw-rank-range', action='store', metavar='RANGE', help="Specifies a comma-separated list of rank numbers or ranges to display, eg. \'1,3,5-9\'")
492		argparser.add_argument('-nn', '--no-nodes', action='store_true', help="Disables displaying the nodes")
493		argparser.add_argument('-vsv', '--viewport-show-vertices', action='store', metavar='RANGE', help="Sets the range of vertices to focus the viewport on, eg. \'1,3,5-9\'")
494		argparser.add_argument('-vp', '--viewport-padding', metavar='PADDING', action='store', help="Sets the padding in coordinate units to apply around the edges when setting the viewport")
495		args = argparser.parse_args()
496
497		if not args.test_execute:
498			import pandas as pd
499			import numpy as np
500			import matplotlib
501			import matplotlib.pyplot as plt
502			from matplotlib.collections import CircleCollection, LineCollection
503			import traceback
504			import sys
505
506			main(args)
507	except ImportError as error:
508		print("Missing import: " + str(error))
509		exit(-1)
510
511
512
513
514
515