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