Belle II Software  release-05-02-19
create_data_sample.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 # Thomas Keck 2016
5 
6 # Create a simple data sample with some variables
7 
8 from basf2 import *
9 from modularAnalysis import *
10 variables = ['p', 'pt', 'pz', 'phi',
11  'daughter(0, p)', 'daughter(0, pz)', 'daughter(0, pt)', 'daughter(0, phi)',
12  'daughter(1, p)', 'daughter(1, pz)', 'daughter(1, pt)', 'daughter(1, phi)',
13  'daughter(2, p)', 'daughter(2, pz)', 'daughter(2, pt)', 'daughter(2, phi)',
14  'chiProb', 'dr', 'dz', 'dphi',
15  'daughter(0, dr)', 'daughter(1, dr)', 'daughter(0, dz)', 'daughter(1, dz)',
16  'daughter(0, dphi)', 'daughter(1, dphi)',
17  'daughter(0, chiProb)', 'daughter(1, chiProb)', 'daughter(2, chiProb)',
18  'daughter(0, kaonID)', 'daughter(0, pionID)', 'daughter(1, kaonID)', 'daughter(1, pionID)',
19  'daughterAngle(0, 1)', 'daughterAngle(0, 2)', 'daughterAngle(1, 2)',
20  'daughter(2, daughter(0, E))', 'daughter(2, daughter(1, E))',
21  'daughter(2, daughter(0, clusterTiming))', 'daughter(2, daughter(1, clusterTiming))',
22  'daughter(2, daughter(0, clusterE9E25))', 'daughter(2, daughter(1, clusterE9E25))',
23  'daughter(2, daughter(0, minC2HDist))', 'daughter(2, daughter(1, minC2HDist))',
24  'daughterInvariantMass(0, 1)', 'daughterInvariantMass(0, 2)', 'daughterInvariantMass(1, 2)']
25 spectators = ['isSignal', 'M']
26 
27 
28 def reconstruction_path(inputfiles):
29  path = create_path()
30  inputMdstList('MC7', inputfiles, path=path)
31  fillParticleLists([('K-', 'kaonID > 0.5'), ('pi+', 'pionID > 0.5'),
32  ('gamma', '[[clusterReg == 1 and E > 0.10] or [clusterReg == 2 and E > 0.09] or '
33  '[clusterReg == 3 and E > 0.16]] and abs(clusterTiming) < 20 and clusterE9E25 > 0.7'
34  ' and minC2HDist > 35')],
35  path=path)
36  reconstructDecay('pi0 -> gamma gamma', '0.1 < M < 1.6', path=path)
37  KFit('pi0', 0.1, fit_type='massvertex', path=path)
38  reconstructDecay('D0 -> K- pi+ pi0', '1.8 < M < 1.9', path=path)
39  KFit('D0', 0.1, path=path)
40  applyCuts('D0', '1.7 < M < 1.9', path=path)
41  matchMCTruth('D0', path=path)
42  return path
43 
44 
45 if __name__ == "__main__":
46 
47  # Create a train, test and validation sample with different MC files
48  # Add your root files here
49  f = 'mdst_002001_prod00000789_task00004203.root'
50  path = reconstruction_path([f])
51  variablesToNtuple('D0', variables + spectators, filename='train.root', treename='tree', path=path)
52  process(path)
53 
54  # Add your root files here
55  path = reconstruction_path([f])
56  variablesToNtuple('D0', variables + spectators, filename='test.root', treename='tree', path=path)
57  process(path)
58 
59  # Add your root files here
60  path = reconstruction_path([f])
61  variablesToNtuple('D0', variables + spectators, filename='validation.root', treename='tree', path=path)
62  process(path)
variablesToNtuple
Definition: variablesToNtuple.py:1