1 files changed, 15 insertions, 2 deletions

diff --git a/util/plot_dh.py b/util/plot_dh.py
index 62a2f86..2aa7d21 100755
--- a/util/plot_dh.py
+++ b/util/plot_dh.py
@@ -1,4 +1,5 @@
 #!/usr/bin/env python
+# -*- coding: UTF-8 -*-
 #
 # Script for plotting ECTester ECDH results.
 #
@@ -23,10 +24,19 @@ if __name__ == "__main__":
 
     opts = parser.parse_args()
 
+    with open(opts.file, "r") as f:
+        header = f.readline()
+    header_names = header.split(";")
+
     hx = lambda x: int(x, 16)
     data = np.genfromtxt(opts.file, delimiter=";", skip_header=1, converters={2: hx, 3: hx, 4: hx}, dtype=np.dtype([("index","u4"), ("time","u4"), ("pub", "O"), ("priv", "O"), ("secret","O")]))
 
-    time_data = map(itemgetter(1), data)
+    if "nano" in header_names[1]:
+        unit = r"$\mu s$"
+        time_data = map(lambda x: x[1]/1000, data)
+    else:
+        unit = r"ms"
+        time_data = map(itemgetter(1), data)
     priv_data = map(itemgetter(2), data)
     pub_data = map(itemgetter(3), data)
     secret_data = map(itemgetter(4), data)
@@ -42,11 +52,14 @@ if __name__ == "__main__":
     axe_hist.axvline(x=time_avg, alpha=0.7, linestyle="dotted", color="red", label="avg = {}".format(time_avg))
     axe_hist.axvline(x=time_median, alpha=0.7, linestyle="dotted", color="green", label="median = {}".format(time_median))
     axe_hist.set_ylabel("count\n(log)")
-    axe_hist.set_xlabel("time (ms)")
+    axe_hist.set_xlabel("time ({})".format(unit))
     axe_hist.xaxis.set_major_locator(ticker.MaxNLocator())
     axe_hist.legend(loc="best")
 
+    fig.text(0.01, 0.02, "Data size: {}".format(len(time_data)), size="small")
+
     if opts.output is None:
         plt.show()
     else:
+        fig.set_size_inches(12, 10)
         plt.savefig(opts.output, dpi=400)