Belle II Software development
vars.py
1# this script sets the aliases for different variable sets
2
3# relative imports
4from .categories import get_available_categories
5from .helpers import read_yaml
6
7
8def set_output_vars():
9 """
10 Sets the aliases for FT output vars
11 """
12
13 # imports
14 import basf2
15 from variables import variables as va
16 import variables.utils as vu
17
18 # aliases for FT output vars
19 va.addAlias('FBDT_qrCombined', 'qrOutput(FBDT)')
20 va.addAlias('FANN_qrCombined', 'qrOutput(FANN)')
21 va.addAlias('qrGNN', 'extraInfo(qrGNN)')
22
23 # alias for FT target var
24 va.addAlias('qrMC', 'isRelatedRestOfEventB0Flavor')
25
26 # aliases for FT category info vars
27 for category in get_available_categories():
28 aliasForQp = f'qp{category}'
29 aliasForTrueTarget = f'hasTrueTarget{category}'
30 aliasForIsRightCategory = f'isRightCategory{category}'
31 va.addAlias(aliasForQp, f'qpCategory({category})')
32 va.addAlias(aliasForTrueTarget, f'hasTrueTargets({category})')
33 va.addAlias(aliasForIsRightCategory, f'isTrueFTCategory({category})')
34
35 # get output vars list
36 filepath = basf2.find_file('data/analysis/output_vars.yaml')
37 flavor_tagging_vars = read_yaml(filepath)["output_vars"]
38
39 # add output vars to flavor tagger collection
40 vu.add_collection(flavor_tagging_vars, 'flavor_tagging')
41
42
43def set_FT_pid_aliases(type="Current", exp_type="Belle2"):
44 """
45 This function adds the pid aliases needed by the flavor tagger.
46 """
47 # imports
48 import basf2
49 from variables import variables as va
50
51 # define a function to set alias
52 def set_alias(alias, var): return va.addAlias(alias, var)
53
54 # get list of aliases
55 filepath = basf2.find_file('data/analysis/pid_alias.yaml')
56 alias_list = read_yaml(filepath)[type]
57
58 for alias, var in alias_list["Common"].items():
59 set_alias(alias, var) # common alias
60 for alias, var in alias_list[exp_type].items():
61 set_alias(alias, var) # exp specific alias
62
63
64def set_GFlat_aliases(categories, usePIDNN):
65 """
66 This function adds aliases for the GNN-based flavor tagger.
67 """
68 # imports
69 from variables import variables as va
70
71 # will be used for target variable 0:B0bar, 1:B0
72 va.addAlias('qrCombined_bit', '(qrCombined+1)/2')
73 alias_list = ['qrCombined_bit']
74
75 var_dict = {
76 # position
77 'dx': 'dx',
78 'dy': 'dy',
79 'dz': 'dz',
80 # mask
81 'E': 'E',
82 # charge,
83 'charge': 'charge',
84 # feature
85 'px_c': 'px*charge',
86 'py_c': 'py*charge',
87 'pz_c': 'pz*charge',
88 'electronID_c': 'electronIDNN*charge' if usePIDNN else 'electronID*charge',
89 'muonID_c': 'muonIDNN*charge' if usePIDNN else 'muonID*charge',
90 'pionID_c': 'pionIDNN*charge' if usePIDNN else 'pionID*charge',
91 'kaonID_c': 'kaonIDNN*charge' if usePIDNN else 'kaonID*charge',
92 'protonID_c': 'protonIDNN*charge' if usePIDNN else 'protonID*charge',
93 'deuteronID_c': 'deuteronIDNN*charge' if usePIDNN else 'deuteronID*charge',
94 'electronID_noSVD_noTOP_c': 'electronID_noSVD_noTOP*charge',
95 }
96
97 # 16 charged particles are used at most
98 for rank in range(1, 17):
99
100 for cat in categories:
101 listName = get_available_categories()[cat].particleList
102 varName = f'QpTrack({listName}, isRightCategory({cat}), isRightCategory({cat}))'
103
104 varWithRank = f'ifNANgiveX(getVariableByRank(pi+:inRoe, FT_p, {varName}, {rank}), 0)'
105 aliasWithRank = f'{cat}_rank{rank}'
106
107 va.addAlias(aliasWithRank, varWithRank)
108 alias_list.append(aliasWithRank)
109
110 for alias, var in var_dict.items():
111 varWithRank = f'ifNANgiveX(getVariableByRank(pi+:inRoe, FT_p, {var}, {rank}), 0)'
112 aliasWithRank = f'{alias}_rank{rank}'
113
114 va.addAlias(aliasWithRank, varWithRank)
115 alias_list.append(aliasWithRank)
116
117 return alias_list
118
119
120def set_masked_vars(mask_name="all"):
121 """
122 Returns variables with a certain mask
123 """
124 # imports
125 from variables import variables as va
126
127 masked_var_dict = {
128 'pMissTag_withMask': f'pMissTag({mask_name})',
129 'cosTPTO_withMask': f'cosTPTO({mask_name})',
130 'ptTracksRoe_withMask': f'ptTracksRoe({mask_name})',
131 'pt2TracksRoe_withMask': f'pt2TracksRoe({mask_name})',
132 }
133
134 for key, value in masked_var_dict.items():
135 va.addAlias(key, value)
136
137 return masked_var_dict
138
139
140def set_and_get_aliases(var_dict):
141 """
142 Sets aliases from a dict of alias, var pair and returns a list of aliases.
143 """
144 from variables import variables as va
145
146 for key, value in var_dict.items():
147 va.addAlias(key, value)
148
149 return list(var_dict.keys())