made plotting a bit nicer

plotting/roofline.py

 import plotly.graph_objs as go
-from dataclasses import dataclass
 import json
 
 try:
@@ -8,18 +7,6 @@ except ImportError:
     # Try backported to PY<37 `importlib_resources`.
     import importlib_resources as pkg_resources
 
-
-@dataclass
-class CPUInfo:
-    name: str
-    max_perf: float  # MFLOP/s
-    max_band: float  # MByte/s
-
-
-SKYLAKE = CPUInfo("Intel(R) Xeon(R) Gold 6148 CPU @ 2.40GHz",
-                  max_perf=5582274.03,
-                  max_band=159321.74)
-
 DEFAULT_MARKER = dict(size=10, color='green')
 
 
@@ -30,14 +17,15 @@ def get_data(file_name="roofline_data.json"):
 
 
 def add_compute_roofline(fig, peak_perf: float, ridge_point: float, max_x: float, *, name='Peak Performance'):
+    text = f'{name}: {peak_perf/1e12:.4f} TFlop/s'
     fig.add_trace(
         go.Scatter(x=[ridge_point, max_x],
                    y=[peak_perf, peak_perf],
                    mode='lines+text',
-                   name=f'{name}: {peak_perf/1e9} GFlop/s',
-                   text=['', f'{name}: {peak_perf/1e9} GFlop/s'],
+                   name=text,
+                   text=['', text],
                    textposition="top left",
-                   line=dict(color='black', width=2))
+                   line=dict(color='black', width=4))
     )
     return fig
 
@@ -47,9 +35,9 @@ def add_memory_roofline(fig, peak_perf: float, ridge_point: float, max_bw: float
         go.Scatter(x=[0, ridge_point],
                    y=[0, peak_perf],
                    mode='lines+text',
-                   name=f"Memory bandwidth ({name}): {max_bw/1e9} GB/s",
+                   name=f"Memory bandwidth ({name}): {max_bw/1e9:.4f} GB/s",
                    textposition="bottom left",
-                   line=dict(color='black', width=2))
+                   line=dict(color='black', width=4))
     )
     return fig
 
@@ -98,26 +86,30 @@ def set_axis(fig):
     return fig
 
 
+def create_figure(title: str):
+    fig = go.Figure()
+    fig = set_axis(fig)
+    fig.update_layout(title=title)
+    return fig
+
+
 def create_plot_by_host(hostname: str, title: str, *, max_x=50):
     data = get_data()[hostname]
-    fig = go.Figure()
     max_band = data["bw"]["load"] * 1e6
     perfs = sorted(data["perf"].items(), key=lambda x: x[1], reverse=True)
-    for name, perf in perfs:
-        add_compute_roofline(fig, perf*1e6, (perf*1e6/max_band), max_x, name=name)
+    fig = create_figure(title)
+    for idx, (name, perf) in enumerate(perfs):
+        perf_in_flops = perf * 1e6
+        add_compute_roofline(fig, perf_in_flops, (perf_in_flops/max_band), max_x, name=name)
+        if idx == 0:
+            add_memory_roofline(fig, perf_in_flops, perf_in_flops/max_band, max_band, "load")
 
-    add_memory_roofline(fig, perfs[0][1]*1e6, perfs[0][1]*1e6/max_band, max_band, "load")
-
-    fig = set_axis(fig)
-    fig.update_layout(title=title)
     return fig
 
 
 def create_plot(peak_perf, max_band, title):
-    fig = go.Figure()
+    fig = create_figure(title)
     fig = add_rooflines(fig, peak_perf, max_band, 100)
-    fig = set_axis(fig)
-    fig.update_layout(title=title)
     return fig