diff --git a/toolkit.py b/toolkit.py index 8d368613e7c8c34b38772b856e0a634e4533185c..5a312a204a434545032310cf87ea7d9fda706914 100755 --- a/toolkit.py +++ b/toolkit.py @@ -926,7 +926,7 @@ class ClassificationProject(object): plt.clf() - def plot_score(self, log=True, plot_opts=dict(bins=50, range=(0, 1))): + def plot_score(self, log=True, plot_opts=dict(bins=50, range=(0, 1)), ylim=None): centers_sig_train, hist_sig_train, _ = self.get_bin_centered_hist(self.scores_train[self.y_train==1].reshape(-1), density=True, weights=self.w_train[self.y_train==1], **plot_opts) centers_bkg_train, hist_bkg_train, _ = self.get_bin_centered_hist(self.scores_train[self.y_train==0].reshape(-1), density=True, weights=self.w_train[self.y_train==0], **plot_opts) centers_sig_test, hist_sig_test, rel_errors_sig_test = self.get_bin_centered_hist(self.scores_test[self.y_test==1].reshape(-1), density=True, weights=self.w_test[self.y_test==1], **plot_opts) @@ -941,6 +941,8 @@ class ClassificationProject(object): ax.errorbar(centers_sig_test, hist_sig_test, fmt="ro", yerr=errors_sig_test, label="signal test") if log: ax.set_yscale("log") + if ylim is not None: + ax.set_ylim(*ylim) ax.set_xlabel("NN output") fig.legend(loc='upper center', framealpha=0.5) fig.savefig(os.path.join(self.project_dir, "scores.pdf"))