4from .categories
import get_available_categories
5from .helpers
import read_yaml
10 Sets the aliases for FT output vars
15 from variables
import variables
as va
19 va.addAlias(
'FBDT_qrCombined',
'qrOutput(FBDT)')
20 va.addAlias(
'FANN_qrCombined',
'qrOutput(FANN)')
21 va.addAlias(
'qrGNN',
'extraInfo(qrGNN)')
24 va.addAlias(
'qrMC',
'isRelatedRestOfEventB0Flavor')
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})')
36 filepath = basf2.find_file(
'data/analysis/output_vars.yaml')
37 flavor_tagging_vars = read_yaml(filepath)[
"output_vars"]
40 vu.add_collection(flavor_tagging_vars,
'flavor_tagging')
43def set_FT_pid_aliases(type="Current", exp_type="Belle2"):
45 This function adds the pid aliases needed by the flavor tagger.
49 from variables
import variables
as va
52 def set_alias(alias, var):
return va.addAlias(alias, var)
55 filepath = basf2.find_file(
'data/analysis/pid_alias.yaml')
56 alias_list = read_yaml(filepath)[type]
58 for alias, var
in alias_list[
"Common"].items():
60 for alias, var
in alias_list[exp_type].items():
64def set_GFlat_aliases(categories, usePIDNN):
66 This function adds aliases for the GNN-based flavor tagger.
69 from variables
import variables
as va
72 va.addAlias(
'qrCombined_bit',
'(qrCombined+1)/2')
73 alias_list = [
'qrCombined_bit']
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',
98 for rank
in range(1, 17):
100 for cat
in categories:
101 listName = get_available_categories()[cat].particleList
102 varName = f
'QpTrack({listName}, isRightCategory({cat}), isRightCategory({cat}))'
104 varWithRank = f
'ifNANgiveX(getVariableByRank(pi+:inRoe, FT_p, {varName}, {rank}), 0)'
105 aliasWithRank = f
'{cat}_rank{rank}'
107 va.addAlias(aliasWithRank, varWithRank)
108 alias_list.append(aliasWithRank)
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}'
114 va.addAlias(aliasWithRank, varWithRank)
115 alias_list.append(aliasWithRank)
120def set_masked_vars(mask_name="all"):
122 Returns variables with a certain mask
125 from variables
import variables
as va
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})',
134 for key, value
in masked_var_dict.items():
135 va.addAlias(key, value)
137 return masked_var_dict
140def set_and_get_aliases(var_dict):
142 Sets aliases from a dict of alias, var pair and returns a list of aliases.
144 from variables
import variables
as va
146 for key, value
in var_dict.items():
147 va.addAlias(key, value)
149 return list(var_dict.keys())