Add parameter to plot function to display mean and stdevs

Signed-off-by: {Johann Heyszl} <[email protected]>

cw/simple_capture_aes_sca.py

@@ -180,7 +180,8 @@ def abort_handler_during_loop(project, sig, frame):
         if USE_HUSKY:
             project.traces.append(trace, dtype=np.uint16)
         if USE_WAVERUNNER:
-            project.traces.append(trace, dtype=np.uint8)
+            # Also use uint16 as dtype so that tvla processing works
+            project.traces.append(trace, dtype=np.uint16)
 
     # Save metadata and entire configuration cfg to project file
     project.settingsDict['datetime'] = datetime.now().strftime("%m/%d/%Y, %H:%M:%S")
@@ -192,6 +193,6 @@ def abort_handler_during_loop(project, sig, frame):
     # Create and show test plot
     if cfg["capture"]["show_plot"]:
         plot.save_plot_to_file(project.waves, None, cfg["capture"]["plot_traces"],
-                               cfg["capture"]["trace_image_filename"])
+                               cfg["capture"]["trace_image_filename"], add_mean_stddev=True)
         print(f'Created plot with {cfg["capture"]["plot_traces"]} traces: '
               f'{Path(cfg["capture"]["trace_image_filename"]).resolve()}')

cw/util/plot.py

@@ -4,13 +4,14 @@
 
 import itertools
 
+import numpy as np
 from bokeh.io import output_file
 from bokeh.models import tools
 from bokeh.palettes import Dark2_5 as palette
 from bokeh.plotting import figure, show
 
 
-def save_plot_to_file(traces, set_indices, num_traces, outfile):
+def save_plot_to_file(traces, set_indices, num_traces, outfile, add_mean_stddev=False):
     """Save plot figure to file."""
     if set_indices is None:
         colors = itertools.cycle(palette)
@@ -29,8 +30,28 @@ def save_plot_to_file(traces, set_indices, num_traces, outfile):
     plot.add_tools(tools.HoverTool())
     for i in range(min(len(traces), num_traces)):
         if set_indices is None:
-            plot.line(xrange, traces[i], line_color=next(colors))
+            if add_mean_stddev:
+                plot.line(xrange, traces[i], line_color='grey')
+            else:
+                plot.line(xrange, traces[i], line_color=next(colors))
         else:
             plot.line(xrange, traces[i], line_color=next(colors), legend_label=str(set_indices[i]))
+
+    if add_mean_stddev:
+        # Add mean and std dev to figure
+        # Convert selected traces to np array before processing
+        traces_new = np.empty((num_traces, len(traces[0])), dtype=np.uint16)
+        for i_trace in range(num_traces):
+            traces_new[i_trace] = traces[i_trace]
+        mean = traces_new.mean(axis=0)
+        std = traces_new.std(axis=0)
+        mean_stddev_upper = mean + std
+        mean_stddev_lower = mean - std
+        plot.line(xrange, mean_stddev_upper, line_color='firebrick',
+                  line_width=2, legend_label='std')
+        plot.line(xrange, mean_stddev_lower, line_color='firebrick',
+                  line_width=2, legend_label='std')
+        plot.line(xrange, mean, line_color='black', line_width=2, legend_label='mean')
+
     output_file(outfile)
     show(plot)