5 Provides class for tracking retention rate of each cut in a skim.
8 __author__ =
"Cyrille Praz, Slavomira Stefkova"
11 import matplotlib.pyplot
as plt
13 from ROOT
import Belle2
19 """Check the retention rate and the number of candidates for a given set of particle lists.
21 The module stores its results in the static variable "summary".
23 To monitor the effect of every module of an initial path, this module should be added after
24 each module of the path. A function was written (`skimExpertFunctions.pathWithRetentionCheck`) to do it:
26 >>> path = pathWithRetentionCheck(particle_lists, path)
28 After the path processing, the result of the RetentionCheck can be printed with
30 >>> RetentionCheck.print_results()
32 or plotted with (check the corresponding documentation)
34 >>> RetentionCheck.plot_retention(...)
36 and the summary dictionary can be accessed through
38 >>> RetentionCheck.summary
42 Cyrille Praz, Slavomira Stefkova
46 module_name (str): name of the module after which the retention rate is measured
47 module_number (int): index of the module after which the retention rate is measured
48 particle_lists (list(str)): list of particle list names which will be tracked by the module
52 output_override =
None
54 def __init__(self, module_name='', module_number=0, particle_lists=[]):
65 type(self).summary[self.
_key] = {}
67 if type(self).output_override
is None:
82 if pl.getListSize() != 0:
98 b2.B2WARNING(
"Belle2.Environment.Instance().getNumberOfEvents() gives 0 or less.")
101 type(self).summary[self.
_key][particle_list] = {
"retention_rate": retention_rate,
108 """ Print the results, should be called after the path processing."""
110 table_headline =
"{:<100}|{:>9}|{:>12}|{:>22}|{:>12}|\n"
111 table_line =
"{:<100}|{:>9.3f}|{:>12}|{:>22}|{:>12}|\n"
115 for module, module_results
in cls.
summary.items():
117 for particle_list, list_results
in module_results.items():
119 if particle_list
not in summary_tables.keys():
121 atLeastOneEntry[particle_list] =
False
123 summary_tables[particle_list] = table_headline.format(
124 "Module",
"Retention",
"# Candidates",
"# Evts with candidates",
"Total # evts")
125 summary_tables[particle_list] +=
"=" * 160 +
"\n"
129 if list_results[
"retention_rate"] > 0
or atLeastOneEntry[particle_list]:
131 atLeastOneEntry[particle_list] =
True
132 if len(module) > 100:
133 module = module[:96] +
"..."
134 summary_tables[particle_list] += table_line.format(module, *list_results.values())
136 for particle_list, summary_table
in summary_tables.items():
137 b2.B2INFO(
"\n" +
"=" * 160 +
"\n" +
138 "Results of the modules RetentionCheck for the list " + particle_list +
".\n" +
140 "Note: the module RetentionCheck is defined in skim/scripts/skimExpertFunctions.py\n" +
144 "End of the results of the modules RetentionCheck for the list " + particle_list +
".\n" +
149 def plot_retention(cls, particle_list, plot_title="", save_as=None, module_name_max_length=80):
150 """ Plot the result of the RetentionCheck for a given particle list.
152 Example of use (to be put after process(path)):
154 >>> RetentionCheck.plot_retention('B+:semileptonic','skim:feiSLBplus','retention_plots/plot.pdf')
158 particle_list (str): particle list name
159 title (str): plot title (overwritten by the -o argument in basf2)
160 save_as (str): output filename (overwritten by the -o argument in basf2)
161 module_name_max_length (int): if the module name length is higher than this value, do not display the full name
166 at_least_one_entry =
False
167 for module, results
in cls.
summary.items():
169 if particle_list
not in results.keys():
170 b2.B2WARNING(particle_list +
" is not present in the results of the RetentionCheck for the module {}."
174 if results[particle_list][
'retention_rate'] > 0
or at_least_one_entry:
175 at_least_one_entry =
True
176 if len(module) > module_name_max_length
and module_name_max_length > 3:
177 module = module[:module_name_max_length - 3] +
"..."
178 module_name.append(module)
179 retention.append(100 * (results[particle_list][
'retention_rate']))
181 if not at_least_one_entry:
182 b2.B2WARNING(particle_list +
" seems to have a zero retention rate when created (if created).")
186 bars = plt.barh(module_name, retention, label=particle_list, color=(0.67, 0.15, 0.31, 0.6))
189 yval = bar.get_width()
190 plt.text(0.5, bar.get_y() + bar.get_height() / 2.0 + 0.1, str(round(yval, 3)))
192 plt.gca().invert_yaxis()
193 plt.xticks(rotation=45)
195 plt.axvline(x=10.0, linewidth=1, linestyle=
"--", color=
'k', alpha=0.5)
196 plt.xlabel(
'Retention Rate [%]')
197 plt.legend(loc=
'lower right')
202 save_as = plot_title +
'.pdf'
204 os.makedirs(os.path.dirname(save_as), exist_ok=
True)
205 plt.title(plot_title)
206 plt.savefig(save_as, bbox_inches=
"tight")
207 b2.B2RESULT(
"Retention rate results for list {} saved in {}."
208 .format(particle_list, os.getcwd() +
"/" + save_as))
211 def pathWithRetentionCheck(particle_lists, path):
212 """ Return a new path with the module RetentionCheck inserted between each module of a given path.
214 This allows for checking how the retention rate is modified by each module of the path.
216 Example of use (to be put just before process(path)):
218 >>> path = pathWithRetentionCheck(['B+:semileptonic'], path)
220 Warning: pathWithRetentionCheck(['B+:semileptonic'], path) does not modify path,
221 it only returns a new one.
223 After the path processing, the result of the RetentionCheck can be printed with
225 >>> RetentionCheck.print_results()
227 or plotted with (check the corresponding documentation)
229 >>> RetentionCheck.plot_retention(...)
231 and the summary dictionary can be accessed through
233 >>> RetentionCheck.summary
237 particle_lists (list(str)): list of particle list names which will be tracked by RetentionCheck
238 path (basf2.Path): initial path (it is not modified, see warning above and example of use)
241 for module_number, module
in enumerate(path.modules()):
242 new_path.add_module(module)
243 if 'ParticleSelector' in module.name():
244 name = module.name() +
'(' + module.available_params()[0].values +
')'
247 new_path.add_module(
RetentionCheck(name, module_number, particle_lists))