Belle II Software light-2406-ragdoll
udst_purge_list.py
1#!/usr/bin/env python3
2
3
10
11# @cond
12
13import basf2
14from mdst import add_mdst_output
15import modularAnalysis as ma
16import vertex as vtx
17from variables import variables as vm
18import argparse
19
20
21def preprocess_roe_data(listName, charge, path):
22 ''' Process/select events on top of the skim processing:
23 -- vertex fit
24 -- select of the positive/negative charge
25 -- extra cuts
26 '''
27 # charge
28 if not (charge == 0):
29 if listName.find("B0") == 0:
30 basf2.B2ERROR(f"Inconsistency between neutral list {listName} and charge requirement {charge}")
31 ma.variablesToEventExtraInfo(listName, {'charge': 'bcharge'}, path=path)
32 vm.addAlias('bcharge', 'eventExtraInfo(bcharge)')
33
34 # Filter:
35 empty_path = basf2.create_path()
36 cfilter = basf2.register_module('VariableToReturnValue')
37 cfilter.param('variable', 'bcharge')
38 cfilter.if_value(f'!={charge}', empty_path)
39 path.add_module(cfilter).set_name("chargeFilter")
40
41 # vertex fit
42 vtx.treeFit(listName, path=path)
43 ma.variablesToEventExtraInfo(listName, {'chiProb': 'chiProb'}, path=path)
44 vm.addAlias('fitProb', 'eventExtraInfo(chiProb)')
45
46 # filter failed fits ...
47 empty_path2 = basf2.create_path()
48 fitfilter = basf2.register_module('VariableToReturnValue')
49 fitfilter.param('variable', 'fitProb')
50 fitfilter.if_value('<0', empty_path2)
51 path.add_module(fitfilter).set_name("fitFilter")
52
53 # require one candidate per event
54 vm.addAlias('nB', f'nParticlesInList({listName})')
55 eventFilter = basf2.register_module('VariableToReturnValue')
56 eventFilter.param('variable', 'nB')
57 eventFilter.if_value('!=1', empty_path2)
58 path.add_module(eventFilter).set_name("eventFilter")
59
60 # store extra momentum/vertex position information
61 Vars = {'px': 'PX', 'py': 'PY', 'pz': 'PZ', 'x': 'X', 'y': 'Y', 'z': 'Z',
62 'IPX': 'IPX', 'IPY': 'IPY', 'IPZ': 'IPZ'}
63 ma.variablesToEventExtraInfo(listName, Vars, path=path)
64
65
66def list_select_roe_data(listName, path):
67
68 filter1 = basf2.register_module("UdstListFilter")
69 filter1.param("listName", listName)
70 filter1.param("keepNotInList", True)
71 path.add_module(filter1).set_name("filter1")
72
73
74def preprocess_signal_mc(listName, charge, path):
75 # charge
76 if not (charge == 0):
77 if listName.find("B0") == 0:
78 basf2.B2ERROR(f"Inconsistency between neutral list {listName} and charge requirement {charge}")
79 ma.variablesToEventExtraInfo(listName, {'charge': 'bcharge'}, path=path)
80 vm.addAlias('bcharge', 'eventExtraInfo(bcharge)')
81 empty_path = basf2.create_path()
82 cfilter = basf2.register_module('VariableToReturnValue')
83 cfilter.param('variable', 'bcharge')
84 cfilter.if_value(f'!={charge}', empty_path)
85 path.add_module(cfilter).set_name("chargeFilter")
86
87 vm.addAlias('nB', f'nParticlesInList({listName})')
88 empty_path2 = basf2.create_path()
89 eventFilter = basf2.register_module('VariableToReturnValue')
90 eventFilter.param('variable', 'nB')
91 eventFilter.if_value('!=1', empty_path2)
92 path.add_module(eventFilter).set_name("eventFilter")
93
94 Vars = {'mcPX': 'PX', 'mcPY': 'PY', 'mcPZ': 'PZ',
95 'mcDecayVertexX': 'X',
96 'mcDecayVertexY': 'Y',
97 'mcDecayVertexZ': 'Z',
98 'IPX': 'IPX', 'IPY': 'IPY', 'IPZ': 'IPZ'
99 }
100 ma.variablesToEventExtraInfo(listName, Vars, path=path)
101
102
103def list_select_signal_mc(listName, path):
104 filter2 = basf2.register_module("UdstListFilter")
105 filter2.param("listName", listName)
106 filter2.param("keepNotInList", False)
107 path.add_module(filter2).set_name("filter2")
108
109
110def prepare_path(FileIn, ListName, icharge, isROE, FileOut):
111 '''
112 Return processing path
113 FileIn -- input file name
114 ListName -- B+ or B0 decay chain of the skim
115 icharge -- charge (0, -1 or +1)
116 isROE -- ROE (data) if true, signal (MC) if false
117 FileOut -- output file name
118 '''
119 main = basf2.Path()
120
121 # input
122 input1 = basf2.register_module('RootInput')
123 input1.param('inputFileName', FileIn)
124 main.add_module(input1).set_name("input1")
125
126 if (isROE):
127 preprocess_roe_data(ListName, icharge, main)
128 list_select_roe_data(ListName, main)
129 else:
130 preprocess_signal_mc(ListName, icharge, main)
131 list_select_signal_mc(ListName, main)
132
133 # output
134 additionalBranches = ['EventExtraInfo']
135
136 add_mdst_output(main, filename=FileOut, additionalBranches=additionalBranches)
137
138 # progress
139 main.add_module('Progress')
140
141 basf2.print_path(main)
142 return main
143
144
145def get_parser():
146 """Handles the command-line argument parsing.
147
148 Returns:
149 argparse.Namespace: The parsed arguments.
150 """
151 parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
152 parser.add_argument("--fileIn", help='Input file name.', required=True)
153 parser.add_argument('--listName', help='Name of the particle list to keep (isSignal) or to remove.', required=True)
154 parser.add_argument("--isSignal", action='store_true', help='If added, the input is signal MC')
155 parser.add_argument("--charge", default='pos',
156 help='Filter charge. Possible values are :code:`pos`, :code:`neg` or :code:`zero`. Default is :code:`pos`')
157 parser.add_argument("--fileOut", default='test.root', help='Output file name. Default: :code:`test.root`')
158
159 return parser
160
161
162
163if __name__ == '__main__':
164
165 args = get_parser().parse_args()
166
167 list1 = args.listName
168 signal = args.isSignal
169 charge = args.charge
170 fileOut = args.fileOut
171 fileIn = args.fileIn
172
173 if charge == 'zero':
174 icharge = 0
175 elif charge == 'pos':
176 icharge = 1
177 elif charge == 'neg':
178 icharge = -1
179
180 main = prepare_path(fileIn, list1, icharge, not signal, fileOut)
181 basf2.process(main)
182
183 print(basf2.statistics)
184
185# @endcond