Belle II Software  release-08-01-10
analysis-lookuptable-creation-energy-bias-correction.py
1 #!/usr/bin/env python3
2 
3 
10 
11 
18 
19 import basf2 as b2
20 import pandas as pd
21 import argparse
22 
23 # add argument to get the input file with energy bias correction values
24 parser = argparse.ArgumentParser()
25 parser.add_argument(
26  '-i',
27  '--input_path',
28  default="/nfs/dust/belle2/user/kovalch/energyBiasCorrFiles/Feb2021.txt",
29  help='Path to input file with energy bias correction values .')
30 args = parser.parse_args()
31 
32 weightsDB = args.input_path
33 df_weightDB = pd.read_csv(weightsDB, delimiter=r"\s+")
34 
35 tableName = "Feb2021:TestEnergy"
36 tableName2 = "Feb2021:TestEnergy2"
37 
38 # Add some bin constructors
39 
40 
41 def make_1D_bin(name, min_val, max_val):
42  return {name: [min_val, max_val]}
43 
44 
45 def get_weight(namew, w, namewerr, werr):
46  return {namew: w, namewerr: werr}
47 
48 
49 def make_2D_bin(bin_x, bin_y):
50  bin_2d = bin_x.copy()
51  bin_2d.update(bin_y)
52  return bin_2d
53 
54 
55 def make_3D_bin(bin_x, bin_y, bin_z):
56  bin_3d = bin_x.copy()
57  bin_3d.update(bin_y).update(bin_z)
58  return bin_3d
59 
60 
61 # Bin IDs will be automatically assigned as follows:
62 #
63 # E|0-2|2-3|3-4|
64 # theta| | | |
65 # -------------
66 # 0-1| 0 | 2 | 4 |
67 # ---------------|
68 # 1-2| 1 | 3 | 5 |
69 # -------------
70 
71 
72 bins_x = []
73 bins_y = []
74 weightInfo = {}
75 weightInfo_x = []
76 tableIDNotSpec = []
77 binID = 0
78 tableIDSpec = []
79 for x in range(0, len(df_weightDB)):
80  weightInfo = {}
81  bins_x = make_1D_bin(df_weightDB['namebin1'].values[x], df_weightDB['minbin1'].values[x], df_weightDB['maxbin1'].values[x])
82  bins_y = make_1D_bin(df_weightDB['namebin2'].values[x], df_weightDB['minbin2'].values[x], df_weightDB['maxbin2'].values[x])
83  weightInfo = get_weight("Weight", df_weightDB['w'].values[x], "StatErr", df_weightDB['werr'].values[x])
84  tableIDNotSpec.append([weightInfo, make_2D_bin(bins_x, bins_y)])
85  binID += 1
86  tableIDSpec.append([[weightInfo, make_2D_bin(bins_x, bins_y)], binID])
87 
88 # And of course let's define out-of-range bin info
89 outOfRangeWeightInfo = {}
90 outOfRangeWeightInfo["Weight"] = -1
91 outOfRangeWeightInfo["StatErr"] = -1
92 
93 # Now, let's configure table creator
94 addtable = b2.register_module('ParticleWeightingLookUpCreator')
95 addtable.param('tableIDSpec', tableIDSpec)
96 addtable.param('outOfRangeWeight', outOfRangeWeightInfo)
97 addtable.param('experimentHigh', -1)
98 addtable.param('experimentLow', 0)
99 addtable.param('runHigh', -1)
100 addtable.param('runLow', 0)
101 addtable.param('tableName', tableName)
102 
103 addtable2 = b2.register_module('ParticleWeightingLookUpCreator')
104 addtable2.param('tableIDNotSpec', tableIDNotSpec)
105 addtable2.param('outOfRangeWeight', outOfRangeWeightInfo)
106 addtable2.param('experimentHigh', -1)
107 addtable2.param('experimentLow', 0)
108 addtable2.param('runHigh', -1)
109 addtable2.param('runLow', 0)
110 addtable2.param('tableName', tableName2)
111 
112 eventinfosetter = b2.register_module('EventInfoSetter')
113 eventinfosetter.param('evtNumList', [10])
114 eventinfosetter.param('runList', [0])
115 eventinfosetter.param('expList', [0])
116 
117 my_path = b2.create_path()
118 my_path.add_module(addtable)
119 my_path.add_module(addtable2)
120 my_path.add_module(eventinfosetter)
121 
122 b2.process(my_path)