Belle II Software  release-05-01-25
B_specific_train.py
1 #!/usr/bin/env python3
2 
3 # William Sutcliffe 2019
4 #
5 # Steering file to train the specfic FEI on Belle II MC, but it can be also easily adapted for converted Belle MC.
6 # This steering file is called several times (so-called stages) during the training process of the FEI.
7 # For reference see Confluence and Thomas Keck's PhD thesis.
8 #
9 # Please adapt for your signal channel. Note that a large amount of MC is needed to train the specific FEI.
10 # I usually use 100 million of signal events for each final state, mixed and charged MC.
11 # This example is for hadronic tagging.
12 
13 import fei
14 import basf2 as b2
15 import modularAnalysis as ma
16 
17 # Create path
18 path = b2.create_path()
19 
20 # Load input ROOT file
21 ma.inputMdst(environmentType='default',
22  filename=b2.find_file('mdst12.root', 'validation', False),
23  path=path)
24 
25 # Max 12 tracks per event - this avoids much computing time.
26 empty_path = b2.create_path()
27 skimfilter = b2.register_module('VariableToReturnValue')
28 skimfilter.param('variable', 'nCleanedTracks(dr < 2 and abs(dz) < 4)')
29 skimfilter.if_value('>12', empty_path, b2.AfterConditionPath.END)
30 path.add_module(skimfilter)
31 
32 # Signal side reconstruction
33 ma.fillParticleList('mu+', 'muonID > 0.8 and dr < 2 and abs(dz) < 4', writeOut=True, path=path)
34 ma.fillParticleList('e+', 'electronID > 0.8 and dr < 2 and abs(dz) < 4', writeOut=True, path=path)
35 ma.fillParticleList(
36  'gamma',
37  '[[clusterReg == 1 and E > 0.10] or [clusterReg == 2 and E > 0.09] or [clusterReg == 3 and E > 0.16]]',
38  writeOut=True,
39  path=path)
40 ma.reconstructDecay(
41  'B+:sig_e -> gamma e+',
42  '1.000 < M < 6.000 and cos(useRestFrame(daughterAngle(0, 1))) < 0.6',
43  dmID=1,
44  writeOut=True,
45  path=path)
46 ma.reconstructDecay(
47  'B+:sig_mu -> gamma mu+',
48  '1.000 < M < 6.000 and cos(useRestFrame(daughterAngle(0, 1))) < 0.6',
49  dmID=2,
50  writeOut=True,
51  path=path)
52 ma.copyLists('B+:sig', ['B+:sig_e', 'B+:sig_mu'], writeOut=True, path=path)
53 ma.looseMCTruth('B+:sig', path=path)
54 ma.rankByHighest('B+:sig', 'daughter(0,E)', outputVariable='PhotonCandidateRank', path=path)
55 ma.buildRestOfEvent('B+:sig', path=path)
56 clean_roe_mask = (
57  'CleanROE',
58  'dr < 2 and abs(dz) < 4',
59  'clusterE9E25 > 0.9 and clusterTiming < 50 and E > 0.09 and trackMatchType==0')
60 ma.appendROEMasks('B+:sig', [clean_roe_mask], path=path)
61 ma.applyCuts('B+:sig', 'roeDeltae(CleanROE) < 2.0 and roeMbc(CleanROE) > 4.8', path=path)
62 
63 skimfilter = b2.register_module('SkimFilter')
64 skimfilter.param('particleLists', ['B+:sig'])
65 empty_path = b2.create_path()
66 skimfilter.if_value('=0', empty_path, b2.AfterConditionPath.END)
67 path.add_module(skimfilter)
68 
69 # Prepare list for the training.
70 path.add_module('MCDecayFinder', decayString='B+ -> e+ nu_e gamma', listName='B+:FEIMC_e', writeOut=True)
71 path.add_module('MCDecayFinder', decayString='B+ -> mu+ nu_mu gamma', listName='B+:FEIMC_mu', writeOut=True)
72 ma.copyLists('B+:FEIMC', ['B+:FEIMC_e', 'B+:FEIMC_mu'], writeOut=True, path=path)
73 
74 
75 # We want the FEI to be only trained on a correctly reconstruced signal side and on wrongly reconstructed background.
76 isSignal = 'isSignalAcceptMissingNeutrino'
77 signalMC = 'eventCached(countInList(B+:FEIMC))'
78 cut = '[[{mc} > 0 and {sig} == 1] or [{mc} == 0 and {sig} != 1]]'.format(mc=signalMC, sig=isSignal)
79 ma.applyCuts('B+:sig', cut, path=path)
80 
81 # Set up FEI configuration specifying the FEI prefix
82 fei_tag = 'my_specFEI'
83 belle_particles = fei.get_default_channels(KLong=False,
84  chargedB=True,
85  neutralB=True,
86  semileptonic=False,
87  B_extra_cut='nRemainingTracksInEvent <= 3',
88  specific=True)
89 
90 # Get FEI path
91 configuration = fei.config.FeiConfiguration(prefix=fei_tag, training=True, monitor=False, cache=-1)
92 
93 
94 # Add FEI path to the path to be processed
95 feistate = fei.get_path(belle_particles, configuration)
96 
97 # FEI training
98 if feistate.stage == 0:
99  # Write out the rest of event, we train only on the rest of event of our signal side.
100  # This is the main difference compared to the generic FEI.
101  rO = b2.register_module('RootOutput')
102  rO.set_name('ROE_RootOutput')
103  rO.param('additionalBranchNames', ['RestOfEvent'])
104  feistate.path.add_module(rO)
105  roe_path = b2.create_path()
106  cond_module = b2.register_module('SignalSideParticleFilter')
107  cond_module.param('particleLists', ['B+:sig'])
108  cond_module.if_true(feistate.path, b2.AfterConditionPath.END)
109  roe_path.add_module(cond_module)
110  path.for_each('RestOfEvent', 'RestOfEvents', roe_path)
111 else:
112  # After stage 0, the training is done only on the written out rest of event.
113  path = b2.create_path()
114  ma.inputMdstList('default', [], path)
115  path.add_path(feistate.path)
116  r1 = b2.register_module('RootOutput')
117  r1.set_name('ROE_RootOutput')
118  r1.param('additionalBranchNames', ['RestOfEvent'])
119  path.add_module(r1)
120 
121 
122 # Process 100 events
123 b2.process(path, max_event=100)