Belle II Software light-2406-ragdoll
particleWeighting.py
1#!/usr/bin/env python3
2
3
10
11from basf2 import B2FATAL, B2RESULT, Path, process, conditions, register_module
12import b2test_utils
13import modularAnalysis as ma
14from variables import variables
15import random
16from ROOT import TFile
17
18
23
24# Add some bin constructors
25
26
27def make_1D_bin(name, min_val, max_val):
28 return {name: [min_val, max_val]}
29
30
31def make_2D_bin(bin_x, bin_y):
32 bin_2d = bin_x.copy()
33 bin_2d.update(bin_y)
34 return bin_2d
35
36
37def make_3D_bin(bin_x, bin_y, bin_z):
38 bin_3d = bin_x.copy()
39 bin_3d.update(bin_y).update(bin_z)
40 return bin_3d
41
42
43def check(ntupleName, treeName):
44 """
45 Verify results make sense.
46 """
47 ntuplefile = TFile(ntupleName)
48 ntuple = ntuplefile.Get(treeName)
49
50 if ntuple.GetEntries() == 0:
51 B2FATAL("No pions saved")
52
53 if not(ntuple.GetEntries("binID > 0 ") > 0):
54 B2FATAL("Binning was applied incorrectly: no pions in physical bins")
55 else:
56 B2RESULT("Bins are defined")
57
58 if not(ntuple.GetEntries("Weight > 0") > 0):
59 B2FATAL("Weights are not applied")
60 else:
61 B2RESULT("Weights are applied")
62
63
69
70
71# Define bin ranges. Bins may be of different size
72# They shouldn't event cover the whole parameter space
73bins_x = [make_1D_bin("p", 0, 2),
74 make_1D_bin("p", 2, 3),
75 make_1D_bin("p", 3, 4)]
76
77bins_y = [make_1D_bin("pz", 0, 1),
78 make_1D_bin("pz", 1, 5)]
79
80
81# Here would be input from the experts.
82# We don't have real calibration tables yet
83# So we use randomized info.
84# Bin IDs will be automatically assigned as follows:
85#
86# p|0-2|2-3|3-4|
87# pz| | | |
88# -------------
89# 0-1| 0 | 2 | 4 |
90# ---------------|
91# 1-5| 1 | 3 | 5 |
92# -------------
93
94tableIDNotSpec = []
95random.seed()
96for xbin in bins_x:
97 for ybin in bins_y:
98 weightInfo = {}
99 weightInfo["Weight"] = float(random.randint(0, 100)) / 100
100 weightInfo["StatErr"] = float(random.randint(100, 200)) / 100
101 weightInfo["SystErr"] = float(10)
102 tableIDNotSpec.append([weightInfo, make_2D_bin(xbin, ybin)])
103
104# Now let's assign binIDs manually (to be able to assign identical IDs for different bins)
105# Let's do this structure of bins (out-of-range bin -1 is assigned later):
106#
107# p|0-2|2-3|3-4|
108# pz| | | |
109# ---------------|
110# | |
111# ------------- |
112# 0-1| 43| 45| | |
113# ------------ | |
114# 1-5| 44| 46| |
115# ------------- |
116# | -1|
117# ---------------|
118tableIDSpec = []
119random.seed()
120binID = 42
121weightInfo = {}
122for xbin in bins_x:
123 for ybin in bins_y:
124 if binID < 46:
125 binID += 1
126 weightInfo = {}
127 weightInfo["Weight"] = float(random.randint(0, 100)) / 100
128 weightInfo["StatErr"] = float(random.randint(100, 200)) / 100
129 weightInfo["SystErr"] = float(10)
130 tableIDSpec.append([[weightInfo, make_2D_bin(xbin, ybin)], binID])
131
132# And of course let's define out-of-range bin info
133outOfRangeWeightInfo = {}
134outOfRangeWeightInfo["Weight"] = -1
135outOfRangeWeightInfo["StatErr"] = -1
136outOfRangeWeightInfo["SystErr"] = -1
137
138# we create payloads so let's switch to an empty, temporary directory
140 conditions.testing_payloads = ["localdb/database.txt"]
141
142 # Now, let's configure table creator
143 addtable = register_module('ParticleWeightingLookUpCreator')
144 addtable.param('tableIDNotSpec', tableIDNotSpec)
145 addtable.param('outOfRangeWeight', outOfRangeWeightInfo)
146 addtable.param('experimentHigh', 1003)
147 addtable.param('experimentLow', 1003)
148 addtable.param('runHigh', 1000)
149 addtable.param('runLow', 0)
150 addtable.param('tableName', "ParticleReweighting:TestMomentum")
151
152 addtable2 = register_module('ParticleWeightingLookUpCreator')
153 addtable2.param('tableIDSpec', tableIDSpec)
154 addtable2.param('outOfRangeWeight', outOfRangeWeightInfo)
155 addtable2.param('experimentHigh', 1003)
156 addtable2.param('experimentLow', 1003)
157 addtable2.param('runHigh', 1000)
158 addtable2.param('runLow', 0)
159 addtable2.param('tableName', "ParticleReweighting:TestMomentum2")
160
161 testpath = Path()
162 testpath.add_module(addtable)
163 testpath.add_module(addtable2)
164 testpath.add_module('EventInfoSetter', evtNumList=[100], runList=[0], expList=[1003])
165
166 # Process the events to create database payloads
167 process(testpath)
168 B2RESULT("Weights are created and loaded to DB")
169
170 # Now, let's test if it weights are applied
171 main = Path()
172 ntupleName = 'particleWeighting.root'
173 treeName = 'pitree'
174 inputfile = b2test_utils.require_file('analysis/tests/mdst.root')
175 ma.inputMdst(inputfile, path=main)
176
177 # use the MC truth information to generate pion lists
178 ma.fillParticleListFromMC('pi+:gen', '', path=main)
179
180 # ID of weight table is taken from B2A904
181 weight_table_id = "ParticleReweighting:TestMomentum"
182
183 # We know what weight info will be added (see B2A904),
184 # so we add aliases and add it to tools
185 variables.addAlias('Weight', 'extraInfo(' + weight_table_id + '_Weight)')
186 variables.addAlias('StatErr', 'extraInfo(' + weight_table_id + '_StatErr)')
187 variables.addAlias('SystErr', 'extraInfo(' + weight_table_id + '_SystErr)')
188 variables.addAlias('binID', 'extraInfo(' + weight_table_id + '_binID)')
189 varsPi = ['p', 'pz', 'Weight', 'StatErr', 'SystErr', 'binID']
190
191 # We configure weighing module
192 reweighter = register_module('ParticleWeighting')
193 reweighter.param('tableName', weight_table_id)
194 reweighter.param('particleList', 'pi+:gen')
195 main.add_module(reweighter)
196
197 # write out the flat ntuple
198 ma.variablesToNtuple('pi+:gen', varsPi, filename=ntupleName, treename=treeName, path=main)
199
200 # Process the events
202
203 check(ntupleName, treeName)
204
205B2RESULT("Weights were applied correctly")
def require_file(filename, data_type="", py_case=None)
Definition: __init__.py:54
def clean_working_directory()
Definition: __init__.py:189
def safe_process(*args, **kwargs)
Definition: __init__.py:236