2 files changed, 53 insertions, 17 deletions

diff --git a/util/plot_gen.py b/util/plot_gen.py
index 0d518e6..4ee1ddc 100755
--- a/util/plot_gen.py
+++ b/util/plot_gen.py
@@ -18,12 +18,13 @@ from matplotlib import ticker, colors
 from copy import deepcopy
 import argparse
 
-from utils import hw, moving_average, plot_hist
+from utils import hw, moving_average, plot_hist, miller_correction
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Plot results of ECTester key generation timing.")
     parser.add_argument("-o", "--output", dest="output", type=argparse.FileType("wb"), help="Write image to [file], do not display.", metavar="file")
     parser.add_argument("--priv", dest="priv", action="store_true", help="Show private key MSB heatmap plot.")
+    parser.add_argument("--entropy", dest="entropy", action="store_true", help="Show estimated entropy of private key MSB conditioned on time of generation.")
     parser.add_argument("--hist", dest="hist", action="store_true", help="Show keygen time histogram.")
     parser.add_argument("--export-hist", dest="export_hist", action="store_true", help="Show export time histogram.")
     parser.add_argument("--avg", dest="avg", action="store_true", help="Show moving average of keygen time.")
@@ -103,21 +104,6 @@ if __name__ == "__main__":
 
     sorted_data = np.sort(data, order="gen_time")
 
-    i = 0
-    entropies = {}
-    while i < len(data):
-        time_val = sorted_data["gen_time"][i]
-        j = i
-        msbs = [0 for _ in range(256)]
-        while j < len(data) and sorted_data["gen_time"][j] == time_val:
-            msbs[(sorted_data["priv"][j] >> (bit_size - 8)) & 0xff] += 1
-            j += 1
-        if j - 100 > i:
-            entropies[time_val] = entropy(msbs, base=2)
-        i = j
-
-    entropy = sum(entropies.values())/len(entropies)
-
     cmap = deepcopy(plt.cm.plasma)
     cmap.set_bad("black")
 
@@ -174,7 +160,23 @@ if __name__ == "__main__":
         fig.colorbar(im, ax=axe_priv_hist)
 
     fig.text(0.01, 0.02, "Data size: {}".format(len(gen_time_data)), size="small")
-    fig.text(0.01, 0.04, "Entropy of privkey MSB(estimated): {:.2f} b".format(entropy), size="small")
+
+    if opts.entropy:
+        i = 0
+        entropies = {}
+        while i < len(data):
+            time_val = sorted_data["gen_time"][i]
+            j = i
+            msbs = [0 for _ in range(256)]
+            while j < len(data) and sorted_data["gen_time"][j] == time_val:
+                msbs[(sorted_data["priv"][j] >> (bit_size - 8)) & 0xff] += 1
+                j += 1
+            if j - 100 > i:
+                entropies[time_val] = miller_correction(entropy(msbs, base=2), j - i, 256)
+            i = j
+    
+        entropy = sum(entropies.values())/len(entropies)
+        fig.text(0.01, 0.04, "Entropy of privkey MSB(estimated): {:.2f} b".format(entropy), size="small")
 
     if opts.output is None:
         plt.show()
diff --git a/util/utils.py b/util/utils.py
new file mode 100644
index 0000000..bddfc35
--- /dev/null
+++ b/util/utils.py
@@ -0,0 +1,34 @@
+import numpy as np
+from matplotlib import ticker
+
+
+def hw(i):
+    res = 0
+    while i:
+        res += 1
+        i &= i - 1
+    return res
+
+
+def moving_average(a, n) :
+    ret = np.cumsum(a, dtype=float)
+    ret[n:] = ret[n:] - ret[:-n]
+    return ret[n - 1:] / n
+
+
+def plot_hist(axes, data, xlabel=None, log=False):
+    time_max = max(data)
+    time_min = min(data)
+    time_avg = np.average(data)
+    time_median = np.median(data)
+    axes.hist(data, bins=time_max - time_min, log=log)
+    axes.axvline(x=time_avg, alpha=0.7, linestyle="dotted", color="blue", label="avg = {}".format(time_avg))
+    axes.axvline(x=time_median, alpha=0.7, linestyle="dotted", color="green", label="median = {}".format(time_median))
+    axes.set_ylabel("count" + ("\n(log)" if log else ""))
+    axes.set_xlabel("time" if xlabel is None else xlabel)
+    axes.xaxis.set_major_locator(ticker.MaxNLocator())
+    axes.legend(loc="best")
+
+
+def miller_correction(entropy, samples, bins):
+    return entropy + (bins - 1)/(2*samples)