Alternative performance metrics in the terminal (#986)

scripts/advanced_statistics.py

@@ -7,74 +7,56 @@
 # ///
 
 import argparse
-import json
-from enum import Enum
 
 import numpy as np
 
-
-class Unit(Enum):
-    SECOND = 1
-    MILLISECOND = 2
-
-    def factor(self):
-        match self:
-            case Unit.SECOND:
-                return 1
-            case Unit.MILLISECOND:
-                return 1e3
-
-    def __str__(self):
-        match self:
-            case Unit.SECOND:
-                return "s"
-            case Unit.MILLISECOND:
-                return "ms"
-
+from plot_utils import METRICS, add_metric_argument, load_results
 
 parser = argparse.ArgumentParser()
 parser.add_argument("file", help="JSON file with benchmark results")
+add_metric_argument(parser)
 parser.add_argument(
     "--time-unit",
-    help="The unit of time.",
-    default="second",
-    action="store",
+    help="Display unit for time metrics (default: second).",
     choices=["second", "millisecond"],
-    dest="unit",
 )
 args = parser.parse_args()
 
-unit = Unit.MILLISECOND if args.unit == "millisecond" else Unit.SECOND
-unit_str = str(unit)
+metric, [results] = load_results(parser, [args.file], args.metric)
+description = METRICS[metric]
+unit, scale = description.unit, description.scale
+if args.time_unit is not None:
+    if not metric.startswith("time_"):
+        parser.error("--time-unit can only be used with a time metric")
+    if args.time_unit == "millisecond":
+        unit, scale = "ms", 0.001
 
-with open(args.file) as f:
-    results = json.load(f)["results"]
 
+def format_value(value):
+    suffix = f" {unit}" if unit else ""
+    return f"{value:.3f}{suffix}"
+
+
+print(f"Metric: {description.label}\n")
 for result in results:
-    summary = result["summary"]["time_wall_clock"]
-    ts = [m["time_wall_clock"]["value"] * unit.factor() for m in result["measurements"]]
+    summary = result["summary"][metric]
+    values = [m[metric]["value"] / scale for m in result["measurements"]]
 
-    p05 = np.percentile(ts, 5)
-    p25 = np.percentile(ts, 25)
-    p75 = np.percentile(ts, 75)
-    p95 = np.percentile(ts, 95)
+    p05, p25, p75, p95 = np.percentile(values, [5, 25, 75, 95])
 
     iqr = p75 - p25
 
     print(f"Command '{result.get('name', result['command'])}'")
     print(f"  runs:   {summary['count']:8d}")
-    print(f"  mean:   {summary['mean'] * unit.factor():8.3f} {unit_str}")
-    if summary["stddev"] is None:
-        print(f"  stddev: {'N/A':>8}")
-    else:
-        print(f"  stddev: {summary['stddev'] * unit.factor():8.3f} {unit_str}")
-    print(f"  median: {summary['median'] * unit.factor():8.3f} {unit_str}")
-    print(f"  min:    {summary['min'] * unit.factor():8.3f} {unit_str}")
-    print(f"  max:    {summary['max'] * unit.factor():8.3f} {unit_str}")
+    for statistic in ["mean", "stddev", "median", "min", "max"]:
+        value = summary[statistic]
+        formatted = "N/A" if value is None else format_value(value / scale)
+        print(f"  {statistic + ':':<8}{formatted:>8}")
     print()
     print("  percentiles:")
-    print(f"     P_05 .. P_95:    {p05:.3f} {unit_str} .. {p95:.3f} {unit_str}")
+    print(f"     P_05 .. P_95:    {format_value(p05)} .. {format_value(p95)}")
     print(
-        f"     P_25 .. P_75:    {p25:.3f} {unit_str} .. {p75:.3f} {unit_str}  (IQR = {iqr:.3f} {unit_str})"
+        f"     P_25 .. P_75:    {format_value(p25)} .. {format_value(p75)}"
+        f"  (IQR = {format_value(iqr)})"
     )
     print()