Belle II Software  release-05-01-25
test_fei_default_channels.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 import unittest
5 import unittest.mock
6 import os
7 import tempfile
8 import atexit
9 import shutil
10 
11 import fei.default_channels
12 
13 import b2bii
14 
15 # @cond
16 
17 
18 class TestGetDefaultChannels(unittest.TestCase):
19  def test_get_default(self):
20  particles = fei.default_channels.get_default_channels()
21  self.assertEqual([p.identifier for p in particles],
22  ['pi+:generic', 'K+:generic', 'p+:generic', 'mu+:generic', 'e+:generic', 'gamma:generic',
23  'pi0:generic', 'K_S0:generic', 'Lambda0:generic', 'Sigma+:generic', 'J/psi:generic',
24  'D0:generic', 'D+:generic', 'D_s+:generic', 'Lambda_c+:generic',
25  'D*0:generic', 'D*+:generic', 'D_s*+:generic',
26  'B0:generic', 'B+:generic',
27  'D0:semileptonic', 'D+:semileptonic', 'D*0:semileptonic', 'D*+:semileptonic',
28  'B0:semileptonic', 'B+:semileptonic'])
29 
30  def test_get_hadronic(self):
31  particles = fei.default_channels.get_default_channels(hadronic=True, semileptonic=False, baryonic=False)
32  self.assertEqual([p.identifier for p in particles],
33  ['pi+:generic', 'K+:generic', 'mu+:generic', 'e+:generic', 'gamma:generic', 'pi0:generic', 'K_S0:generic',
34  'J/psi:generic', 'D0:generic', 'D+:generic', 'D_s+:generic', 'D*0:generic', 'D*+:generic',
35  'D_s*+:generic', 'B0:generic', 'B+:generic'])
36 
37  def test_get_semileptonic(self):
38  particles = fei.default_channels.get_default_channels(hadronic=False, semileptonic=True, baryonic=False)
39  self.assertEqual([p.identifier for p in particles],
40  ['pi+:generic', 'K+:generic', 'mu+:generic', 'e+:generic', 'gamma:generic', 'pi0:generic', 'K_S0:generic',
41  'J/psi:generic', 'D0:generic', 'D+:generic', 'D_s+:generic', 'D*0:generic', 'D*+:generic',
42  'D_s*+:generic', 'D0:semileptonic', 'D+:semileptonic', 'D*0:semileptonic', 'D*+:semileptonic',
43  'B0:semileptonic', 'B+:semileptonic'])
44 
45  def test_get_charged(self):
46  particles = fei.default_channels.get_default_channels(chargedB=True, neutralB=False, baryonic=False)
47  self.assertEqual([p.identifier for p in particles],
48  ['pi+:generic', 'K+:generic', 'mu+:generic', 'e+:generic', 'gamma:generic', 'pi0:generic', 'K_S0:generic',
49  'J/psi:generic', 'D0:generic', 'D+:generic', 'D_s+:generic', 'D*0:generic', 'D*+:generic',
50  'D_s*+:generic', 'B+:generic', 'D0:semileptonic', 'D+:semileptonic', 'D*0:semileptonic',
51  'D*+:semileptonic', 'B+:semileptonic'])
52 
53  def test_get_neutral(self):
54  particles = fei.default_channels.get_default_channels(chargedB=False, neutralB=True, baryonic=False)
55  self.assertEqual([p.identifier for p in particles],
56  ['pi+:generic', 'K+:generic', 'mu+:generic', 'e+:generic', 'gamma:generic', 'pi0:generic', 'K_S0:generic',
57  'J/psi:generic', 'D0:generic', 'D+:generic', 'D_s+:generic', 'D*0:generic', 'D*+:generic',
58  'D_s*+:generic', 'B0:generic', 'D0:semileptonic', 'D+:semileptonic', 'D*0:semileptonic',
59  'D*+:semileptonic', 'B0:semileptonic'])
60 
61  def test_get_klong(self):
62  particles = fei.default_channels.get_default_channels(hadronic=False, semileptonic=False, KLong=True, baryonic=False)
63  self.assertEqual([p.identifier for p in particles],
64  ['pi+:generic', 'K+:generic', 'mu+:generic', 'e+:generic', 'gamma:generic', 'pi0:generic', 'K_S0:generic',
65  'J/psi:generic', 'D0:generic', 'D+:generic', 'D_s+:generic', 'D*0:generic', 'D*+:generic',
66  'D_s*+:generic', 'K_L0:generic', 'D0:KL', 'D+:KL', 'D_s+:KL', 'D*0:KL', 'D*+:KL', 'D_s*+:KL',
67  'B0:KL', 'B+:KL'])
68 
69  def test_B_extra_cut(self):
70  particles = fei.default_channels.get_default_channels(B_extra_cut='nRemainingTracksInROE == 0', KLong=True, baryonic=False)
71  self.assertEqual([p.identifier for p in particles],
72  ['pi+:generic', 'K+:generic', 'mu+:generic', 'e+:generic', 'gamma:generic', 'pi0:generic', 'K_S0:generic',
73  'J/psi:generic', 'D0:generic', 'D+:generic', 'D_s+:generic', 'D*0:generic', 'D*+:generic',
74  'D_s*+:generic', 'B0:generic', 'B+:generic',
75  'K_L0:generic', 'D0:KL', 'D+:KL', 'D_s+:KL', 'D*0:KL', 'D*+:KL', 'D_s*+:KL', 'B0:KL', 'B+:KL',
76  'D0:semileptonic', 'D+:semileptonic', 'D*0:semileptonic', 'D*+:semileptonic',
77  'B0:semileptonic', 'B+:semileptonic',
78  ])
79  self.assertEqual(particles[14].preCutConfig.userCut, 'Mbc > 5.2 and abs(deltaE) < 0.5 and [nRemainingTracksInROE == 0]')
80  self.assertEqual(particles[15].preCutConfig.userCut, 'Mbc > 5.2 and abs(deltaE) < 0.5 and [nRemainingTracksInROE == 0]')
81  self.assertEqual(particles[23].preCutConfig.userCut, 'nRemainingTracksInROE == 0')
82  self.assertEqual(particles[24].preCutConfig.userCut, 'nRemainingTracksInROE == 0')
83  self.assertEqual(particles[29].preCutConfig.userCut, 'nRemainingTracksInROE == 0')
84  self.assertEqual(particles[30].preCutConfig.userCut, 'nRemainingTracksInROE == 0')
85 
86  def test_get_specific(self):
87  particles = fei.default_channels.get_default_channels(specific=True, KLong=True, baryonic=False)
88  self.assertEqual([p.identifier for p in particles],
89  ['pi+:generic', 'K+:generic', 'mu+:generic', 'e+:generic', 'gamma:generic', 'pi0:generic', 'K_S0:generic',
90  'J/psi:generic', 'D0:generic', 'D+:generic', 'D_s+:generic', 'D*0:generic', 'D*+:generic',
91  'D_s*+:generic', 'B0:generic', 'B+:generic',
92  'K_L0:generic', 'D0:KL', 'D+:KL', 'D_s+:KL', 'D*0:KL', 'D*+:KL', 'D_s*+:KL', 'B0:KL', 'B+:KL',
93  'D0:semileptonic', 'D+:semileptonic', 'D*0:semileptonic', 'D*+:semileptonic',
94  'B0:semileptonic', 'B+:semileptonic'])
95  self.assertEqual(particles[0].preCutConfig.userCut, '[dr < 2] and [abs(dz) < 4] and isInRestOfEvent > 0.5')
96  self.assertEqual(particles[1].preCutConfig.userCut, '[dr < 2] and [abs(dz) < 4] and isInRestOfEvent > 0.5')
97  self.assertEqual(particles[2].preCutConfig.userCut, '[dr < 2] and [abs(dz) < 4] and isInRestOfEvent > 0.5')
98  self.assertEqual(particles[3].preCutConfig.userCut, '[dr < 2] and [abs(dz) < 4] and isInRestOfEvent > 0.5')
99  self.assertEqual(particles[4].preCutConfig.userCut, '[[clusterReg == 1 and E > 0.10] or [clusterReg == 2 and E > 0.09]'
100  ' or [clusterReg == 3 and E > 0.16]] and isInRestOfEvent > 0.5')
101  self.assertEqual(particles[6].preCutConfig.userCut, '0.4 < M < 0.6')
102  self.assertEqual(particles[6].channels[2].preCutConfig.userCut, '0.4 < M < 0.6 and isInRestOfEvent > 0.5')
103  self.assertEqual(particles[16].preCutConfig.userCut, 'isInRestOfEvent > 0.5')
104 
105  def test_get_specific_converted(self):
106  b2bii.setB2BII()
107  particles = fei.default_channels.get_default_channels(specific=True, KLong=True, baryonic=False)
108  self.assertEqual([p.identifier for p in particles],
109  ['pi+:generic', 'K+:generic', 'mu+:generic', 'e+:generic', 'gamma:generic', 'pi0:generic', 'K_S0:generic',
110  'J/psi:generic', 'D0:generic', 'D+:generic', 'D_s+:generic', 'D*0:generic', 'D*+:generic',
111  'D_s*+:generic', 'B0:generic', 'B+:generic',
112  'K_L0:generic', 'D0:KL', 'D+:KL', 'D_s+:KL', 'D*0:KL', 'D*+:KL', 'D_s*+:KL', 'B0:KL', 'B+:KL',
113  'D0:semileptonic', 'D+:semileptonic', 'D*0:semileptonic', 'D*+:semileptonic',
114  'B0:semileptonic', 'B+:semileptonic'])
115  self.assertEqual(particles[0].preCutConfig.userCut, '[dr < 2] and [abs(dz) < 4] and isInRestOfEvent > 0.5')
116  self.assertEqual(particles[1].preCutConfig.userCut, '[dr < 2] and [abs(dz) < 4] and isInRestOfEvent > 0.5')
117  self.assertEqual(particles[2].preCutConfig.userCut, '[dr < 2] and [abs(dz) < 4] and isInRestOfEvent > 0.5')
118  self.assertEqual(particles[3].preCutConfig.userCut, '[dr < 2] and [abs(dz) < 4] and isInRestOfEvent > 0.5')
119  self.assertEqual(particles[4].preCutConfig.userCut, 'goodBelleGamma == 1 and clusterBelleQuality == 0 '
120  'and isInRestOfEvent > 0.5')
121  self.assertEqual(particles[5].preCutConfig.userCut, '0.08 < InvM < 0.18 and isInRestOfEvent > 0.5')
122  self.assertEqual(particles[6].preCutConfig.userCut, '0.4 < M < 0.6 and isInRestOfEvent > 0.5')
123  self.assertEqual(particles[16].preCutConfig.userCut, 'isInRestOfEvent > 0.5')
124  b2bii.unsetB2BII()
125 
126  def test_get_converted(self):
127  b2bii.setB2BII()
128  particles = fei.default_channels.get_default_channels(baryonic=False)
129  self.assertEqual([p.identifier for p in particles],
130  ['pi+:generic', 'K+:generic', 'mu+:generic', 'e+:generic', 'gamma:generic', 'pi0:generic', 'K_S0:generic',
131  'J/psi:generic', 'D0:generic', 'D+:generic', 'D_s+:generic', 'D*0:generic', 'D*+:generic',
132  'D_s*+:generic', 'B0:generic', 'B+:generic',
133  'D0:semileptonic', 'D+:semileptonic', 'D*0:semileptonic', 'D*+:semileptonic',
134  'B0:semileptonic', 'B+:semileptonic'])
135 
136  chargedVariables = ['eIDBelle',
137  'atcPIDBelle(3,2)',
138  'atcPIDBelle(4,2)', 'atcPIDBelle(4,3)',
139  'muIDBelle',
140  'p', 'pt', 'pz', 'dr', 'dz', 'chiProb', 'extraInfo(preCut_rank)']
141  self.assertEqual(particles[0].preCutConfig.userCut, '[dr < 2] and [abs(dz) < 4]')
142  self.assertEqual(particles[0].mvaConfig.variables, chargedVariables)
143  self.assertEqual(particles[0].preCutConfig.bestCandidateVariable, 'atcPIDBelle(2,3)')
144  self.assertEqual(particles[1].preCutConfig.userCut, '[dr < 2] and [abs(dz) < 4]')
145  self.assertEqual(particles[1].mvaConfig.variables, chargedVariables)
146  self.assertEqual(particles[1].preCutConfig.bestCandidateVariable, 'atcPIDBelle(3,2)')
147  self.assertEqual(particles[2].preCutConfig.userCut, '[dr < 2] and [abs(dz) < 4]')
148  self.assertEqual(particles[2].mvaConfig.variables, chargedVariables)
149  self.assertEqual(particles[2].preCutConfig.bestCandidateVariable, 'muIDBelle')
150  self.assertEqual(particles[3].preCutConfig.userCut, '[dr < 2] and [abs(dz) < 4]')
151  self.assertEqual(particles[3].mvaConfig.variables, chargedVariables)
152  self.assertEqual(particles[3].preCutConfig.bestCandidateVariable, 'eIDBelle')
153  self.assertEqual(particles[4].preCutConfig.userCut, 'goodBelleGamma == 1 and clusterBelleQuality == 0')
154  self.assertEqual(particles[6].mvaConfig.variables, ['dr', 'dz', 'distance', 'significanceOfDistance', 'chiProb',
155  'M', 'abs(dM)',
156  'useCMSFrame(E)', 'daughterAngle(0,1)',
157  'cosAngleBetweenMomentumAndVertexVector',
158  'extraInfo(preCut_rank)', 'extraInfo(goodKs)', 'extraInfo(ksnbVLike)',
159  'extraInfo(ksnbNoLam)', 'extraInfo(ksnbStandard)'])
160  b2bii.unsetB2BII()
161 
162 
163 class TestGetFRChannels(unittest.TestCase):
164  def test_get_default(self):
165  particles = fei.default_channels.get_fr_channels()
166  self.assertEqual([p.identifier for p in particles],
167  ['pi+:generic', 'K+:generic', 'mu+:generic', 'e+:generic', 'gamma:generic', 'pi0:generic', 'K_S0:generic',
168  'J/psi:generic', 'D0:generic', 'D+:generic', 'D_s+:generic', 'D*0:generic', 'D*+:generic',
169  'D_s*+:generic', 'B0:generic', 'B+:generic'])
170  self.assertEqual(len(particles[0].channels), 1)
171  self.assertEqual(len(particles[1].channels), 1)
172  self.assertEqual(len(particles[2].channels), 1)
173  self.assertEqual(len(particles[3].channels), 1)
174  self.assertEqual(len(particles[4].channels), 2)
175  self.assertEqual(len(particles[5].channels), 1)
176  self.assertEqual(len(particles[6].channels), 3)
177  self.assertEqual(len(particles[7].channels), 2)
178  self.assertEqual(len(particles[8].channels), 10)
179  self.assertEqual(len(particles[9].channels), 7)
180  self.assertEqual(len(particles[10].channels), 8)
181  self.assertEqual(len(particles[11].channels), 2)
182  self.assertEqual(len(particles[12].channels), 2)
183  self.assertEqual(len(particles[13].channels), 1)
184  self.assertEqual(len(particles[14].channels), 15)
185  self.assertEqual(len(particles[15].channels), 17)
186 
187 
188 if __name__ == '__main__':
189  tempdir = tempfile.mkdtemp()
190  os.chdir(tempdir)
191  # main() never returns, so install exit handler to do our cleanup
192  atexit.register(shutil.rmtree, tempdir)
193  unittest.main()
194 
195 # @endcond
fei.default_channels
Definition: default_channels.py:1
b2bii.unsetB2BII
def unsetB2BII()
Definition: b2bii.py:17
b2bii.setB2BII
def setB2BII()
Definition: b2bii.py:13
fei.default_channels.get_fr_channels
def get_fr_channels(convertedFromBelle=False)
Definition: default_channels.py:1049
fei.default_channels.get_default_channels
def get_default_channels(B_extra_cut=None, hadronic=True, semileptonic=True, KLong=False, baryonic=True, chargedB=True, neutralB=True, specific=False, removeSLD=False)
Definition: default_channels.py:27