Awesome graphs

Author: devExcale

src/graph.py

@@ -1,8 +1,9 @@
 import argparse
-from typing import Dict
+from typing import Dict, Literal
 
 import matplotlib.pyplot as plt
 import pandas
+from matplotlib.colors import TABLEAU_COLORS
 
 parser = argparse.ArgumentParser(
 	description="Generate a graph with the execution times of the MPI/OMP MinHash implementations."
@@ -27,13 +28,17 @@
 	help="Path to the directory where the PNG file will be saved (dot for current directory)",
 )
 
+colors = list(TABLEAU_COLORS.values())
+
 
 def draw_graph(
 		dists: Dict[str, pandas.DataFrame],
 		x_label: str,
 		y_label: str,
 		title: str = None,
-		save_path: str = None
+		save_path: str = None,
+		peakline: Literal['min', 'max', None] = None,
+		legend_loc: str = 'best',
 ):
 	"""
 	Draws a graph from the given data.
@@ -52,14 +57,56 @@ def draw_graph(
 	# Add Style
 	plt.style.use('classic')
 
-	# Plot the data
-	for name, df in dists.items():
-		if len(df) > 1:
+	# Iterate over non-empty distributions
+	_dists = filter(lambda t: len(t[1]) > 0, dists.items())
+
+	for i, (name, df) in enumerate(_dists):
+		len_df = len(df)
+
+		# Get dist color
+		color = colors[i % len(colors)]
+
+		if len_df > 1:
 			# Proper distribution, plot it
-			plt.plot(df[x_label], df[y_label], label=name, marker='o', markersize=3)
-		elif len(df) == 1:
-			# Single point, plot it as a line
-			plt.axline((0, df[y_label].values[0]), slope=0, label=name)
+			plt.plot(
+				df[x_label],
+				df[y_label],
+				label=name,
+				marker='o',
+				markersize=3,
+				color=color,
+			)
+
+		# Check if we need to plot a peakline
+		if peakline:
+			peakline = str(peakline).lower()
+
+			# Get peak value
+			if peakline == 'min':
+				y = df[y_label].min()
+			elif peakline == 'max':
+				y = df[y_label].max()
+			else:
+				raise ValueError(f"Invalid peakline value: {peakline}")
+
+			# Plot the mix/max line
+			if len_df == 1:
+				# Plot with label
+				plt.axline(
+					(0, y),
+					slope=0,
+					label=name,
+					color=color,
+					linestyle='dashed',
+				)
+			else:
+				# Plot without label
+				plt.axline(
+					(0, y),
+					slope=0,
+					color=color,
+					linestyle='dashed',
+				)
 
 	# Add labels
 	plt.title(title)
@@ -70,7 +117,7 @@ def draw_graph(
 	plt.gcf().set_size_inches(10, 5)
 
 	# Add legend
-	plt.legend()
+	plt.legend(loc=legend_loc)
 
 	# Add Grid
 	plt.grid()
@@ -113,23 +160,72 @@ def get_time_dataframes(csv_path: str) -> Dict[str, pandas.DataFrame]:
 	return dists
 
 
+def get_speedup_dataframes(time_dists: Dict[str, pandas.DataFrame]) -> Dict[str, pandas.DataFrame]:
+	"""
+	Computes the speedup for each library and returns the results as a dict {lib:DataFrame}.
+
+	Args:
+	time_df: DataFrame containing the time data.
+
+	Returns:
+	DataFrame containing the speedup for each library.
+	"""
+
+	# Copy the DataFrame
+	speedup_df = time_dists.copy()
+
+	# Get baseline time
+	if "NONE" in time_dists:
+		baseline_time = time_dists["NONE"]["time_elapsed"].item()
+	else:
+		baseline_time = None
+		for df in time_dists.values():
+			if len(df) <= 0:
+				continue
+			t = df["time_elapsed"].min()
+			if baseline_time is None or t < baseline_time:
+				baseline_time = t
+
+	# Compute the speedup
+	for lib, df in speedup_df.items():
+		df["speedup"] = df["time_elapsed"].map(lambda time: baseline_time / time)
+
+	# Return the DataFrame
+	return speedup_df
+
+
 def main():
 	# Parse the command-line arguments
 	args = parser.parse_args()
 
 	# Read the CSV files
-	data = get_time_dataframes(f"{args.in_csv_path}/time_{args.dataset}.csv")
+	exec_dists = get_time_dataframes(f"{args.in_csv_path}/time_{args.dataset}.csv")
+	speedup_dists = get_speedup_dataframes(exec_dists)
 
 	# Compute save path
-	save_path = f"{args.out_png_path}/time_{args.dataset}.svg" if args.out_png_path else None
+	path_exec = f"{args.out_png_path}/exec_{args.dataset}.svg" if args.out_png_path else None
+	path_speedup = f"{args.out_png_path}/speedup_{args.dataset}.svg" if args.out_png_path else None
 
-	# Draw a graph
+	# Draw execution graph
 	draw_graph(
-		dists=data,
+		dists=exec_dists,
 		x_label="n_processes",
 		y_label="time_elapsed",
 		title=f"Execution time ({args.dataset})",
-		save_path=save_path
+		save_path=path_exec,
+		peakline="min",
+		legend_loc='upper right',
+	)
+
+	# Draw speedup graph
+	draw_graph(
+		dists=speedup_dists,
+		x_label="n_processes",
+		y_label="speedup",
+		title=f"Speedup ({args.dataset})",
+		save_path=path_speedup,
+		peakline="max",
+		legend_loc='lower right',
 	)