Belle II Software light-2406-ragdoll
tmva_bdt.py
1#!/usr/bin/env python3
2
3
10
11import basf2_mva
12import basf2_mva_util
13import time
14
15if __name__ == "__main__":
16 from basf2 import conditions, find_file
17 # NOTE: do not use testing payloads in production! Any results obtained like this WILL NOT BE PUBLISHED
18 conditions.testing_payloads = [
19 'localdb/database.txt'
20 ]
21 train_file = find_file("mva/train_D0toKpipi.root", "examples")
22 test_file = find_file("mva/test_D0toKpipi.root", "examples")
23
24 training_data = basf2_mva.vector(train_file)
25 testing_data = basf2_mva.vector(test_file)
26
27 variables = ['M', 'p', 'pt', 'pz',
28 'daughter(0, p)', 'daughter(0, pz)', 'daughter(0, pt)',
29 'daughter(1, p)', 'daughter(1, pz)', 'daughter(1, pt)',
30 'daughter(2, p)', 'daughter(2, pz)', 'daughter(2, pt)',
31 'chiProb', 'dr', 'dz',
32 'daughter(0, dr)', 'daughter(1, dr)',
33 'daughter(0, dz)', 'daughter(1, dz)',
34 'daughter(0, chiProb)', 'daughter(1, chiProb)', 'daughter(2, chiProb)',
35 'daughter(0, kaonID)', 'daughter(0, pionID)',
36 'daughterInvM(0, 1)', 'daughterInvM(0, 2)', 'daughterInvM(1, 2)']
37
38 # Train a MVA method and directly upload it to the database
39 general_options = basf2_mva.GeneralOptions()
40 general_options.m_datafiles = training_data
41 general_options.m_treename = "tree"
42 general_options.m_identifier = "TMVA"
43 general_options.m_variables = basf2_mva.vector(*variables)
44 general_options.m_target_variable = "isSignal"
45
46 tmva_bdt_options = basf2_mva.TMVAOptionsClassification()
47 tmva_bdt_options.m_config = ("!H:!V:CreateMVAPdfs:NTrees=100:BoostType=Grad:Shrinkage=0.2:UseBaggedBoost:"
48 "BaggedSampleFraction=0.5:nCuts=1024:MaxDepth=3:IgnoreNegWeightsInTraining")
49
50 training_start = time.time()
51 basf2_mva.teacher(general_options, tmva_bdt_options)
52 training_stop = time.time()
53
54 training_time = training_stop - training_start
55 method = basf2_mva_util.Method(general_options.m_identifier)
56
57 inference_start = time.time()
58 p, t = method.apply_expert(testing_data, general_options.m_treename)
59 inference_stop = time.time()
60 inference_time = inference_stop - inference_start
61
63 print("TMVA", training_time, inference_time, auc)
def calculate_auc_efficiency_vs_background_retention(p, t, w=None)