diff options
Diffstat (limited to 'util/utils.py')
| -rw-r--r-- | util/utils.py | 16 |
1 files changed, 11 insertions, 5 deletions
diff --git a/util/utils.py b/util/utils.py index bddfc35..d6b9aed 100644 --- a/util/utils.py +++ b/util/utils.py @@ -16,18 +16,24 @@ def moving_average(a, n) : return ret[n - 1:] / n -def plot_hist(axes, data, xlabel=None, log=False): +def plot_hist(axes, data, xlabel=None, log=False, avg=True, median=True, bins=None, **kwargs): 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)) + if bins is None: + bins = time_max - time_min + 1 + hist = axes.hist(data, bins=bins, log=log, **kwargs) + if avg: + axes.axvline(x=time_avg, alpha=0.7, linestyle="dotted", color="blue", label="avg = {}".format(time_avg)) + if median: + 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") + if avg or median: + axes.legend(loc="best") + return hist def miller_correction(entropy, samples, bins): |