From 9f14f46631d1a989fdb030278f445c5a79c18dc9 Mon Sep 17 00:00:00 2001
From: reiniscirpons <rc234@st-andrews.ac.uk>
Date: Thu, 2 Nov 2023 01:32:52 +0000
Subject: [PATCH] Basic rewrite of benchmark plot generator

---
 etc/bench_plot.py | 109 ++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 91 insertions(+), 18 deletions(-)

diff --git a/etc/bench_plot.py b/etc/bench_plot.py
index f619bf42..d878a534 100755
--- a/etc/bench_plot.py
+++ b/etc/bench_plot.py
@@ -4,6 +4,7 @@
 import re
 import statistics as stats
 import sys
+from math import isqrt
 
 import matplotlib
 import numpy as np
@@ -50,8 +51,9 @@ def time_unit(Y):
         Y = [y / 1000 for y in Y]
     return time_units[index], Y
 
+
 def add_plot(xml_fnam, num_bars=4):
-    global color;
+    global color
     current_bar = 0
     Y = []
     Y_for_comparison = None
@@ -74,7 +76,7 @@ def add_plot(xml_fnam, num_bars=4):
             Y,
             1,
             align="center",
-            color=color[:len(Y)],
+            color=color[: len(Y)],
         )
         total_cols += len(Y) + 1
     plt.yscale("log", nonpositive="clip")
@@ -85,7 +87,6 @@ def add_plot(xml_fnam, num_bars=4):
     # print(Y)
     # width = 1
 
-
     # plt.axhline(
     #     stats.mean(Y),
     #     color=color[current_bar],
@@ -116,24 +117,96 @@ def add_plot(xml_fnam, num_bars=4):
     #     plt.ylabel("Time (relative)")
     #     plt.legend(loc="upper left")
 
-def check_filename(xml_fnam):
-    if len(xml_fnam.split(".")) < 2:
+
+def determine_subplot_layout(nr_plots: int) -> tuple[int, int]:
+    """Determine the number of rows and columns from number of plots."""
+    nr_plot_rows = isqrt(nr_plots)
+    nr_plot_cols = nr_plot_rows
+    if nr_plot_rows * nr_plot_cols < nr_plots:
+        nr_plot_cols += 1
+    while nr_plot_rows * nr_plot_cols < nr_plots:
+        nr_plot_rows += 1
+    return nr_plot_rows, nr_plot_cols
+
+
+def process_result(result_soup) -> tuple[str, float]:
+    """Extract data from a single xml result entry.
+
+    Returns
+    -------
+    result_name: str
+        The test case name
+    result_time: float
+        The test case time in nanoseconds
+    """
+    result_name = result_soup["name"]
+    if "name" not in result_soup.attrs:
+        raise ValueError(
+            f"Malformed benchmark file, result record does not contain 'name': {result_soup}"
+        )
+    result_mean_soup = result_soup.find("mean")
+    if result_mean_soup is None:
+        raise ValueError(
+            f"Malformed benchmark file, result record does not contain 'mean': {result_soup}"
+        )
+    if "value" not in result_mean_soup.attrs:
         raise ValueError(
-            f"expected filename of form x.xml found {xml_fnam}"
+            f"Malformed benchmark file, result 'mean' record does not contain 'value': {result_mean_soup}"
         )
+    result_time = float(result_mean_soup["value"]) / 1  # time in nanoseconds
+    return result_name, result_time
+
+
+def make_ax(ax, test_case_soup):
+    if "name" not in test_case_soup.attrs:
+        raise ValueError(
+            f"Malformed benchmark file, test_case record does not contain 'name': {test_case_soup}"
+        )
+    results = test_case_soup.find_all("BenchmarkResults")
+    result_names, result_times = zip(*map(process_result, reversed(results)))
+    bars = ax.barh(
+        result_names,
+        result_times,
+        align="center",
+        color=color[: len(result_times)],
+    )
+    test_name = test_case_soup["name"]
+    ax.set_title(f'Benchmark "{test_name}" runtime')
+    ax.set_xlabel(f"ns")
+    return ax
+
+
+def make_fig(benchmark_soup):
+    test_cases = benchmark_soup.find_all("TestCase")
+    nr_plots = len(test_cases)
+    nr_plot_rows, nr_plot_cols = determine_subplot_layout(nr_plots)
+    fig, axs = plt.subplots(nr_plot_rows, nr_plot_cols)
+    for test_case_soup, ax in zip(test_cases, axs.flat):
+        ax = make_ax(ax, test_case_soup)
+    return fig
+
+
+def check_filename(xml_fnam):
+    if len(xml_fnam.split(".")) < 2:
+        raise ValueError(f"expected filename of form x.xml found {xml_fnam}")
+
 
+if __name__ == "__main__":
+    args = sys.argv[1:]
 
-from sys import argv
+    for x in args:
+        check_filename(x)
+        # TODO more arg checks
 
-args = sys.argv[1:]
+    for x in args:
+        with open(x, "r") as in_file:
+            xml_text = in_file.read()
+        soup = BeautifulSoup(xml_text, "xml")
+        fig = make_fig(soup)
+        plt.show()
 
-for x in args:
-    check_filename(x)
-    # TODO more arg checks
-for x in args:
-    add_plot(x)
-xml_fnam = args[0]
-png_fnam = "".join(xml_fnam.split(".")[:-1]) + ".png"
-print("Writing {} . . .".format(png_fnam))
-plt.savefig(png_fnam, format="png", dpi=300)
-sys.exit(0)
+    xml_fnam = args[0]
+    png_fnam = "".join(xml_fnam.split(".")[:-1]) + ".png"
+    print("Writing {} . . .".format(png_fnam))
+    plt.savefig(png_fnam, format="png", dpi=300)
+    sys.exit(0)