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