From 442fb2eaf332eb592a67cfdc10af84f197a861fc Mon Sep 17 00:00:00 2001
From: EdmundGoodman <egoodman3141@gmail.com>
Date: Mon, 26 Feb 2024 00:17:54 +0000
Subject: [PATCH] Add support for splitting line plots into series by metrics

---
 src/hpc_multibench/test_bench.py    | 18 +++++++++++++-----
 src/hpc_multibench/yaml_model.py    |  1 +
 yaml_examples/kudu/simple_rust.yaml | 12 +++++++++++-
 3 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/src/hpc_multibench/test_bench.py b/src/hpc_multibench/test_bench.py
index 2e8b7fa..04c76ab 100755
--- a/src/hpc_multibench/test_bench.py
+++ b/src/hpc_multibench/test_bench.py
@@ -253,15 +253,23 @@ def report(self) -> None:
         # Extract the outputs into the data format needed for the line plot
         for plot in self.bench_model.analysis.line_plots:
             # TODO: Could extraction function out into analysis file?
-            data: dict[str, list[tuple[float, float]]] = {
-                run_name: [] for run_name in self.run_configuration_models
-            }
+            data: dict[tuple[str, ...], list[tuple[float, float]]] = {}
+
             for run_configuration, output in run_outputs.values():
                 if output is not None:
                     metrics = self.extract_metrics(output)
                     if metrics is None:
                         continue
-                    data[run_configuration.name].append(
+
+                    split_data: list[str] = [
+                        f"{split_metric}={metrics[split_metric]}"
+                        for split_metric in plot.split_by
+                    ]
+                    series_name = (run_configuration.name, *split_data)
+
+                    if series_name not in data:
+                        data[series_name] = []
+                    data[series_name].append(
                         (
                             float(metrics[plot.x]),
                             float(metrics[plot.y]),
@@ -270,7 +278,7 @@ def report(self) -> None:
 
             for name, results in data.items():
                 print(name, results)
-                plt.plot(*zip(*results, strict=True), marker="x", label=name)
+                plt.plot(*zip(*results, strict=True), marker="x", label=",".join(name))
             plt.xlabel(plot.x)
             plt.ylabel(plot.y)
             plt.title(plot.title)
diff --git a/src/hpc_multibench/yaml_model.py b/src/hpc_multibench/yaml_model.py
index 7180f48..4f7cfaf 100755
--- a/src/hpc_multibench/yaml_model.py
+++ b/src/hpc_multibench/yaml_model.py
@@ -54,6 +54,7 @@ class LinePlotModel(BaseModel):
     title: str
     x: str
     y: str
+    split_by: list[str] = []
 
 
 # class RooflinePlotModel(BaseModel):
diff --git a/yaml_examples/kudu/simple_rust.yaml b/yaml_examples/kudu/simple_rust.yaml
index 3a061b8..826d71d 100644
--- a/yaml_examples/kudu/simple_rust.yaml
+++ b/yaml_examples/kudu/simple_rust.yaml
@@ -166,7 +166,6 @@ benches:
           y: "Wall time (s)"
 
   "mpi-config-sweep":
-    enabled: False
     run_configurations:
       - "cpp-mpi"
       - "rust-mpi"
@@ -190,12 +189,18 @@ benches:
         "Total flops": "FLOPS Summary:[\\s\\S]*Total\\s*: ([\\d\\.]+)[\\s\\S]*\nMFLOPS Summary"
         "Total mflops": "MFLOPS Summary:[\\s\\S]*Total\\s*: ([\\d\\.]+)"
         "Wall time (s)": "real\\s([\\d\\.]+)\nuser"
+        "Nodes": "=== RUN INSTANTIATION ===\n\\{.*sbatch_config: \\{.*nodes: (\\d+).*\\}"
+        "Tasks per Node": "=== RUN INSTANTIATION ===\n\\{.*sbatch_config: \\{.*ntasks-per-node: (\\d+).*\\}"
       line_plots:
         - title: "MPI Configuration Sweep"
           x: "Mesh x size"
           y: "Wall time (s)"
+          split_by:
+            - "Nodes"
+            - "Tasks per Node"
 
   "hybrid-single-node-config-sweep":
+    enabled: False
     run_configurations:
       - "cpp-hybrid"
       - "rust-hybrid"
@@ -218,7 +223,12 @@ benches:
         "Total flops": "FLOPS Summary:[\\s\\S]*Total\\s*: ([\\d\\.]+)[\\s\\S]*\nMFLOPS Summary"
         "Total mflops": "MFLOPS Summary:[\\s\\S]*Total\\s*: ([\\d\\.]+)"
         "Wall time (s)": "real\\s([\\d\\.]+)\nuser"
+        "Tasks per Node": "=== RUN INSTANTIATION ===\n\\{.*sbatch_config: \\{.*ntasks-per-node: (\\d+).*\\}"
+        "Threads": "=== RUN INSTANTIATION ===\n\\{.*sbatch_config: \\{.*cpus-per-task: (\\d+).*\\}"
       line_plots:
         - title: "Parallel Implementation Comparison"
           x: "Mesh x size"
           y: "Wall time (s)"
+          split_by:
+            - "Tasks per Node"
+            - "Threads"