Belle II Software light-2406-ragdoll
B2A702-ContinuumSuppression_MVATrain.py
1#!/usr/bin/env python3
2
3
10
11
25
26import basf2 as b2
27import basf2_mva
28import subprocess
29
30if __name__ == "__main__":
31
32 # Note that the target variable 'isNotContinuum' needs to be
33 # saved in your train.root and test.root files, along with the
34 # trainingVars, listed again here (see B2A701).
35 import os
36 if not os.getenv('BELLE2_EXAMPLES_DATA_DIR'):
37 b2.B2FATAL("You need the example data installed. Run `b2install-data example` in terminal for it.")
38
39 # Use this path to run over Bd_KsPi0 reconstructed signal and qqbar skims.
40 path = os.getenv('BELLE2_EXAMPLES_DATA_DIR')+'/mva/'
41
42 train_data = path + 'train.root'
43 test_data = path + 'test.root'
44 apply_signal_data = path + 'apply_signal.root'
45 apply_qqbar_data = path + 'apply_qqbar.root'
46
47 # Define the variables for training.
48 # For details, please see the Continuum suppression section at https://software.belle2.org
49 # Note that KSFWVariables takes the optional additional argument FS1, to return the variables calculated from the
50 # signal-B final state particles.
51 # CleoCone also takes the optional additional argument ROE, to return the cones calculated from ROE particles only.
52 trainVars = [
53 'R2',
54 'thrustBm',
55 'thrustOm',
56 'cosTBTO',
57 'cosTBz',
58 'KSFWVariables(et)',
59 'KSFWVariables(mm2)',
60 'KSFWVariables(hso00)',
61 'KSFWVariables(hso01)',
62 'KSFWVariables(hso02)',
63 'KSFWVariables(hso03)',
64 'KSFWVariables(hso04)',
65 'KSFWVariables(hso10)',
66 'KSFWVariables(hso12)',
67 'KSFWVariables(hso14)',
68 'KSFWVariables(hso20)',
69 'KSFWVariables(hso22)',
70 'KSFWVariables(hso24)',
71 'KSFWVariables(hoo0)',
72 'KSFWVariables(hoo1)',
73 'KSFWVariables(hoo2)',
74 'KSFWVariables(hoo3)',
75 'KSFWVariables(hoo4)',
76 'CleoConeCS(1)',
77 'CleoConeCS(2)',
78 'CleoConeCS(3)',
79 'CleoConeCS(4)',
80 'CleoConeCS(5)',
81 'CleoConeCS(6)',
82 'CleoConeCS(7)',
83 'CleoConeCS(8)',
84 'CleoConeCS(9)'
85 ]
86
87 general_options = basf2_mva.GeneralOptions()
88 general_options.m_datafiles = basf2_mva.vector(train_data)
89 general_options.m_treename = "tree"
90 general_options.m_identifier = "MVAFastBDT.root"
91 general_options.m_variables = basf2_mva.vector(*trainVars)
92 general_options.m_target_variable = "isNotContinuumEvent"
93 fastbdt_options = basf2_mva.FastBDTOptions()
94
95 # Train a MVA method and store the weightfile (MVAFastBDT.root) locally.
96 basf2_mva.teacher(general_options, fastbdt_options)
97
98 # Evaluate training.
99 subprocess.call('basf2_mva_evaluate.py '
100 ' -train ' + train_data +
101 ' -data ' + test_data +
102 ' -id ' + 'MVAFastBDT.root' +
103 ' --output qqbarSuppressionEvaluation.pdf',
104 shell=True
105 )
106
107 # If you're only interested in the network output distribution, then
108 # comment these in to apply the trained methods on an independent sample
109 # (produced in B2A701 if you ran with the `apply_signal` and `apply_qqbar` options).
110 #
111 basf2_mva.expert(
112 basf2_mva.vector('MVAFastBDT.root'),
113 basf2_mva.vector(apply_signal_data),
114 'tree',
115 'MVAExpert_signal.root')
116 basf2_mva.expert(
117 basf2_mva.vector('MVAFastBDT.root'),
118 basf2_mva.vector(apply_qqbar_data),
119 'tree',
120 'MVAExpert_qqbar.root')