Belle II Software light-2406-ragdoll
create_data_sample.py
1#!/usr/bin/env python3
2
3
10
11# Create a simple data sample with some variables
12
13import basf2 as b2
14import modularAnalysis as ma
15import vertex as vx
16variables = ['p', 'pt', 'pz', 'phi',
17 'daughter(0, p)', 'daughter(0, pz)', 'daughter(0, pt)', 'daughter(0, phi)',
18 'daughter(1, p)', 'daughter(1, pz)', 'daughter(1, pt)', 'daughter(1, phi)',
19 'daughter(2, p)', 'daughter(2, pz)', 'daughter(2, pt)', 'daughter(2, phi)',
20 'chiProb', 'dr', 'dz', 'dphi',
21 'daughter(0, dr)', 'daughter(1, dr)', 'daughter(0, dz)', 'daughter(1, dz)',
22 'daughter(0, dphi)', 'daughter(1, dphi)',
23 'daughter(0, chiProb)', 'daughter(1, chiProb)', 'daughter(2, chiProb)',
24 'daughter(0, kaonID)', 'daughter(0, pionID)', 'daughter(1, kaonID)', 'daughter(1, pionID)',
25 'daughterAngle(0, 1)', 'daughterAngle(0, 2)', 'daughterAngle(1, 2)',
26 'daughter(2, daughter(0, E))', 'daughter(2, daughter(1, E))',
27 'daughter(2, daughter(0, clusterTiming))', 'daughter(2, daughter(1, clusterTiming))',
28 'daughter(2, daughter(0, clusterE9E25))', 'daughter(2, daughter(1, clusterE9E25))',
29 'daughter(2, daughter(0, minC2TDist))', 'daughter(2, daughter(1, minC2TDist))',
30 'daughterInvM(0, 1)', 'daughterInvM(0, 2)', 'daughterInvM(1, 2)']
31spectators = ['isSignal', 'M']
32
33
34def reconstruction_path(inputfiles):
35 path = b2.create_path()
36 ma.inputMdstList(inputfiles, path=path)
37 ma.fillParticleLists([('K-', 'kaonID > 0.5'), ('pi+', 'pionID > 0.5'),
38 ('gamma', '[[clusterReg == 1 and E > 0.10] or [clusterReg == 2 and E > 0.09] or '
39 '[clusterReg == 3 and E > 0.16]] and abs(clusterTiming) < 20 and clusterE9E25 > 0.7'
40 ' and minC2TDist > 35')],
41 path=path)
42 ma.reconstructDecay('pi0 -> gamma gamma', '0.1 < M < 1.6', path=path)
43 vx.kFit('pi0', 0.1, fit_type='massvertex', path=path)
44 ma.reconstructDecay('D0 -> K- pi+ pi0', '1.8 < M < 1.9', path=path)
45 vx.kFit('D0', 0.1, path=path)
46 ma.applyCuts('D0', '1.7 < M < 1.9', path=path)
47 ma.matchMCTruth('D0', path=path)
48 ma.applyCuts('D0', 'isNAN(isSignal) == False', path=path)
49 return path
50
51
52if __name__ == "__main__":
53
54 # Create a train, test and validation sample with different MC files
55 path = reconstruction_path([b2.find_file('ccbar_sample_to_train.root', 'examples', False)])
56 ma.variablesToNtuple('D0', variables + spectators, filename='train.root', treename='tree', path=path)
57 b2.process(path, 100000)
58
59 path = reconstruction_path([b2.find_file('ccbar_sample_to_test.root', 'examples', False)])
60 ma.variablesToNtuple('D0', variables + spectators, filename='test.root', treename='tree', path=path)
61 b2.process(path, 100000)
62
63 path = reconstruction_path([b2.find_file('ccbar_sample_to_test.root', 'examples', False)])
64 ma.variablesToNtuple('D0', variables + spectators, filename='validation.root', treename='tree', path=path)
65 b2.process(path, 100000)