1#!/usr/bin/env python3
2
3import sys
4import basf2 as b2
5import modularAnalysis as ma
6import stdV0s
7import flavorTagger as ft
8from variables import variables as vm # shorthand for VariableManager
9import variables.collections as vc
10import variables.utils as vu
11
12# get input file number from the command line
13filenumber = sys.argv[1]
14
15# set analysis global tag (needed for flavor tagging)
16b2.conditions.prepend_globaltag("analysis_tools_release-04-02")
17
18# create path
19main = b2.Path()
20
21# load input data from mdst/udst file
22ma.inputMdstList(
23 environmentType="default",
24 filelist=[b2.find_file(f"starterkit/2021/1111540100_eph3_BGx0_{filenumber}.root", "examples")],
25 path=main,
26)
27
28# fill final state particle lists
29ma.fillParticleList(
30 "e+:uncorrected",
31 "electronID > 0.1 and dr < 0.5 and abs(dz) < 2 and thetaInCDCAcceptance",
32 path=main,
33)
34stdV0s.stdKshorts(path=main)
35
36# apply Bremsstrahlung correction to electrons
37vm.addAlias(
38 "goodFWDGamma", "passesCut(clusterReg == 1 and clusterE > 0.075)"
39)
40vm.addAlias(
41 "goodBRLGamma", "passesCut(clusterReg == 2 and clusterE > 0.05)"
42)
43vm.addAlias(
44 "goodBWDGamma", "passesCut(clusterReg == 3 and clusterE > 0.1)"
45)
46vm.addAlias(
47 "goodGamma", "passesCut(goodFWDGamma or goodBRLGamma or goodBWDGamma)"
48)
49ma.fillParticleList("gamma:brems", "goodGamma", path=main)
50ma.correctBrems("e+:corrected", "e+:uncorrected", "gamma:brems", path=main)
51vm.addAlias("isBremsCorrected", "extraInfo(bremsCorrected)")
52
53# combine final state particles to form composite particles
54ma.reconstructDecay(
55 "J/psi:ee -> e+:corrected e-:corrected ?addbrems",
56 cut="dM < 0.11",
57 path=main,
58)
59
60# combine J/psi and KS candidates to form B0 candidates
61ma.reconstructDecay(
62 "B0 -> J/psi:ee K_S0:merged",
63 cut="Mbc > 5.2 and abs(deltaE) < 0.15",
64 path=main,
65)
66
67# match reconstructed with MC particles
68ma.matchMCTruth("B0", path=main)
69
70# build the rest of the event
71ma.buildRestOfEvent("B0", fillWithMostLikely=True, path=main)
72track_based_cuts = "thetaInCDCAcceptance and pt > 0.075 and dr < 5 and abs(dz) < 10"
73ecl_based_cuts = "thetaInCDCAcceptance and E > 0.05"
74roe_mask = ("my_mask", track_based_cuts, ecl_based_cuts)
75ma.appendROEMasks("B0", [roe_mask], path=main)
76
77# call flavor tagging
78ft.flavorTagger(["B0"], path=main)
79
80# perform best candidate selection
81b2.set_random_seed("Belle II StarterKit")
82ma.rankByHighest("B0", variable="random", numBest=1, path=main)
83
84# Create list of variables to save into the output file
85b_vars = []
86
87standard_vars = vc.kinematics + vc.mc_kinematics + vc.mc_truth
88b_vars += vc.deltae_mbc
89b_vars += standard_vars
90
91# ROE variables
92roe_kinematics = ["roeE()", "roeM()", "roeP()", "roeMbc()", "roeDeltae()"]
93roe_multiplicities = [
94 "nROE_Charged()",
95 "nROE_Photons()",
96 "nROE_NeutralHadrons()",
97]
98b_vars += roe_kinematics + roe_multiplicities
99# Let's also add a version of the ROE variables that includes the mask:
100for roe_variable in roe_kinematics + roe_multiplicities:
101 # e.g. instead of 'roeE()' (no mask) we want 'roeE(my_mask)'
102 roe_variable_with_mask = roe_variable.replace("()", "(my_mask)")
103 b_vars.append(roe_variable_with_mask)
104
105b_vars += ft.flavor_tagging
106
107# Variables for final states (electrons, positrons, pions)
108fs_vars = vc.pid + vc.track + vc.track_hits + standard_vars
109b_vars += vu.create_aliases_for_selected(
110 fs_vars + ["isBremsCorrected"],
111 "B0 -> [J/psi -> ^e+ ^e-] K_S0",
112 prefix=["ep", "em"],
113)
114b_vars += vu.create_aliases_for_selected(
115 fs_vars, "B0 -> J/psi [K_S0 -> ^pi+ ^pi-]", prefix=["pip", "pim"]
116)
117# Variables for J/Psi, KS
118jpsi_ks_vars = vc.inv_mass + standard_vars
119b_vars += vu.create_aliases_for_selected(jpsi_ks_vars, "B0 -> ^J/psi ^K_S0")
120# Add the J/Psi mass calculated with uncorrected electrons:
121vm.addAlias(
122 "Jpsi_M_uncorrected", "daughter(0, daughterCombination(M,0:0,1:0))"
123)
124b_vars += ["Jpsi_M_uncorrected"]
125# Also add kinematic variables boosted to the center of mass frame (CMS)
126# for all particles
127cmskinematics = vu.create_aliases(
128 vc.kinematics, "useCMSFrame({variable})", "CMS"
129)
130b_vars += vu.create_aliases_for_selected(
131 cmskinematics, "^B0 -> [^J/psi -> ^e+ ^e-] [^K_S0 -> ^pi+ ^pi-]"
132)
133
134vm.addAlias(
135 "withBremsCorrection",
136 "passesCut(passesCut(ep_isBremsCorrected == 1) or passesCut(em_isBremsCorrected == 1))",
137)
138b_vars += ["withBremsCorrection"]
139
140# Save variables to an output file (ntuple)
141ma.variablesToNtuple(
142 "B0",
143 variables=b_vars,
144 filename="Bd2JpsiKS.root",
145 treename="tree",
146 path=main,
147)
148
149# Start the event loop (actually start processing things)
150b2.process(main)
151
152# print out the summary
153print(b2.statistics)