Belle II Software  release-05-01-25
all_classifiers.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 # Thomas Keck 2016
5 
6 import basf2_mva
7 
8 import os
9 import tempfile
10 import subprocess
11 import sys
12 import shutil
13 
14 variables = ['p', 'pz', 'daughter(0, p)', 'daughter(0, pz)', 'daughter(1, p)', 'daughter(1, pz)',
15  'chiProb', 'dr', 'dz', 'daughter(0, dr)', 'daughter(1, dr)', 'daughter(0, chiProb)', 'daughter(1, chiProb)',
16  'daughter(0, kaonID)', 'daughter(0, pionID)', 'daughterAngle(0, 1)']
17 
18 if __name__ == "__main__":
19 
20  # Skip test if files are not available
21  if not (os.path.isfile('train.root') and os.path.isfile('test.root')):
22  sys.exit(0)
23 
24  general_options = basf2_mva.GeneralOptions()
25  general_options.m_datafiles = basf2_mva.vector("train.root")
26  general_options.m_treename = "tree"
27  general_options.m_variables = basf2_mva.vector(*variables)
28  general_options.m_target_variable = "isSignal"
29 
30  methods = [
31  ('Trivial.xml', basf2_mva.TrivialOptions()),
32  ('Python.xml', basf2_mva.PythonOptions()),
33  ('FastBDT.xml', basf2_mva.FastBDTOptions()),
34  ('TMVAClassification.xml', basf2_mva.TMVAOptionsClassification()),
35  ('FANN.xml', basf2_mva.FANNOptions()),
36  ]
37 
38  olddir = os.getcwd()
39  with tempfile.TemporaryDirectory() as tempdir:
40  os.symlink(os.path.abspath('train.root'), tempdir + '/' + os.path.basename('train.root'))
41  os.symlink(os.path.abspath('test.root'), tempdir + '/' + os.path.basename('test.root'))
42  os.chdir(tempdir)
43 
44  for identifier, specific_options in methods:
45  general_options.m_identifier = identifier
46  basf2_mva.teacher(general_options, specific_options)
47 
48  basf2_mva.expert(basf2_mva.vector(*[i for i, _ in methods]),
49  basf2_mva.vector('train.root'), 'tree', 'expert.root')
50 
51  subprocess.call('basf2_mva_evaluate.py -o latex.pdf -train train.root -data test.root -i ' +
52  ' '.join([i for i, _ in methods]), shell=True)
53 
54  os.chdir(olddir)
55  shutil.copyfile(tempdir + '/latex.pdf', 'latex.pdf')