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