~maddevelopers/mg5amcnlo/2.9.4

« back to all changes in this revision

Viewing changes to tests/unit_tests/various/test_decay.py

pass to v2.0.0

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
################################################################################
 
2
#
 
3
# Copyright (c) 2009 The MadGraph Development team and Contributors
 
4
#
 
5
# This file is a part of the MadGraph 5 project, an application which 
 
6
# automatically generates Feynman diagrams and matrix elements for arbitrary
 
7
# high-energy processes in the Standard Model and beyond.
 
8
#
 
9
# It is subject to the MadGraph license which should accompany this 
 
10
# distribution.
 
11
#
 
12
# For more information, please visit: http://madgraph.phys.ucl.ac.be
 
13
#
 
14
################################################################################
 
15
"""Unit test Library for the objects in decay module."""
 
16
from __future__ import division
 
17
 
 
18
import copy
 
19
import os
 
20
import sys
 
21
import time
 
22
import math
 
23
import cmath
 
24
 
 
25
import tests.unit_tests as unittest
 
26
import madgraph.core.base_objects as base_objects
 
27
import models.import_ufo as import_ufo
 
28
import models.model_reader as model_reader
 
29
import madgraph.iolibs.save_model as save_model
 
30
import madgraph.iolibs.drawing_eps as drawing_eps
 
31
import madgraph.iolibs.import_v4 as import_v4
 
32
from madgraph import MG5DIR
 
33
import mg5decay.decay_objects as decay_objects
 
34
import tests.input_files.import_vertexlist as import_vertexlist
 
35
import madgraph.core.diagram_generation as diagram_generation
 
36
import madgraph.various.diagram_symmetry as diagram_symmetry
 
37
 
 
38
 
 
39
_file_path = os.path.split(os.path.dirname(os.path.realpath(__file__)))[0]
 
40
 
 
41
#===============================================================================
 
42
# DecayParticleTest
 
43
#===============================================================================
 
44
class Test_DecayParticle(unittest.TestCase):
 
45
    """Test class for the DecayParticle object"""
 
46
 
 
47
    mydict = {}
 
48
    mypart = None
 
49
    my_2bodyvertexlist = base_objects.VertexList()
 
50
    my_3bodyvertexlist = base_objects.VertexList()
 
51
    my_2bodyvertexlist_wrongini = base_objects.VertexList()
 
52
    my_3bodyvertexlist_wrongini = base_objects.VertexList()
 
53
 
 
54
    def setUp(self):
 
55
 
 
56
        #Import a model from my_testmodel
 
57
        self.sm_path = import_ufo.find_ufo_path('sm')
 
58
        self.my_testmodel_base = import_ufo.import_model(self.sm_path)
 
59
        self.my_testmodel = decay_objects.DecayModel(self.my_testmodel_base, 
 
60
                                                     True)
 
61
        param_path = os.path.join(_file_path,'../input_files/param_card_sm.dat')
 
62
        self.my_testmodel.read_param_card(param_path)
 
63
        #print len(self.my_testmodel_base.get('interactions')), 
 
64
        #len(self.my_testmodel.get('interactions'))
 
65
 
 
66
 
 
67
        # Simplify the model
 
68
 
 
69
        particles = self.my_testmodel.get('particles')
 
70
        #print 'Here\n', self.my_testmodel['particles']
 
71
        interactions = self.my_testmodel.get('interactions')
 
72
        inter_list = copy.copy(interactions)
 
73
        
 
74
        no_want_pid = [1, 2, 3, 4, 13, 14, 15, 16, 21, 23]
 
75
        for pid in no_want_pid:
 
76
            particles.remove(self.my_testmodel.get_particle(pid))
 
77
 
 
78
        for inter in inter_list:
 
79
            if any([p.get('pdg_code') in no_want_pid for p in \
 
80
                        inter.get('particles')]):
 
81
                interactions.remove(inter)
 
82
 
 
83
        # Set a new name
 
84
 
 
85
        self.my_testmodel.set('name', 'my_smallsm')
 
86
        self.my_testmodel.set('particles', particles)
 
87
        self.my_testmodel.set('interactions', interactions)
 
88
 
 
89
 
 
90
        #Setup the vertexlist for my_testmodel and save this model
 
91
 
 
92
        import_vertexlist.make_vertexlist(self.my_testmodel)
 
93
        #save_model.save_model(os.path.join(MG5DIR, 'tests/input_files', 
 
94
        #self.my_testmodel['name']), self.my_testmodel)
 
95
 
 
96
 
 
97
        # Setup vertexlist for test
 
98
 
 
99
        full_vertexlist = import_vertexlist.full_vertexlist
 
100
 
 
101
        # Take w+ > t b~ and w+ > e+ ve (correct ones)
 
102
        #      w- > t~ b and t~ > w- b~ (wrong initial particle)
 
103
        self.my_2bodyvertexlist = base_objects.VertexList()
 
104
        self.my_2bodyvertexlist_wrongini = base_objects.VertexList()
 
105
        self.my_3bodyvertexlist = base_objects.VertexList()
 
106
        self.my_3bodyvertexlist_wrongini = base_objects.VertexList()
 
107
        self.my_3bodyvertexlist_radiactive = base_objects.VertexList()
 
108
 
 
109
        for index, vertex in full_vertexlist.items():
 
110
            legs_set = set([l['id'] for l in vertex['legs']])
 
111
            if legs_set == set([-5, 6, 24]) :
 
112
                self.my_2bodyvertexlist.append(vertex)
 
113
            elif legs_set == set([-11, 12, 24]):
 
114
                self.my_2bodyvertexlist.append(vertex)
 
115
            elif legs_set == set([-6, 5, -24]):
 
116
                self.my_2bodyvertexlist_wrongini.append(vertex)
 
117
            elif legs_set == set([-5, -24, -6]):
 
118
                self.my_2bodyvertexlist_wrongini.append(vertex)
 
119
 
 
120
        # Artificial 3-body vertices
 
121
        fake_vertex = copy.deepcopy(self.my_2bodyvertexlist[0])
 
122
        fake_vertex['legs'].append(base_objects.Leg({'id':22}))
 
123
        fake_vertex2 = copy.deepcopy(self.my_2bodyvertexlist[1])
 
124
        fake_vertex2['legs'].append(base_objects.Leg({'id': 11}))
 
125
        self.my_3bodyvertexlist.extend([fake_vertex, fake_vertex2])
 
126
 
 
127
        # Artificial 3-body vertices with wrong initial particle
 
128
        fake_vertex3 = copy.deepcopy(self.my_2bodyvertexlist_wrongini[0])
 
129
        fake_vertex3['legs'].append(base_objects.Leg({'id':12}))
 
130
        self.my_3bodyvertexlist_wrongini.append(fake_vertex3)
 
131
 
 
132
        # Artificial 3-body vertices with radiation
 
133
        fake_vertex4 = copy.deepcopy(self.my_2bodyvertexlist[0])
 
134
        fake_vertex4['legs'].append(base_objects.Leg({'id':24}))
 
135
        self.my_3bodyvertexlist_radiactive.append(fake_vertex4)
 
136
 
 
137
        
 
138
        # Testing particle
 
139
        self.mydict = {'name':'w+',
 
140
                      'antiname':'w-',
 
141
                      'spin':3,
 
142
                      'color':1,
 
143
                      'mass':'MW',
 
144
                      'width':'WW',
 
145
                      'texname':'W+',
 
146
                      'antitexname':'W-',
 
147
                      'line':'wavy',
 
148
                      'charge': 1.00,
 
149
                      'pdg_code': 24,
 
150
                      'propagating':True,
 
151
                      'is_part': True,
 
152
                      'self_antipart': False,
 
153
                       # decay_vertexlist must have two lists, one for on-shell,
 
154
                       # one for off-shell
 
155
                      'decay_vertexlist': {\
 
156
                           (2, False): self.my_2bodyvertexlist,
 
157
                           (2, True) : self.my_2bodyvertexlist,
 
158
                           (3, False): self.my_3bodyvertexlist,
 
159
                           (3, True) : self.my_3bodyvertexlist},
 
160
                       'is_stable': False,
 
161
                       'vertexlist_found': False,
 
162
                       'max_vertexorder': 0,
 
163
                       'apx_decaywidth': 0.,
 
164
                       'apx_decaywidth_err': 0.
 
165
                       }
 
166
 
 
167
        self.mypart = decay_objects.DecayParticle(self.mydict)
 
168
 
 
169
 
 
170
 
 
171
    def test_setgetinit_correct(self):
 
172
        """Test __init__, get, and set functions of DecayParticle
 
173
           mypart should give the dict as my dict
 
174
        """
 
175
        
 
176
        mypart2 = decay_objects.DecayParticle()
 
177
 
 
178
        # To avoid the error raised when setting the vertexlist
 
179
        # because of the wrong particle id.
 
180
        mypart2.set('pdg_code', self.mydict['pdg_code'])
 
181
        for key in self.mydict:
 
182
            # Test for the __init__ assign values as mydict
 
183
            self.assertEqual(self.mydict[key], self.mypart[key])
 
184
 
 
185
            # Test the set function
 
186
            mypart2.set(key, self.mydict[key])
 
187
            #print key, mypart2[key]
 
188
            self.assertEqual(mypart2[key], self.mydict[key])
 
189
 
 
190
        for key in self.mypart:
 
191
            # Test the get function return the value as in mypart
 
192
            # Note: for apx_decaywidth_err, .get will call 
 
193
            # estimate_decaywidth_error to recalculate
 
194
            # so the result will be one for zero width
 
195
            self.assertEqual(self.mypart.get(key), self.mypart[key])
 
196
 
 
197
 
 
198
    def test_setgetinit_exceptions(self):
 
199
        """Test the exceptions raised by __init__, get, and set."""
 
200
        
 
201
        myNondict = 1.
 
202
        myWrongdict = self.mydict
 
203
        myWrongdict['Wrongkey'] = 'wrongvalue'
 
204
 
 
205
        # Test __init__
 
206
        self.assertRaises(AssertionError, decay_objects.DecayParticle,myNondict)
 
207
        # Do not check if the input dict has wrong dict!
 
208
        #self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
209
        #                  decay_objects.DecayParticle, myWrongdict)
 
210
                          
 
211
        # Test get
 
212
        self.assertRaises(AssertionError, self.mypart.get, myNondict)
 
213
 
 
214
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
215
                          self.mypart.get, 'WrongParameter')
 
216
                          
 
217
        # Test set
 
218
        self.assertRaises(AssertionError, self.mypart.set, myNondict, 1)
 
219
 
 
220
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
221
                          self.mypart.set, 'WrongParameter', 1)
 
222
 
 
223
    def test_values_for_prop(self):
 
224
        """Test filters for DecayParticle properties."""
 
225
 
 
226
        test_values = [
 
227
                       {'prop':'name',
 
228
                        'right_list':['h', 'e+', 'e-', 'u~',
 
229
                                      'k++', 'k--', 'T', 'u+~'],
 
230
                        'wrong_list':['', 'x ', 'e?', '{}']},
 
231
                       {'prop':'spin',
 
232
                        'right_list':[1, 2, 3, 4, 5],
 
233
                        'wrong_list':[-1, 0, 'a', 6]},
 
234
                       {'prop':'color',
 
235
                        'right_list':[1, 3, 6, 8],
 
236
                        'wrong_list':[2, 0, 'a', 23, -1, -3, -6]},
 
237
                       {'prop':'mass',
 
238
                        'right_list':['me', 'zero', 'mm2'],
 
239
                        'wrong_list':['m+', '', ' ', 'm~']},
 
240
                       {'prop':'pdg_code',
 
241
                        #The last pdg_code must be 6 to be consistent with
 
242
                        #vertexlist
 
243
                        'right_list':[1, 12, 80000000, -1, 24],
 
244
                        'wrong_list':[1.2, 'a']},
 
245
                       {'prop':'line',
 
246
                        'right_list':['straight', 'wavy', 'curly', 'dashed'],
 
247
                        'wrong_list':[-1, 'wrong']},
 
248
                       {'prop':'charge',
 
249
                        'right_list':[1., -1., -2. / 3., 0.],
 
250
                        'wrong_list':[1, 'a']},
 
251
                       {'prop':'propagating',
 
252
                        'right_list':[True, False],
 
253
                        'wrong_list':[1, 'a', 'true', None]},
 
254
                       {'prop':'is_part',
 
255
                        #Restore the is_part to be consistent with vertexlist
 
256
                        'right_list':[True, False, True],
 
257
                        'wrong_list':[1, 'a', 'true', None]},
 
258
                       {'prop':'self_antipart',
 
259
                        'right_list':[True, False],
 
260
                        'wrong_list':[1, 'a', 'true', None]},
 
261
                       {'prop':'is_stable',
 
262
                        'right_list':[True, False],
 
263
                        'wrong_list':[1, 'a', 'true', None]},
 
264
                       {'prop':'vertexlist_found',
 
265
                        'right_list':[True, False],
 
266
                        'wrong_list':[1, 'a', 'true', None]},
 
267
                       {'prop':'max_vertexorder',
 
268
                        'right_list':[3, 4, 0],
 
269
                        'wrong_list':['a', 'true', None]},
 
270
                       {'prop':'apx_decaywidth',
 
271
                        'right_list':[3., 4.5, 0.2],
 
272
                        'wrong_list':['a', [12,2], None]},
 
273
                       {'prop':'apx_decaywidth_err',
 
274
                        'right_list':[3., 4.5, 0.2],
 
275
                        'wrong_list':['a', [12,2], None]},
 
276
                       {'prop':'2body_massdiff',
 
277
                        'right_list':[3., 4.5, 0.2],
 
278
                        'wrong_list':['a', [12,2], None]},
 
279
                       {'prop':'decay_vertexlist',
 
280
                        'right_list':[{(2, False):self.my_2bodyvertexlist,
 
281
                                       (2, True) :self.my_2bodyvertexlist,
 
282
                                       (3, False):self.my_3bodyvertexlist,
 
283
                                       (3, True) :self.my_3bodyvertexlist}],
 
284
                        'wrong_list':[1, 
 
285
                                      {'a': self.my_2bodyvertexlist},
 
286
                                      {(24, 2, False): self.my_2bodyvertexlist},
 
287
                                      {(5, True):self.my_2bodyvertexlist,
 
288
                                       (5, False):self.my_3bodyvertexlist},
 
289
                                      {(2, 'Not bool'):self.my_2bodyvertexlist},
 
290
                                      {(2, False): 'hey'},
 
291
                                      {(2, False): self.my_2bodyvertexlist, 
 
292
                                       (2, True) : self.my_3bodyvertexlist},
 
293
                                      {(2, False):self.my_2bodyvertexlist_wrongini, 
 
294
                                       (2, True): self.my_2bodyvertexlist,
 
295
                                       (3, False):self.my_3bodyvertexlist,
 
296
                                       (3, True): self.my_3bodyvertexlist},
 
297
                                      {(2, False):self.my_2bodyvertexlist, 
 
298
                                       (2, True): self.my_2bodyvertexlist,
 
299
                                       (3, False):self.my_3bodyvertexlist_wrongini,
 
300
                                       (3, True): self.my_3bodyvertexlist},
 
301
                                      {(2, False):self.my_2bodyvertexlist, 
 
302
                                       (2, True): self.my_2bodyvertexlist,
 
303
                                       (3, False):self.my_3bodyvertexlist,
 
304
                                       (3, True): self.my_3bodyvertexlist_radiactive}
 
305
                                      
 
306
                                     ]},
 
307
                       ]
 
308
 
 
309
        temp_part = self.mypart
 
310
 
 
311
        for test in test_values:
 
312
            for x in test['right_list']:
 
313
                self.assertTrue(temp_part.set(test['prop'], x))
 
314
            for x in test['wrong_list']:
 
315
                self.assertFalse(temp_part.set(test['prop'], x))
 
316
 
 
317
    def test_getsetvertexlist_correct(self):
 
318
        """Test the get and set for vertexlist is correct"""
 
319
        temp_part = self.mypart
 
320
 
 
321
        # Reset the off-shell '2_body_decay_vertexlist'
 
322
        templist = self.my_2bodyvertexlist
 
323
        templist.extend(templist)
 
324
        temp_part.set_vertexlist(2, False, templist)
 
325
 
 
326
        # Test for equality from get_vertexlist
 
327
        self.assertEqual(temp_part.get_vertexlist(2, False),
 
328
                         templist)
 
329
 
 
330
        # Test the result for keys don't exist
 
331
        self.assertEqual(temp_part.get_vertexlist(4, True),
 
332
                         [])
 
333
 
 
334
        # Reset the on-shell '2_body_decay_vertexlist'
 
335
        templist.extend(templist)
 
336
        temp_part.set_vertexlist(2, True, templist)
 
337
 
 
338
        # Test for equality from get_vertexlist
 
339
        self.assertEqual(temp_part.get_vertexlist(2, True),
 
340
                         templist)
 
341
 
 
342
 
 
343
        # Reset the off-shell '3_body_decay_vertexlist'
 
344
        templist = self.my_3bodyvertexlist
 
345
        templist.extend(templist)
 
346
        temp_part.set_vertexlist(3, False, templist)
 
347
 
 
348
        # Test for equality from get_vertexlist
 
349
        self.assertEqual(temp_part.get_vertexlist(3, False),
 
350
                         templist)
 
351
 
 
352
 
 
353
        # Reset the on-shell '3_body_decay_vertexlist'
 
354
        templist.extend(templist)
 
355
        temp_part.set_vertexlist(3, True, templist)
 
356
 
 
357
        # Test for equality from get_vertexlist
 
358
        self.assertEqual(temp_part.get_vertexlist(3, True),
 
359
                         templist)
 
360
 
 
361
    def test_getsetvertexlist_exceptions(self):
 
362
        """Test for the exceptions raised by the get_ or set_vertexlist"""
 
363
 
 
364
        # Test of get_ and set_vertexlist
 
365
        # Test the exceptions raised from partnum and onshell
 
366
        for wrongpartnum in ['string', 1.5]:
 
367
            self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
368
                              self.mypart.get_vertexlist, wrongpartnum, True)
 
369
            self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
370
                              self.mypart.set_vertexlist, wrongpartnum, True,
 
371
                              self.my_2bodyvertexlist, self.my_testmodel)
 
372
 
 
373
        for wrongbool in [15, 'NotBool']:           
 
374
            self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
375
                              self.mypart.get_vertexlist, 2, wrongbool)
 
376
            self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
377
                              self.mypart.set_vertexlist, 3, wrongbool,
 
378
                              self.my_3bodyvertexlist, self.my_testmodel)
 
379
 
 
380
        
 
381
 
 
382
        # Test the exceptions raised from value in set_vertexlist
 
383
        # Test for non vertexlist objects
 
384
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
385
                          self.mypart.set_vertexlist,
 
386
                          2, True, ['not', 'Vertexlist'], self.my_testmodel)
 
387
 
 
388
        #Test for vertexlist not consistent with partnum
 
389
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
390
                          self.mypart.set_vertexlist,
 
391
                          2, True, self.my_3bodyvertexlist, self.my_testmodel)
 
392
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
393
                          self.mypart.set_vertexlist,
 
394
                          3, True, self.my_2bodyvertexlist, self.my_testmodel)
 
395
 
 
396
        # Test for vertexlist not consistent with initial particle
 
397
        # for both number and id
 
398
        # Use the vertexlist from test_getsetvertexlist_exceptions
 
399
 
 
400
        Wrong_vertexlist = [self.my_2bodyvertexlist_wrongini,
 
401
                            self.my_3bodyvertexlist_wrongini,
 
402
                            self.my_3bodyvertexlist_radiactive]
 
403
 
 
404
        for item in Wrong_vertexlist:
 
405
            for partnum in [2,3]:
 
406
                self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError
 
407
                             , self.mypart.set_vertexlist, partnum, False, item)
 
408
                
 
409
 
 
410
    def test_reset_decaywidth(self):
 
411
        """ Test for the reset and update of apx_decaywidth, branching ratio
 
412
            for amplitudes, and apx_decaywidth_err."""
 
413
 
 
414
        self.my_testmodel.find_all_channels(3)
 
415
        tquark = self.my_testmodel.get_particle(6)
 
416
        amp = tquark.get_amplitudes(2)[0]
 
417
        old_width = amp.get('apx_decaywidth')
 
418
 
 
419
        # Add a new channel to tquark
 
420
        tquark.get_amplitudes(2).append(amp)
 
421
 
 
422
        # Total width doubles. Update the width, br should remain the same.
 
423
        amp.get('apx_br')
 
424
        tquark.update_decay_attributes(True, True, False)
 
425
        self.assertAlmostEqual(tquark.get('apx_decaywidth'),old_width*2)
 
426
        self.assertAlmostEqual(amp.get('apx_br'), 1)
 
427
 
 
428
        # Update the br now. It should be 0.5.
 
429
        tquark.update_decay_attributes(False, True, True)
 
430
        self.assertAlmostEqual(tquark.get('apx_decaywidth'),old_width*2)
 
431
        self.assertAlmostEqual(amp.get('apx_br'), 0.5)
 
432
 
 
433
        # Test the estimate_width_error
 
434
        w = self.my_testmodel.get_particle(24)
 
435
        width_err = w.get('apx_decaywidth_err')
 
436
        offshell_clist = w.get_channels(3, False)
 
437
        w.get_channels(3, False).extend(offshell_clist)
 
438
 
 
439
        # Test if the width_err doubles after the update
 
440
        w.update_decay_attributes(False, True, False)
 
441
        self.assertAlmostEqual(w.get('apx_decaywidth_err'),width_err*2)
 
442
 
 
443
    def test_find_vertexlist(self):
 
444
        """ Test for the find_vertexlist function and 
 
445
            the get_max_vertexorder"""
 
446
        
 
447
        # Test validity of arguments
 
448
        # Test if the calling particle is in the model
 
449
        extra_part = copy.copy(self.mypart)
 
450
        extra_part.set('pdg_code', 2)
 
451
        extra_part.set('name', 'u')
 
452
 
 
453
 
 
454
        # Test the return of get_max_vertexorder if  vertexlist_found = False
 
455
        self.assertEqual(None, extra_part.get_max_vertexorder())
 
456
 
 
457
        #print self.my_testmodel
 
458
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
459
                          extra_part.find_vertexlist, self.my_testmodel)
 
460
 
 
461
 
 
462
        # Test if option is boolean
 
463
        wronglist=[ 'a', 5, {'key': 9}, [1,5]]
 
464
        for wrongarg in wronglist:
 
465
            self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
466
                              self.mypart.find_vertexlist, self.my_testmodel,\
 
467
                              wrongarg)
 
468
 
 
469
 
 
470
        # Test find_vertexlist in t quark, w+ boson, photon
 
471
 
 
472
        tquark = decay_objects.DecayParticle(self.my_testmodel.get_particle(6),
 
473
                                             True)
 
474
        tquark.find_vertexlist(self.my_testmodel)
 
475
 
 
476
        wboson_p = decay_objects.DecayParticle(\
 
477
            self.my_testmodel.get_particle(24), True)
 
478
        wboson_p.find_vertexlist(self.my_testmodel)
 
479
 
 
480
        # photon is known as stable after stable particles are found internally
 
481
        # in find_vertexlist
 
482
        photon = decay_objects.DecayParticle(self.my_testmodel.get_particle(22),
 
483
                                             True)
 
484
        photon.find_vertexlist(self.my_testmodel)
 
485
 
 
486
        # Get vertex list
 
487
        blank_vlist = base_objects.VertexList()
 
488
        # t > b w+ (t~ b w+)
 
489
        top_vlist_2_on = base_objects.VertexList()        
 
490
        # w+ > t b~
 
491
        w_vlist_2_on = base_objects.VertexList()
 
492
        # w+ > e+ ve        
 
493
        w_vlist_2_off = base_objects.VertexList()
 
494
 
 
495
        for index, vertex in import_vertexlist.full_vertexlist.items():
 
496
            legs_set = set([l['id'] for l in vertex['legs']])
 
497
            if legs_set == set([5, 24, 6]) :
 
498
                top_vlist_2_on.append(vertex)
 
499
            elif legs_set == set([-11, 12, 24]):
 
500
                w_vlist_2_on.append(vertex)
 
501
            elif legs_set == set([-5, 6, 24]):
 
502
                w_vlist_2_off.append(vertex)
 
503
        
 
504
        rightlist_top = [blank_vlist, top_vlist_2_on, blank_vlist, blank_vlist]
 
505
        rightlist_w =[w_vlist_2_off, w_vlist_2_on, blank_vlist, blank_vlist] 
 
506
        rightlist_a =[blank_vlist, blank_vlist, blank_vlist, blank_vlist]
 
507
 
 
508
        i=0
 
509
        for partnum in [2,3]:
 
510
            for onshell in [False, True]:
 
511
                self.assertEqual(tquark.get_vertexlist(partnum, onshell),
 
512
                                 rightlist_top[i])
 
513
                self.assertEqual(wboson_p.get_vertexlist(partnum, onshell),
 
514
                                 rightlist_w[i])
 
515
                self.assertEqual(photon.get_vertexlist(partnum, onshell),
 
516
                                 rightlist_a[i])
 
517
                i +=1
 
518
 
 
519
 
 
520
        # Test get_max_vertexorder()
 
521
        self.assertEqual(0, photon.get_max_vertexorder())
 
522
        self.assertEqual(2, tquark.get_max_vertexorder())
 
523
        self.mypart['vertexlist_found'] = True
 
524
        self.assertEqual(3, self.mypart.get_max_vertexorder())
 
525
 
 
526
 
 
527
    def test_setget_channel(self):
 
528
        """ Test of the get_channel set_channel functions (and the underlying
 
529
            check_channels.)"""
 
530
 
 
531
        # Prepare the channel
 
532
        full_vertexlist = import_vertexlist.full_vertexlist
 
533
 
 
534
        vert_0 = base_objects.Vertex({'id': 0, 'legs': base_objects.LegList([\
 
535
                 base_objects.Leg({'id':25, 'number':1, 'state': False}),
 
536
                 base_objects.Leg({'id':25, 'number':2})])})
 
537
 
 
538
        for index, vertex in import_vertexlist.full_vertexlist.items():
 
539
            legs_set = set([l['id'] for l in vertex['legs']])
 
540
            # h > t t~
 
541
            if legs_set == set([-6, 6, 25]) :
 
542
                vert_1 = copy.deepcopy(vertex)
 
543
            # t > b w+ = (t~ b w+)
 
544
            elif legs_set == set([-5, -24, -6]):
 
545
                vert_2 = copy.deepcopy(vertex)
 
546
            # t~ > b~ w- = (t b~ w-)
 
547
            elif legs_set == set([5, 24, 6]):
 
548
                vert_3 = copy.deepcopy(vertex)
 
549
 
 
550
        vert_1['legs'][0]['number'] = 2
 
551
        vert_1['legs'][1]['number'] = 3
 
552
        vert_1['legs'][2]['number'] = 2
 
553
 
 
554
        vert_2['legs'][0]['number'] = 2
 
555
        vert_2['legs'][1]['number'] = 4
 
556
        vert_2['legs'][2]['number'] = 2
 
557
 
 
558
        vert_3['legs'][0]['number'] = 3
 
559
        vert_3['legs'][1]['number'] = 5
 
560
        vert_3['legs'][2]['number'] = 3
 
561
 
 
562
        h_tt_bbww = decay_objects.Channel({'vertices': \
 
563
                                           base_objects.VertexList([
 
564
                                           vert_3, vert_2, 
 
565
                                           vert_1, vert_0])})
 
566
        channellist = decay_objects.ChannelList([h_tt_bbww])
 
567
 
 
568
 
 
569
        # Test set and get
 
570
        higgs = self.my_testmodel.get_particle(25)
 
571
        higgs.set('decay_channels', {(4, True): channellist})
 
572
        self.assertEqual(higgs.get('decay_channels'), {(4, True): channellist})
 
573
                
 
574
        # Test set_channel and get_channel
 
575
        higgs = self.my_testmodel.get_particle(25)
 
576
        higgs.set_channels(4, True, [h_tt_bbww])
 
577
        self.assertEqual(higgs.get_channels(4, True), channellist)
 
578
 
 
579
 
 
580
        # Test for exceptions
 
581
 
 
582
        # Wrong final particle number
 
583
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError, 
 
584
                          higgs.set_channels, 'non_int', True, [h_tt_bbww])
 
585
        # Test from the filter function
 
586
        self.assertFalse(higgs.set('decay_channels', 
 
587
                                   {('non_int', True): channellist}))
 
588
        # Wrong onshell
 
589
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError, 
 
590
                          higgs.get_channels, 3, 5)
 
591
        # Wrong channel
 
592
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError, 
 
593
                          higgs.set_channels, 3, True, ['non', 'channellist'])
 
594
        # Wrong initial particle
 
595
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError, 
 
596
                          self.my_testmodel.get_particle(24).set_channels, 3,
 
597
                          True, [h_tt_bbww])
 
598
        # Wrong onshell condition (h is lighter than ww pair)
 
599
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError, 
 
600
                          higgs.set_channels, 3, True, [h_tt_bbww],
 
601
                          self.my_testmodel)
 
602
        non_sm = copy.copy(higgs)
 
603
        non_sm.set('pdg_code', 26)
 
604
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError, 
 
605
                          higgs.set_channels, 3, False, [h_tt_bbww],
 
606
                          self.my_testmodel)
 
607
 
 
608
 
 
609
    def test_get_max_level(self):
 
610
        """ Test the get_max_level function. """
 
611
 
 
612
        higgs = self.my_testmodel.get_particle(25)
 
613
        higgs.find_channels(2, self.my_testmodel)
 
614
        self.assertEqual(higgs.get_max_level(), 2)
 
615
        higgs.find_channels_nextlevel(self.my_testmodel)
 
616
        self.assertEqual(higgs.get_max_level(), 3)
 
617
        higgs.find_channels_nextlevel(self.my_testmodel)
 
618
        higgs.set_amplitudes(4, decay_objects.DecayAmplitudeList())
 
619
        self.assertEqual(higgs.get_max_level(), 4)
 
620
 
 
621
#===============================================================================
 
622
# TestDecayParticleList
 
623
#===============================================================================
 
624
class Test_DecayParticleList(unittest.TestCase):
 
625
    """Test the DecayParticleList"""
 
626
    def setUp(self):
 
627
        self.mg5_part = base_objects.Particle({'pdg_code':6, 'is_part':True})
 
628
        self.mg5_partlist = base_objects.ParticleList([self.mg5_part]*5)
 
629
 
 
630
    def test_convert(self):
 
631
        #Test the conversion in __init__
 
632
        decay_partlist = decay_objects.DecayParticleList(self.mg5_partlist)
 
633
        for i in range(0, 5):
 
634
            self.assertTrue(isinstance(decay_partlist[i], 
 
635
                                       decay_objects.DecayParticle))
 
636
 
 
637
        #Test the conversion in append
 
638
        decay_partlist.append(self.mg5_part)
 
639
        self.assertTrue(isinstance(decay_partlist[-1], 
 
640
                                   decay_objects.DecayParticle))
 
641
        self.assertTrue(isinstance(decay_partlist,
 
642
                                   decay_objects.DecayParticleList))
 
643
        
 
644
        #Test the conversion in generate_dict
 
645
        for num, part in decay_partlist.generate_dict().items():
 
646
            self.assertTrue(isinstance(part, decay_objects.DecayParticle))
 
647
 
 
648
 
 
649
#===============================================================================
 
650
# TestDecayModel
 
651
#===============================================================================
 
652
class Test_DecayModel(unittest.TestCase):
 
653
    """Test class for the DecayModel object"""
 
654
 
 
655
    def setUp(self):
 
656
        """Set up decay model"""
 
657
        
 
658
        if not hasattr(Test_DecayModel, 'base_model'):
 
659
            Test_DecayModel.base_model = import_ufo.import_model('mssm')
 
660
            Test_DecayModel.my_testmodel_base = import_ufo.import_model('sm')
 
661
        
 
662
        #Full SM DecayModel
 
663
        self.decay_model = decay_objects.DecayModel(self.base_model, True)
 
664
 
 
665
        #My_small sm DecayModel
 
666
        self.my_testmodel = decay_objects.DecayModel(self.my_testmodel_base, True)
 
667
        param_path = os.path.join(_file_path,'../input_files/param_card_sm.dat')
 
668
        self.my_testmodel.read_param_card(param_path)
 
669
 
 
670
        # Simplify the model
 
671
        particles = self.my_testmodel.get('particles')
 
672
        interactions = self.my_testmodel.get('interactions')
 
673
        inter_list = copy.copy(interactions)
 
674
        no_want_pid = [1, 2, 3, 4, 13, 14, 15, 16, 21, 23, 25]
 
675
        for pid in no_want_pid:
 
676
            particles.remove(self.my_testmodel.get_particle(pid))
 
677
 
 
678
        for particle in particles:
 
679
            particle.set('charge', 0)
 
680
 
 
681
        for inter in inter_list:
 
682
            if any([p.get('pdg_code') in no_want_pid for p in \
 
683
                        inter.get('particles')]):
 
684
                interactions.remove(inter)
 
685
 
 
686
        # Set a new name
 
687
        self.my_testmodel.set('name', 'my_smallsm')
 
688
        self.my_testmodel.set('particles', particles)
 
689
        self.my_testmodel.set('interactions', interactions)
 
690
 
 
691
        import_vertexlist.make_vertexlist(self.my_testmodel)
 
692
 
 
693
        #import madgraph.iolibs.export_v4 as export_v4
 
694
        #writer = export_v4.UFO_model_to_mg4(self.base_model,'temp')
 
695
        #writer.build()
 
696
 
 
697
    def test_read_param_card(self):
 
698
        """Test reading a param card"""
 
699
        param_path = os.path.join(_file_path, '../input_files/param_card_mssm.dat')
 
700
        self.decay_model.read_param_card(os.path.join(param_path))
 
701
 
 
702
        for param in sum([self.base_model.get('parameters')[key] for key \
 
703
                              in self.base_model.get('parameters')], []):
 
704
            if param.name:
 
705
                value = eval("decay_objects.%s" % param.name)
 
706
                self.assertTrue(isinstance(value, int) or \
 
707
                                    isinstance(value, float) or \
 
708
                                    isinstance(value, complex))
 
709
 
 
710
    def test_setget(self):
 
711
        """ Test the set and get for special properties"""
 
712
 
 
713
        self.my_testmodel.set('vertexlist_found', True)
 
714
        self.assertEqual(self.my_testmodel.get('vertexlist_found'), True)
 
715
 
 
716
        self.my_testmodel.set('vertexlist_found', False)
 
717
        self.assertRaises(self.my_testmodel.PhysicsObjectError,
 
718
                          self.my_testmodel.filter, 'max_vertexorder', 'a')
 
719
        self.assertRaises(self.my_testmodel.PhysicsObjectError,
 
720
                          self.my_testmodel.filter, 'stable_particles', 
 
721
                          [self.my_testmodel.get('particles'), ['a']])
 
722
        self.assertRaises(decay_objects.DecayModel.PhysicsObjectError,
 
723
                          self.my_testmodel.filter, 'vertexlist_found', 4)
 
724
                          
 
725
 
 
726
    def test_particles_type(self):
 
727
        """Test if the DecayModel can convert the assign particle into
 
728
           decay particle"""
 
729
 
 
730
        # Test the particle is DecayParticle during generator stage
 
731
        # Test the default_setup first
 
732
        temp_model = decay_objects.DecayModel()
 
733
        self.assertTrue(isinstance(temp_model.get('particles'),
 
734
                              decay_objects.DecayParticleList))
 
735
 
 
736
 
 
737
        # Test the embeded set in __init__
 
738
        self.assertTrue(isinstance(self.decay_model.get('particles'), 
 
739
                                   decay_objects.DecayParticleList))
 
740
 
 
741
 
 
742
        # Test the conversion into DecayParticle explicitly
 
743
        # by the set function
 
744
        mg5_particlelist = self.base_model['particles']
 
745
 
 
746
        result = self.decay_model.set('particles', mg5_particlelist, True)
 
747
 
 
748
 
 
749
        # Using ParticleList to set should be fine, the result is converted
 
750
        # into DecayParticleList.
 
751
        self.assertTrue(result)
 
752
        self.assertTrue(isinstance(self.decay_model['particles'],
 
753
                              decay_objects.DecayParticleList))
 
754
 
 
755
 
 
756
        # particle_dict should contain DecayParticle
 
757
        self.assertTrue(isinstance(self.decay_model.get('particle_dict')[6],
 
758
                                   decay_objects.DecayParticle))
 
759
 
 
760
 
 
761
        # Test if the set function returns correctly when assign a bad value
 
762
        try:
 
763
            self.assertFalse(self.decay_model.set('particles', 
 
764
                                                  'NotParticleList'))
 
765
        except:
 
766
            self.assertRaises(AssertionError, 
 
767
                              self.decay_model.set, 
 
768
                              'particles', 'NotParticleList')
 
769
 
 
770
 
 
771
        # Test if the particls in interaction is converted to DecayParticle
 
772
        self.assertTrue(isinstance(self.decay_model['interactions'][-1]['particles'], decay_objects.DecayParticleList))
 
773
                        
 
774
 
 
775
    def test_find_vertexlist(self):
 
776
        """Test of the find_vertexlist"""
 
777
 
 
778
        # Set the mass of bottom to be nonzero for the test of radiation
 
779
        self.my_testmodel.get_particle(5)['mass'] = '4.7'
 
780
        # Test the exception of get_max_vertexorder
 
781
        self.assertEqual(None, self.my_testmodel.get_max_vertexorder())
 
782
 
 
783
 
 
784
        # Test find vertex exception when vertices are repeated
 
785
        # (vertices already be stored in DecayParicles in the 1st search)
 
786
        self.my_testmodel.find_vertexlist()
 
787
        self.my_testmodel['vertexlist_found'] = False
 
788
        self.assertRaises(decay_objects.DecayModel.PhysicsObjectError,
 
789
                          self.my_testmodel.find_vertexlist)
 
790
 
 
791
        # Ignore it by specifying it in argument.
 
792
        self.my_testmodel['vertexlist_found'] = False
 
793
        self.my_testmodel.find_vertexlist(True)
 
794
 
 
795
        # Reset everything
 
796
        self.my_testmodel['vertexlist_found'] = False
 
797
        for p in self.my_testmodel['particles']:
 
798
            p['decay_vertexlist'] = {}
 
799
 
 
800
 
 
801
 
 
802
        # Test find vertexlist
 
803
        self.my_testmodel.find_vertexlist()
 
804
        empty = []
 
805
 
 
806
        # General test by comparing to import_vertexlist
 
807
        full_vertexlist = import_vertexlist.full_vertexlist_newindex
 
808
        for part in self.my_testmodel.get('particles'):
 
809
            for partnum in [2, 3]:
 
810
                for onshell in [True, False]:
 
811
                    #print part.get_pdg_code(), partnum, onshell
 
812
                    # Use get(key, default) to prevent keys do not exist
 
813
                    result = part.get_vertexlist(partnum, onshell)
 
814
                    goal = full_vertexlist.get((part.get_pdg_code(),
 
815
                                                partnum, onshell), empty)
 
816
 
 
817
                    self.assertEqual(result, goal)
 
818
 
 
819
                    # Test the order of legs
 
820
                    for vert in part.get_vertexlist(partnum, onshell):
 
821
                        lids = [l['id'] for l in vert['legs'][:-1]]
 
822
                        goal_lids = sorted(lids, reverse= True)
 
823
                        self.assertEqual(goal_lids, lids)
 
824
                    
 
825
        # Specific test to top, w+, photon (very similar to another 
 
826
        # test_find_vertexlist in DecayParticle)
 
827
        tquark = decay_objects.DecayParticle(self.my_testmodel.get_particle(6),
 
828
                                             True)
 
829
        wboson_p = decay_objects.DecayParticle(\
 
830
            self.my_testmodel.get_particle(24), True)
 
831
        photon = decay_objects.DecayParticle(self.my_testmodel.get_particle(22),
 
832
                                             True)
 
833
 
 
834
        # Get vertex list
 
835
        # t > b w+ (t~ b w+)
 
836
        top_vlist_2_on = base_objects.VertexList()        
 
837
        # w+ > t b~
 
838
        w_vlist_2_on = base_objects.VertexList()
 
839
        # w+ > e+ ve        
 
840
        w_vlist_2_off = base_objects.VertexList()
 
841
 
 
842
        for index, vertex in import_vertexlist.full_vertexlist.items():
 
843
            legs_set = set([l['id'] for l in vertex['legs']])
 
844
            if legs_set == set([5, 24, 6]) :
 
845
                top_vlist_2_on.append(vertex)
 
846
            elif legs_set == set([-11, 12, 24]):
 
847
                w_vlist_2_on.append(vertex)
 
848
            elif legs_set == set([-5, 6, 24]):
 
849
                w_vlist_2_off.append(vertex)
 
850
        
 
851
        rightlist_top = [empty, top_vlist_2_on, empty, empty]
 
852
        rightlist_w =[w_vlist_2_off, w_vlist_2_on, empty, empty] 
 
853
        rightlist_a =[empty, empty, empty, empty]
 
854
 
 
855
        i=0
 
856
        for partnum in [2,3]:
 
857
            for onshell in [False, True]:
 
858
                self.assertEqual(tquark.get_vertexlist(partnum, onshell),
 
859
                                 rightlist_top[i])
 
860
                self.assertEqual(wboson_p.get_vertexlist(partnum, onshell),
 
861
                                 rightlist_w[i])
 
862
                self.assertEqual(photon.get_vertexlist(partnum, onshell),
 
863
                                 rightlist_a[i])
 
864
                i +=1
 
865
 
 
866
        return
 
867
 
 
868
 
 
869
 
 
870
 
 
871
    def test_find_mssm_decay_groups(self):
 
872
        """Test finding the decay groups of the MSSM"""
 
873
 
 
874
        mssm = import_ufo.import_model('mssm')
 
875
        decay_mssm = decay_objects.DecayModel(mssm, True)
 
876
        decay_mssm.find_decay_groups()
 
877
        goal_groups = [[25, 35, 36, 37],
 
878
                       [1000001, 1000002, 1000003, 1000004, 1000005, 1000006, 1000011, 1000012, 1000013, 1000014, 1000015, 1000016, 1000021, 1000022, 1000023, 1000024, 1000025, 1000035, 1000037, 2000001, 2000002, 2000003, 2000004, 2000005, 2000006, 2000011, 2000013, 2000015]]
 
879
 
 
880
        # find_decay_groups_general should be run automatically
 
881
        for i, group in enumerate(decay_mssm['decay_groups']):
 
882
            self.assertEqual(sorted([p.get('pdg_code') for p in group]),
 
883
                             goal_groups[i])
 
884
 
 
885
    def test_find_mssm_decay_groups_modified_mssm(self):
 
886
        """Test finding the decay groups of the MSSM"""
 
887
 
 
888
        mssm = import_ufo.import_model('mssm')
 
889
        particles = mssm.get('particles')
 
890
        no_want_particle_codes = [1000022, 1000023, 1000024, -1000024, 
 
891
                                  1000025, 1000035, 1000037, -1000037]
 
892
        no_want_particles = [p for p in particles if p.get('pdg_code') in \
 
893
                                 no_want_particle_codes]
 
894
 
 
895
        for particle in no_want_particles:
 
896
            particles.remove(particle)
 
897
 
 
898
        interactions = mssm.get('interactions')
 
899
        inter_list = copy.copy(interactions)
 
900
 
 
901
        for interaction in inter_list:
 
902
            if any([p.get('pdg_code') in no_want_particle_codes for p in \
 
903
                        interaction.get('particles')]):
 
904
                interactions.remove(interaction)
 
905
        
 
906
        mssm.set('particles', particles)
 
907
        mssm.set('interactions', interactions)
 
908
        decay_mssm = decay_objects.DecayModel(mssm, True)
 
909
 
 
910
        decay_mssm.find_decay_groups()
 
911
        goal_groups = set([(25, 35, 36, 37),
 
912
                           (1000001, 1000002, 1000003, 1000004, 1000005, 
 
913
                            1000006, 1000021, 2000001, 2000002, 2000003, 
 
914
                            2000004, 2000005, 2000006), 
 
915
                           (1000011, 1000012), 
 
916
                           (1000013, 1000014), 
 
917
                           (1000015, 1000016, 2000015), 
 
918
                           (2000011,), 
 
919
                           (2000013,)])
 
920
 
 
921
        self.assertEqual(set([tuple(sorted([p.get('pdg_code') for p in \
 
922
                                                group])) \
 
923
                                  for group in decay_mssm['decay_groups']]),
 
924
                         goal_groups)
 
925
 
 
926
 
 
927
    def test_find_mssm_decay_groups_general(self):
 
928
        """Test finding the decay groups of the MSSM"""
 
929
 
 
930
        mssm = import_ufo.import_model('mssm')
 
931
        decay_mssm = decay_objects.DecayModel(mssm, True)
 
932
        # Read data to find massless SM-like particle
 
933
        param_path = os.path.join(_file_path,
 
934
                                  '../input_files/param_card_mssm.dat')
 
935
        decay_mssm.read_param_card(param_path)
 
936
 
 
937
        # In first group, 15 and from 23 are calculated,
 
938
        # others are massless default.
 
939
        goal_groups = [[1,2,3,4,11,12,13,14, 15, 16,21,22, 23, 24, 25, 35, 36, 37], 
 
940
                       [1000001, 1000002, 1000003, 1000004, 1000011, 1000012, 1000013, 1000014, 1000015, 1000016, 1000021, 1000022, 1000023, 1000024, 1000025, 1000035, 1000037, 2000001, 2000002, 2000003, 2000004, 2000011, 2000013, 2000015],
 
941
                       [1000005, 1000006, 2000005, 2000006], 
 
942
                       [5, 6]]
 
943
        goal_stable_particle_ids = set([(1,2,3,4,11,12,13,14,16,21,22),
 
944
                                        (5,),
 
945
                                        (1000022,)])
 
946
        for i, group in enumerate(decay_mssm.get('decay_groups')):
 
947
            pdg = sorted([p.get('pdg_code') for p in group])
 
948
            self.assertTrue( pdg in goal_groups)
 
949
        self.assertEqual(len(goal_groups), i+1)
 
950
            
 
951
        # Test if all useless interactions are deleted.
 
952
        for inter in decay_mssm['reduced_interactions']:
 
953
            self.assertTrue(len(inter['particles']))
 
954
 
 
955
        # Reset decay_groups, test the auto run from find_stable_particles
 
956
        decay_mssm['decay_groups'] = []
 
957
        
 
958
        self.assertEqual(set([tuple(sorted([p.get('pdg_code') for p in plist])) for plist in decay_mssm.get('stable_particles')]), goal_stable_particle_ids)
 
959
 
 
960
            
 
961
 
 
962
    def test_find_mssm_decay_groups_modified_mssm_general(self):
 
963
        """Test finding the decay groups of the MSSM using general way.
 
964
           Test to get decay_groups and stable_particles from get."""
 
965
 
 
966
        # Setup the mssm with parameters read in.
 
967
        mssm = import_ufo.import_model('mssm')
 
968
        decay_mssm = decay_objects.DecayModel(mssm, True)
 
969
        particles = decay_mssm.get('particles')
 
970
        param_path = os.path.join(_file_path,
 
971
                                  '../input_files/param_card_mssm.dat')
 
972
        decay_mssm.read_param_card(param_path)
 
973
        
 
974
        # Set no want particles
 
975
        no_want_particle_codes = [1000022, 1000023, 1000024, -1000024, 
 
976
                                  1000025, 1000035, 1000037, -1000037]
 
977
        no_want_particles = [p for p in particles if p.get('pdg_code') in \
 
978
                                 no_want_particle_codes]
 
979
 
 
980
        for particle in no_want_particles:
 
981
            particles.remove(particle)
 
982
 
 
983
        interactions = decay_mssm.get('interactions')
 
984
        inter_list = copy.copy(interactions)
 
985
 
 
986
        for interaction in inter_list:
 
987
            if any([p.get('pdg_code') in no_want_particle_codes for p in \
 
988
                        interaction.get('particles')]):
 
989
                interactions.remove(interaction)
 
990
        
 
991
        decay_mssm.set('particles', particles)
 
992
        decay_mssm.set('interactions', interactions)
 
993
 
 
994
        # Set sd4, sd5 quark mass the same as b quark, so that 
 
995
        # degeneracy happens and can be tested 
 
996
        # (both particle and anti-particle must be changed)
 
997
        # This reset of particle mass must before the reset of particles
 
998
        # so that the particles of all interactions can change simutaneuosly.
 
999
        decay_mssm.get_particle(2000003)['mass'] = \
 
1000
            decay_mssm.get_particle(5).get('mass')
 
1001
        decay_mssm.get_particle(2000001)['mass'] = \
 
1002
            decay_mssm.get_particle(5).get('mass')
 
1003
        decay_mssm.get_particle(1000012)['mass'] = \
 
1004
            decay_mssm.get_particle(1000015).get('mass')
 
1005
 
 
1006
        decay_mssm.get_particle(-2000003)['mass'] = \
 
1007
            decay_mssm.get_particle(5).get('mass')
 
1008
        decay_mssm.get_particle(-2000001)['mass'] = \
 
1009
            decay_mssm.get_particle(5).get('mass')
 
1010
 
 
1011
 
 
1012
        # New interactions that mix different groups
 
1013
        new_interaction = base_objects.Interaction({\
 
1014
                'id': len(decay_mssm.get('interactions'))+1,
 
1015
                'particles': base_objects.ParticleList(
 
1016
                             [decay_mssm.get_particle(1000001),
 
1017
                              decay_mssm.get_particle(1000011),
 
1018
                              decay_mssm.get_particle(1000003),
 
1019
                              decay_mssm.get_particle(1000013),
 
1020
                              # This new SM particle should be removed
 
1021
                              # In the reduction level
 
1022
                              decay_mssm.get_particle(2000013),
 
1023
                              decay_mssm.get_particle(1000015)])})
 
1024
        new_interaction_add_sm = base_objects.Interaction({\
 
1025
                'id': len(decay_mssm.get('interactions'))+2,
 
1026
                'particles': base_objects.ParticleList(
 
1027
                             [decay_mssm.get_particle(25),
 
1028
                              decay_mssm.get_particle(2000013)])})
 
1029
 
 
1030
        decay_mssm.get('interactions').append(new_interaction)
 
1031
        decay_mssm.get('interactions').append(new_interaction_add_sm)
 
1032
 
 
1033
        goal_groups = set([(1,2,3,4,11,12,13,14, 15, 16,21,22,
 
1034
                            23, 24, 25, 35, 36, 37, 2000013), # 15 and from 23
 
1035
                           # are calculated, others are massless default
 
1036
                           (1000005, 1000006, 2000005, 2000006),
 
1037
                           (1000015, 1000016, 2000015),                        
 
1038
                           (1000001, 1000002, 1000003, 1000004, 
 
1039
                            1000021, 2000001, 2000002, 2000003, 2000004),
 
1040
                           (5, 6),
 
1041
                           (1000011, 1000012), 
 
1042
                           (1000013, 1000014), 
 
1043
                           # 2000013 originally should be here, but the
 
1044
                           # the new_interaction_add_sm change it to SM group
 
1045
                           (2000011,)
 
1046
                           ])
 
1047
 
 
1048
        # the stable_candidates that should appear in 1st stage of
 
1049
        # find stable_particles
 
1050
        goal_stable_candidates = [[1,2,3,4,11,12,13,14,16,21,22],
 
1051
                                  [1000006], [1000015], [2000001, 2000003],
 
1052
                                  [5], [1000012], [1000014],[2000011]]
 
1053
        # Group 1,3,4 mixed; group 2, 5, 6 mixed
 
1054
        goal_stable_particle_ids = set([(1,2,3,4,11,12,13,14,16,21,22),
 
1055
                                        (1000015,),
 
1056
                                        # 5 mass = squark
 
1057
                                        # will be set later
 
1058
                                        (2000001, 2000003),
 
1059
                                        (5,),
 
1060
                                        (1000012,),
 
1061
                                        (1000014,),
 
1062
                                        # all sleptons are combine
 
1063
                                        (2000011,)])
 
1064
 
 
1065
        # Get the decay_groups (this should run find_decay_groups_general)
 
1066
        # automatically.
 
1067
        mssm_decay_groups = decay_mssm.get('decay_groups')
 
1068
 
 
1069
        self.assertEqual(set([tuple(sorted([p.get('pdg_code') for p in \
 
1070
                                            group])) \
 
1071
                              for group in mssm_decay_groups]),
 
1072
                         goal_groups)
 
1073
 
 
1074
 
 
1075
        # Test if all useless interactions are deleted.
 
1076
        for inter in decay_mssm['reduced_interactions']:
 
1077
            self.assertTrue(len(inter['particles']))
 
1078
 
 
1079
 
 
1080
        # Test stable particles
 
1081
        # Reset the decay_groups, test the auto-run of find_decay_groups_general
 
1082
        decay_mssm['decay_groups'] = []
 
1083
        decay_mssm.find_stable_particles()
 
1084
 
 
1085
        self.assertEqual(set([tuple(sorted([p.get('pdg_code') for p in plist])) for plist in decay_mssm['stable_particles']]), goal_stable_particle_ids)
 
1086
        
 
1087
 
 
1088
        # Test the assignment of is_stable to particles
 
1089
        goal_stable_pid = [1,2,3,4,5,11,12,13,14,16,21,22,1000012,1000014,
 
1090
                           1000015, 2000001, 2000003, 2000011]
 
1091
        self.assertEqual(sorted([p.get_pdg_code() \
 
1092
                                     for p in decay_mssm.get('particles') \
 
1093
                                     if p.get('is_stable')]), goal_stable_pid)
 
1094
        self.assertTrue(decay_mssm.get_particle(-goal_stable_pid[0]).get('is_stable'))
 
1095
 
 
1096
 
 
1097
        # Test the advance search of stable particles
 
1098
        for p in decay_mssm['particles']:
 
1099
            p['is_stable'] = False
 
1100
        decay_mssm['stable_particles'] = []
 
1101
        decay_mssm.find_stable_particles_advance()
 
1102
        self.assertEqual(sorted([p.get_pdg_code() \
 
1103
                                     for p in decay_mssm.get('particles') \
 
1104
                                     if p.get('is_stable')]), goal_stable_pid)
 
1105
 
 
1106
        goal_stable_particles_ad = set([(1,),(2,),(3,),(4,),(5,),
 
1107
                                        (11,),(12,),(13,),(14,),(16,),
 
1108
                                        (21,),(22,),
 
1109
                                        (1000012,),(1000014,),(1000015,),
 
1110
                                        (2000001,),( 2000003,),( 2000011,)])
 
1111
        self.assertEqual(set([tuple(sorted([p.get('pdg_code') for p in plist])) for plist in decay_mssm['stable_particles']]), goal_stable_particles_ad)
 
1112
 
 
1113
 
 
1114
        
 
1115
    def test_find_full_sm_decay_groups(self):
 
1116
        """ Test the algorithm in find stable particle in full sm.
 
1117
            First, test the full sm with massive neutrinos.
 
1118
            Second, set the neutrinos as massless."""
 
1119
 
 
1120
        # Import the full sm with param_card
 
1121
        sm_path = import_ufo.find_ufo_path('sm')
 
1122
        full_sm_base = import_ufo.import_full_model(sm_path)
 
1123
        full_sm = decay_objects.DecayModel(full_sm_base, True)
 
1124
        param_path = os.path.join(_file_path,
 
1125
                                  '../input_files/param_card_full_sm.dat')
 
1126
        full_sm.read_param_card(param_path)
 
1127
        
 
1128
        # Stage 1: Find stable particles with 
 
1129
        #          nonzero neutrino masses, quark masses
 
1130
        # Turned on light quark and neutrino masses
 
1131
        neutrinos = [12, 14, 16]
 
1132
        lightquarks1 = [1, 3]
 
1133
        for npid in neutrinos:
 
1134
            full_sm.get_particle(npid)['mass'] = '1E-6'
 
1135
            full_sm.get_particle(-npid)['mass'] = '1E-6'
 
1136
        for qpid in lightquarks1:
 
1137
            full_sm.get_particle(qpid)['mass'] = '5'
 
1138
        full_sm.get_particle(2)['mass'] = '1'
 
1139
 
 
1140
        # the stable particles are the ones with lightest mass in
 
1141
        # their group
 
1142
        goal_groups_1 = set([(21,22, 23,25), # 23 and 25 are calculated
 
1143
                             # others are massless default
 
1144
                             (1,), (2,), (3,), (4,), (5,), (6,),
 
1145
                             (11,), (12,), (13,), (14,), (15,), (16,),
 
1146
                             (24,)])
 
1147
        goal_stable_particles_1 = set([(21,22),
 
1148
                                       (2,),
 
1149
                                       (11,), (12,), (14,), (16,)])
 
1150
 
 
1151
        # Test decay groups
 
1152
        self.assertEqual(set([tuple(sorted([p.get('pdg_code') for p in \
 
1153
                                            group])) \
 
1154
                                  for group in full_sm.get('decay_groups')]),
 
1155
                         goal_groups_1)
 
1156
        # Test stable particles
 
1157
        self.assertEqual(set([tuple(sorted([p.get('pdg_code') for p in \
 
1158
                                            group])) \
 
1159
                                 for group in full_sm.get('stable_particles')]),
 
1160
                         goal_stable_particles_1)
 
1161
 
 
1162
 
 
1163
        # Stage 2: turn off the neutrino mass
 
1164
        full_sm.get_particle(12)['mass'] = 'ZERO'
 
1165
        full_sm.get_particle(14)['mass'] = 'ZERO'
 
1166
        full_sm.get_particle(16)['mass'] = 'ZERO'
 
1167
        full_sm.get_particle(-12)['mass'] = 'ZERO'
 
1168
        full_sm.get_particle(-14)['mass'] = 'ZERO'
 
1169
        full_sm.get_particle(-16)['mass'] = 'ZERO'
 
1170
 
 
1171
        full_sm['decay_groups'] = []
 
1172
        full_sm['stable_particles'] = []
 
1173
        for p in full_sm['particles']:
 
1174
            p['is_stable'] = False
 
1175
 
 
1176
        goal_groups_2 = set([(12,14,16,21,22, 23,25), # 23,25 are
 
1177
                             # calculated, others are massless
 
1178
                             (1,), (2,), (3,), (4,), (5,), (6,),
 
1179
                             (11,13,15,24)])
 
1180
        goal_stable_particles_2 = set([(12,14,16,21,22),
 
1181
                                       (2,),
 
1182
                                       (11,)])
 
1183
        goal_stable_pid_2 = [2, 11, 12,14,16,21,22]
 
1184
 
 
1185
        # Test decay groups
 
1186
        self.assertEqual(set([tuple(sorted([p.get('pdg_code') for p in \
 
1187
                                            group])) \
 
1188
                                  for group in full_sm.get('decay_groups')]),
 
1189
                         goal_groups_2)
 
1190
 
 
1191
        # Test stable particles
 
1192
        self.assertEqual(set([tuple(sorted([p.get('pdg_code') for p in \
 
1193
                                            group])) \
 
1194
                                 for group in full_sm.get('stable_particles')]),
 
1195
                         goal_stable_particles_2)
 
1196
 
 
1197
        # Test the assignment of is_stable
 
1198
        self.assertEqual(sorted([p.get_pdg_code() \
 
1199
                                     for p in full_sm.get('particles') \
 
1200
                                     if p.get('is_stable')]), goal_stable_pid_2)
 
1201
 
 
1202
 
 
1203
 
 
1204
    def test_find_full_sm_decay_groups_advance(self):
 
1205
        """ Test the algorithm in find_stable_particles_advance in full sm.
 
1206
            First, test the full sm with massive neutrinos.
 
1207
            Second, set the neutrinos as massless."""
 
1208
 
 
1209
        # Import the full sm with param_card
 
1210
        sm_path = import_ufo.find_ufo_path('sm')
 
1211
        full_sm_base = import_ufo.import_full_model(sm_path)
 
1212
        full_sm = decay_objects.DecayModel(full_sm_base, True)
 
1213
        param_path = os.path.join(_file_path,
 
1214
                                  '../input_files/param_card_full_sm.dat')
 
1215
        full_sm.read_param_card(param_path)
 
1216
        
 
1217
        # Stage 1: Find stable particles with 
 
1218
        #          nonzero neutrino masses, quark masses
 
1219
        # Turned on light quark and neutrino masses
 
1220
        neutrinos = [12, 14, 16]
 
1221
        lightquarks1 = [1, 3]
 
1222
        for npid in neutrinos:
 
1223
            full_sm.get_particle(npid)['mass'] = '1E-6'
 
1224
            full_sm.get_particle(-npid)['mass'] = '1E-6'
 
1225
        for qpid in lightquarks1:
 
1226
            full_sm.get_particle(qpid)['mass'] = '5'
 
1227
        full_sm.get_particle(2)['mass'] = '1'
 
1228
 
 
1229
        # Find stable particles with given spectrum
 
1230
        goal_stable_pid_1 = [2, 11, 12,14,16,21,22]
 
1231
        full_sm.find_stable_particles_advance()
 
1232
        
 
1233
        # Test the assignment of is_stable
 
1234
        self.assertEqual(sorted([p.get_pdg_code() \
 
1235
                                     for p in full_sm.get('particles') \
 
1236
                                     if p.get('is_stable')]), goal_stable_pid_1)
 
1237
        
 
1238
        
 
1239
        # Stage 2: turn off the neutrino mass
 
1240
        full_sm.get_particle(12)['mass'] = 'ZERO'
 
1241
        full_sm.get_particle(14)['mass'] = 'ZERO'
 
1242
        full_sm.get_particle(16)['mass'] = 'ZERO'
 
1243
        full_sm.get_particle(-12)['mass'] = 'ZERO'
 
1244
        full_sm.get_particle(-14)['mass'] = 'ZERO'
 
1245
        full_sm.get_particle(-16)['mass'] = 'ZERO'
 
1246
 
 
1247
        full_sm['decay_groups'] = []
 
1248
        full_sm['stable_particles'] = []
 
1249
        for p in full_sm['particles']:
 
1250
            p['is_stable'] = False
 
1251
        full_sm.find_stable_particles_advance()
 
1252
 
 
1253
        goal_stable_pid_2 = [2, 11, 12,14,16,21,22]
 
1254
        
 
1255
        # Test the assignment of is_stable
 
1256
        self.assertEqual(sorted([p.get_pdg_code() \
 
1257
                                     for p in full_sm.get('particles') \
 
1258
                                     if p.get('is_stable')]), goal_stable_pid_2)
 
1259
            
 
1260
 
 
1261
 
 
1262
    def test_running_couplings(self):
 
1263
        """ Test the running coupling constants in DecayModel.
 
1264
            param_card uses:
 
1265
            Block SMINPUTS 
 
1266
            1 1.279340e+02 # aEWM1 
 
1267
            2 1.166370e-05 # Gf 
 
1268
            3 1.180000e-01 # aS 
 
1269
            4 9.118760e+01 # MMZ 
 
1270
            """
 
1271
 
 
1272
 
 
1273
        # Read mssm
 
1274
        model_base = import_ufo.import_model('mssm')
 
1275
        model = decay_objects.DecayModel(model_base, True)
 
1276
        param_path = os.path.join(_file_path,'../input_files/param_card_mssm.dat')
 
1277
        model.read_param_card(param_path)
 
1278
        #print decay_objects.MZ
 
1279
 
 
1280
        # Test for exception
 
1281
        self.assertRaises(decay_objects.DecayModel.PhysicsObjectError,
 
1282
                          self.my_testmodel.running_externals, "not i")
 
1283
 
 
1284
        # Test for running_externals
 
1285
        # No reference value for q higher than top quark mass
 
1286
 
 
1287
        # Set b quark mass to be consistent with SM model
 
1288
        decay_objects.MB = 4.7
 
1289
        
 
1290
        # q=400., Nf = 5 quarks
 
1291
        model.running_externals(400., 1)
 
1292
        self.assertAlmostEqual(decay_objects.aS, 0.0972887598, 5)
 
1293
        model.running_externals(400.)
 
1294
        self.assertAlmostEqual(decay_objects.aS, 0.096563954696, 5)
 
1295
 
 
1296
        # q=170., Nf = 5 quarks
 
1297
        model.running_externals(170., 1)
 
1298
        self.assertAlmostEqual(decay_objects.aS, 0.1082883202, 6)
 
1299
        model.running_externals(170.)
 
1300
        self.assertAlmostEqual(decay_objects.aS, 0.10788637604, 6)
 
1301
 
 
1302
        # q=10., Nf = 5 quarks
 
1303
        model.running_externals(10., 1)
 
1304
        self.assertAlmostEqual(decay_objects.aS, 0.1730836377, 6)
 
1305
        model.running_externals(10.)
 
1306
        self.assertAlmostEqual(decay_objects.aS, 0.17787426379, 6)
 
1307
 
 
1308
        # q=4., Nf = 4 quarks
 
1309
        model.running_externals(4., 1)
 
1310
        self.assertAlmostEqual(decay_objects.aS, 0.215406025, 6)
 
1311
        model.running_externals(4.)
 
1312
        self.assertAlmostEqual(decay_objects.aS, 0.22772914557, 6)
 
1313
 
 
1314
        # q=1., Nf = 3 quarks
 
1315
        model.running_externals(1., 1)
 
1316
        self.assertAlmostEqual(decay_objects.aS, 0.36145974008, 5)
 
1317
        model.running_externals(1.)
 
1318
        self.assertAlmostEqual(decay_objects.aS, 0.45187971053, 5)
 
1319
 
 
1320
        # Test for running_internals
 
1321
        temp_aS = 100
 
1322
        decay_objects.aS = temp_aS
 
1323
        #Ru11_old = decay_objects.Ru11
 
1324
        Ru11_old = decay_objects.RUU1x1
 
1325
        # coupling of no running dependence
 
1326
        try:
 
1327
            coup0 = model['couplings'][('aEWM1',)][0]
 
1328
        except KeyError:
 
1329
            coup0 = model['couplings'][()][0]        
 
1330
        # coupling depend on aS
 
1331
        coup_aS = model['couplings'][('aS',)][0]
 
1332
        # coupling depend on both aS and aEWM1
 
1333
        try:
 
1334
            coup_both = model['couplings'][('aS', 'aEWM1')][0]
 
1335
        except KeyError:
 
1336
            coup_both = model['couplings'][('aEWM1', 'aS')][0]
 
1337
 
 
1338
        coup0_old = eval('decay_objects.'+coup0.name)
 
1339
        #print coup_aS.name
 
1340
        #print coup_both.name
 
1341
        model.running_internals()
 
1342
 
 
1343
 
 
1344
        # Test for parameters
 
1345
 
 
1346
        # Ru11 should not change
 
1347
        self.assertAlmostEqual(decay_objects.RUU1x1, Ru11_old)
 
1348
 
 
1349
        # G should change
 
1350
        self.assertAlmostEqual(decay_objects.G, \
 
1351
                                   2*cmath.sqrt(temp_aS)*cmath.sqrt(cmath.pi))
 
1352
 
 
1353
 
 
1354
        # Test for couplings
 
1355
 
 
1356
        # GC_365 should not change
 
1357
        self.assertAlmostEqual(eval('decay_objects.'+coup0.name), coup0_old)
 
1358
 
 
1359
        # Both of GC_114 ('aS',) and GC_15 ('aEWSM1', 'aS') should change
 
1360
        self.assertAlmostEqual(eval('decay_objects.'+coup_aS.name), \
 
1361
                                   -decay_objects.G)
 
1362
 
 
1363
        # copying the expr of 
 
1364
        self.assertAlmostEqual(eval('decay_objects.'+coup_both.name), \
 
1365
                                   (-2*decay_objects.ee*complex(0,1)*decay_objects.G*decay_objects.I12x33)/3. - (2*decay_objects.ee*complex(0,1)*decay_objects.G*decay_objects.I13x33)/3.)
 
1366
 
 
1367
#===============================================================================
 
1368
# Test_Channel
 
1369
#===============================================================================
 
1370
class Test_Channel(unittest.TestCase):
 
1371
    """ Test for the channel object"""
 
1372
 
 
1373
    my_channel = decay_objects.Channel()
 
1374
    h_tt_bbmmvv = decay_objects.Channel()
 
1375
 
 
1376
    def setUp(self):
 
1377
        """ Set up necessary objects for the test"""
 
1378
 
 
1379
        if not hasattr(self, 'my_testmodel_base'):
 
1380
            self.my_testmodel_base = import_ufo.import_model('sm')            
 
1381
 
 
1382
        #Import a model from my_testmodel
 
1383
        self.my_testmodel = decay_objects.DecayModel(self.my_testmodel_base, True)
 
1384
        param_path = os.path.join(_file_path,'../input_files/param_card_sm.dat')
 
1385
        self.my_testmodel.read_param_card(param_path)
 
1386
 
 
1387
        # Simplify the model
 
1388
        particles = self.my_testmodel.get('particles')
 
1389
        interactions = self.my_testmodel.get('interactions')
 
1390
        inter_list = copy.copy(interactions)
 
1391
 
 
1392
        # Pids that will be removed
 
1393
        no_want_pid = [1, 2, 3, 4, 15, 16, 21]
 
1394
        for pid in no_want_pid:
 
1395
            particles.remove(self.my_testmodel.get_particle(pid))
 
1396
 
 
1397
        for inter in inter_list:
 
1398
            if any([p.get('pdg_code') in no_want_pid for p in \
 
1399
                        inter.get('particles')]):
 
1400
                interactions.remove(inter)
 
1401
 
 
1402
        # Set a new name
 
1403
        self.my_testmodel.set('name', 'my_smallsm')
 
1404
        self.my_testmodel.set('particles', particles)
 
1405
        self.my_testmodel.set('interactions', interactions)
 
1406
        
 
1407
        # Set up vertexlist for my_testmodel
 
1408
        self.my_testmodel.find_vertexlist()
 
1409
 
 
1410
        #Setup the vertexlist for my_testmodel and save this model (optional)
 
1411
        import_vertexlist.make_vertexlist(self.my_testmodel)
 
1412
        # Report results
 
1413
        #import_vertexlist
 
1414
 
 
1415
        # Save files to check vertices number
 
1416
        # check input_files/model_name/interaction.py for vertex number
 
1417
        #save_model.save_model(os.path.join(MG5DIR, 'tests/input_files', 
 
1418
        #                      self.my_testmodel['name']), self.my_testmodel)
 
1419
    
 
1420
        full_vertexlist = import_vertexlist.full_vertexlist
 
1421
 
 
1422
        # Get vertices
 
1423
        for index, vertex in full_vertexlist.items():
 
1424
            legs_set = set([l['id'] for l in vertex['legs']])
 
1425
            if legs_set == set([-6, 6, 25]) :
 
1426
                # h > t t~
 
1427
                vert_1 = copy.deepcopy(vertex)
 
1428
            elif legs_set == set([5, 24, 6]):
 
1429
                # t~ > b~ w- (decay of antiparticle)
 
1430
                # t > b w+ 
 
1431
                vert_2 = copy.deepcopy(vertex)
 
1432
                vert_3 = copy.deepcopy(vertex)
 
1433
                vert_6 = copy.deepcopy(vertex)
 
1434
            elif legs_set == set([-13, 14, 24]):
 
1435
                # w- > mu- vm~ (decay of antiparticle)
 
1436
                # w+ > mu+ vm
 
1437
                vert_4 = copy.deepcopy(vertex)
 
1438
                vert_5 = copy.deepcopy(vertex)
 
1439
 
 
1440
 
 
1441
        # h > t t~ > b b~ w+ w-
 
1442
        vert_1['legs'][0]['number'] = 2
 
1443
        vert_1['legs'][1]['number'] = 3
 
1444
        vert_1['legs'][2]['number'] = 1
 
1445
        vert_1['legs'][2]['state'] = False
 
1446
 
 
1447
        # t~ > b~ w- (decay of antiparticle)
 
1448
        vert_2['legs'][0]['number'] = 2
 
1449
        vert_2['legs'][1]['number'] = 4
 
1450
        vert_2['legs'][2]['number'] = 2
 
1451
        vert_2['legs'][0]['id'] = -5
 
1452
        vert_2['legs'][1]['id'] = -24
 
1453
        vert_2['legs'][2]['id'] = -6
 
1454
        vert_2['id'] = self.my_testmodel['conj_int_dict'][33]
 
1455
 
 
1456
        # t > b w+
 
1457
        vert_3['legs'][0]['number'] = 3
 
1458
        vert_3['legs'][1]['number'] = 5
 
1459
        vert_3['legs'][2]['number'] = 3
 
1460
 
 
1461
        # w- > mu vm~ (decay of antiparticle)
 
1462
        vert_4['legs'][0]['number'] = 4
 
1463
        vert_4['legs'][1]['number'] = 6
 
1464
        vert_4['legs'][2]['number'] = 4
 
1465
        vert_4['legs'][0]['id'] = 13
 
1466
        vert_4['legs'][1]['id'] = -14
 
1467
        vert_4['legs'][2]['id'] = -24
 
1468
        vert_4['id'] = self.my_testmodel['conj_int_dict'][vert_4['id']]
 
1469
 
 
1470
        # w+ > mu vm~
 
1471
        vert_5['legs'][0]['number'] = 5
 
1472
        vert_5['legs'][1]['number'] = 7
 
1473
        vert_5['legs'][2]['number'] = 5
 
1474
 
 
1475
        #temp_vertices = base_objects.VertexList
 
1476
        self.h_tt_bbmmvv = decay_objects.Channel({'vertices': \
 
1477
                                             base_objects.VertexList([
 
1478
                                             vert_5, vert_4, vert_3, vert_2, \
 
1479
                                             vert_1])})
 
1480
 
 
1481
        #print self.h_tt_bbmmvv.nice_string()
 
1482
        #pic = drawing_eps.EpsDiagramDrawer(self.h_tt_bbmmvv, 'h_tt_bbmmvv', self.my_testmodel)
 
1483
        #pic.draw()
 
1484
 
 
1485
        # t > b w+
 
1486
        vert_6['legs'][0]['number'] = 2
 
1487
        vert_6['legs'][1]['number'] = 3
 
1488
        vert_6['legs'][2]['number'] = 1
 
1489
        # change initial id to anti particle id
 
1490
        vert_6['legs'][2]['id'] = -6
 
1491
 
 
1492
        self.t_bw = decay_objects.Channel({'vertices': \
 
1493
                                           base_objects.VertexList([vert_6])})
 
1494
        #print self.t_bw.nice_string()
 
1495
 
 
1496
    def test_get_initialfinal(self):
 
1497
        """ test the get_initial_id and get_final_legs"""
 
1498
 
 
1499
        # Test the get_initial_id
 
1500
        # Raise error when neither ini_pid is found nor model is given
 
1501
        self.h_tt_bbmmvv['ini_pid'] = 0
 
1502
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
1503
                          self.h_tt_bbmmvv.get_initial_id)
 
1504
 
 
1505
        # After ini_pid is found, no model is needed.
 
1506
        self.assertEqual(self.h_tt_bbmmvv.get_anti_initial_id(), 25)
 
1507
        self.assertEqual(self.h_tt_bbmmvv.get_initial_id(self.my_testmodel), 25)
 
1508
        self.assertEqual(self.h_tt_bbmmvv.get_initial_id(), 25)
 
1509
 
 
1510
        # Test for mother particle which is not self conjugate.
 
1511
        self.t_bw['ini_pid'] = 0
 
1512
        self.assertEqual(self.t_bw.get_anti_initial_id(), -6)
 
1513
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
1514
                          self.t_bw.get_initial_id)
 
1515
        self.assertEqual(self.t_bw.get_initial_id(self.my_testmodel), 6)
 
1516
        self.assertEqual(self.t_bw.get_initial_id(), 6)
 
1517
 
 
1518
 
 
1519
        # Test the get_final_legs
 
1520
        vertexlist = self.h_tt_bbmmvv.get('vertices')
 
1521
        goal_final_legs = base_objects.LegList([vertexlist[0]['legs'][0],
 
1522
                                                vertexlist[0]['legs'][1],
 
1523
                                                vertexlist[1]['legs'][0],
 
1524
                                                vertexlist[1]['legs'][1],
 
1525
                                                vertexlist[2]['legs'][0],
 
1526
                                                vertexlist[3]['legs'][0]])
 
1527
        self.assertEqual(self.h_tt_bbmmvv.get_final_legs(), goal_final_legs)
 
1528
 
 
1529
        # Test get_final_legs with force argument
 
1530
        # Change final_legs manually
 
1531
        self.h_tt_bbmmvv['final_legs'].append(vertexlist[3]['legs'][0])
 
1532
        goal_final_legs.append(vertexlist[3]['legs'][0])
 
1533
        self.assertEqual(self.h_tt_bbmmvv.get_final_legs(), goal_final_legs)
 
1534
        self.assertEqual(self.h_tt_bbmmvv.get_final_legs(True), 
 
1535
                         goal_final_legs[:-1])
 
1536
 
 
1537
    def test_get_onshell(self):
 
1538
        """ test the get_onshell function"""
 
1539
        vertexlist = self.h_tt_bbmmvv.get('vertices')
 
1540
        h_tt_bbww = decay_objects.Channel({'vertices': \
 
1541
                                           base_objects.VertexList(\
 
1542
                                           vertexlist[2:])})
 
1543
        #print h_tt_bbww.nice_string()
 
1544
        # Test for on shell decay ( h > b b~ mu+ mu- vm vm~)
 
1545
        self.assertTrue(self.h_tt_bbmmvv.get_onshell(self.my_testmodel))
 
1546
 
 
1547
        # Test for off-shell decay (h > b b~ w+ w-)
 
1548
        # Raise the mass of higgs
 
1549
        decay_objects.MH = 220
 
1550
        self.assertTrue(h_tt_bbww.get_onshell(self.my_testmodel))
 
1551
 
 
1552
    def test_initial_setups(self):
 
1553
        """ test the intial_setups function"""
 
1554
 
 
1555
        self.h_tt_bbmmvv.initial_setups(self.my_testmodel, True)
 
1556
        # Test for existence of various properties
 
1557
        self.assertTrue(self.h_tt_bbmmvv['onshell'])
 
1558
        self.assertTrue(self.h_tt_bbmmvv['ini_pid'])
 
1559
        self.assertTrue(self.h_tt_bbmmvv['final_legs'])
 
1560
        self.assertTrue(self.h_tt_bbmmvv['orders'])
 
1561
 
 
1562
        t = self.my_testmodel.get_particle(6)
 
1563
        t.find_channels(2, self.my_testmodel)
 
1564
        c = t.get_channels(2, True)[0]
 
1565
 
 
1566
        # Test for existence of various properties
 
1567
        self.assertTrue(c['onshell'])
 
1568
        self.assertTrue(c['ini_pid'])
 
1569
        self.assertTrue(c['final_legs'])
 
1570
        self.assertTrue(c['orders'])
 
1571
 
 
1572
    def test_helper_find_channels(self):
 
1573
        """ Test of the find_channels function of DecayParticle.
 
1574
            Also the test for some helper function for find_channels."""
 
1575
 
 
1576
        higgs = self.my_testmodel.get_particle(25)
 
1577
        t = self.my_testmodel.get_particle(6)
 
1578
 
 
1579
        vertexlist = self.h_tt_bbmmvv.get('vertices')
 
1580
        h_tt_bbww = decay_objects.Channel({'vertices': \
 
1581
                                           base_objects.VertexList(\
 
1582
                                           vertexlist[2:])})
 
1583
        h_tt_bbww.calculate_orders(self.my_testmodel)
 
1584
        self.my_testmodel.find_vertexlist()
 
1585
 
 
1586
        """ Test the connect_channel_vertex """
 
1587
        # 1. connect the w- in h_tt_bbww with w_muvm to get h_tt_bwbmuvm
 
1588
        #    the interaction id should change from w+ decay into w- decay
 
1589
        # 2. connect the w+ in h_tt_bwbmuvm with w_muvm to get h_tt_bbmmvv
 
1590
        h_tt_bwbmuvm = decay_objects.Channel({'vertices': \
 
1591
                                              base_objects.VertexList(\
 
1592
                                              vertexlist[1:])})
 
1593
        # run initial setups
 
1594
        h_tt_bwbmuvm.initial_setups(self.my_testmodel, True)
 
1595
        self.h_tt_bbmmvv.initial_setups(self.my_testmodel, True)
 
1596
 
 
1597
 
 
1598
        # decay of w+ > mu+ vm
 
1599
        for v in self.my_testmodel.get_particle(24).get_vertexlist(2, True):
 
1600
            if set([l['id'] for l in v['legs']]) == set([24, -13, 14]):
 
1601
                w_muvm = v
 
1602
 
 
1603
        # Connect h > w+ w- b b~  with w+ > mu+ vm 
 
1604
        new_channel = higgs.connect_channel_vertex(h_tt_bbww, 3, w_muvm,
 
1605
                                                   self.my_testmodel)
 
1606
 
 
1607
        #print 'source:', h_tt_bbww.nice_string()
 
1608
        #print 'result:', new_channel.nice_string(),\
 
1609
        #    '\ngoal  :', h_tt_bwbmuvm.nice_string(), '\n'
 
1610
        #print self.h_tt_bbmmvv.nice_string()
 
1611
        self.assertEqual(new_channel, h_tt_bwbmuvm)
 
1612
        
 
1613
        # Test the change of legs in mother channel doesn't change resulting
 
1614
        # channel
 
1615
 
 
1616
        # Change final_legs in new_channel
 
1617
        for index, l in enumerate(new_channel['final_legs']):
 
1618
            if l['id'] == 24:
 
1619
                l_index = index
 
1620
                l_num = l['number']
 
1621
                l['number'] = 7
 
1622
        #print new_channel.get_final_legs()
 
1623
 
 
1624
        self.assertFalse(new_channel.get_final_legs()\
 
1625
                             == h_tt_bwbmuvm.get_final_legs())
 
1626
        # reset final_legs to be normal one
 
1627
        new_channel['final_legs'][l_index]['number'] = l_num
 
1628
 
 
1629
        # Test of further connection
 
1630
        #print [l['id'] for l in h_tt_bwbmuvm.get_final_legs()]
 
1631
        new_channel = higgs.connect_channel_vertex(h_tt_bwbmuvm, 3, w_muvm,
 
1632
                                                   self.my_testmodel)
 
1633
        #print new_channel.nice_string(), self.h_tt_bbmmvv.nice_string()
 
1634
        self.assertEqual(new_channel, self.h_tt_bbmmvv)
 
1635
        
 
1636
 
 
1637
 
 
1638
        """ Test of check_idlegs """
 
1639
        # based on t > w+ b, add one more w+
 
1640
        # w+, b, t
 
1641
        temp_vert = copy.deepcopy(vertexlist[2])
 
1642
        temp_vert['legs'].insert(2, temp_vert['legs'][0])
 
1643
        temp_vert['legs'][2]['number'] = 3
 
1644
        #print temp_vert
 
1645
        self.assertEqual(decay_objects.Channel.check_idlegs(vertexlist[2]), {})
 
1646
        self.assertEqual(decay_objects.Channel.check_idlegs(temp_vert),
 
1647
                         {24: [0, 2]})
 
1648
 
 
1649
        # Test of get_idpartlist
 
1650
        # add two b quarks
 
1651
        # result: w+, b, w+, b, b, t
 
1652
        temp_vert2 = copy.deepcopy(temp_vert)
 
1653
        temp_vert2['legs'].insert(3, temp_vert['legs'][1])
 
1654
        temp_vert2['legs'].insert(4, temp_vert['legs'][1])
 
1655
        #print temp_vert2
 
1656
 
 
1657
        # Create a non-sensible channel to test the get_idpartlist
 
1658
        idpart_channel = decay_objects.Channel({'vertices': \
 
1659
                                       base_objects.VertexList([temp_vert])})
 
1660
        idpart_channel = higgs.connect_channel_vertex(idpart_channel, 
 
1661
                                                      1, temp_vert2, 
 
1662
                                                      self.my_testmodel)
 
1663
        #print idpart_c.nice_string()
 
1664
        self.assertEqual(idpart_channel.get_idpartlist(),
 
1665
                         {(1, temp_vert['id'], 24): [0, 2], 
 
1666
                          (0, temp_vert['id'], 24): [0, 2],
 
1667
                          (0, temp_vert['id'], 5):  [1,3,4]})
 
1668
        self.assertTrue(idpart_channel.get('has_idpart'))
 
1669
 
 
1670
 
 
1671
 
 
1672
        """Test of check_channels_equiv """
 
1673
        # Create several fake vertices for test
 
1674
        # h > t t~ t t~ t~
 
1675
        vert_1_id = copy.deepcopy(vertexlist[4])
 
1676
 
 
1677
        vert_1_id.set('id', 800)
 
1678
        vert_1_id.get('legs').insert(2, copy.copy(vert_1_id.get('legs')[0]))
 
1679
        vert_1_id.get('legs').insert(3, copy.copy(vert_1_id.get('legs')[1]))
 
1680
        vert_1_id.get('legs').insert(4, copy.copy(vert_1_id.get('legs')[1]))
 
1681
        vert_1_id.get('legs')[2]['number'] = 4
 
1682
        vert_1_id.get('legs')[3]['number'] = 5
 
1683
        vert_1_id.get('legs')[4]['number'] = 6
 
1684
        #print vert_1_id
 
1685
 
 
1686
        # t > w+ b~ w+
 
1687
        vert_2_id = copy.deepcopy(vertexlist[2])
 
1688
        vert_2_id.set('id', 900)
 
1689
        vert_2_id.get('legs').insert(2, copy.copy(vert_2_id.get('legs')[0]))
 
1690
        vert_2_id.get('legs')[2]['number'] = 0
 
1691
        #print vert_2_id
 
1692
 
 
1693
        # w+ > mu vm is given by vertex[0]
 
1694
 
 
1695
        # w+ > mu ve
 
1696
        w_muve = copy.deepcopy(vertexlist[0])
 
1697
        w_muve.set('id', 1100)
 
1698
        w_muve.get('legs')[0].set('id', -13)
 
1699
        w_muve.get('legs')[1].set('id', 12)
 
1700
 
 
1701
        # w+ > e- vm
 
1702
        w_evmu = copy.deepcopy(vertexlist[0])
 
1703
        w_evmu.set('id', 1200)
 
1704
        w_evmu.get('legs')[0].set('id', -11)
 
1705
        w_evmu.get('legs')[1].set('id', 14)
 
1706
 
 
1707
 
 
1708
        # Update informations in model
 
1709
        for new_inter_id in [800, 900, 1000, 1100, 1200]:
 
1710
            self.my_testmodel.get('interactions').append(\
 
1711
                base_objects.Interaction({'id':new_inter_id }))
 
1712
            self.my_testmodel.get('interactions').append(\
 
1713
                base_objects.Interaction({'id':new_inter_id+1}))
 
1714
            self.my_testmodel['conj_int_dict'][new_inter_id] = new_inter_id+1
 
1715
 
 
1716
        self.my_testmodel.reset_dictionaries()
 
1717
 
 
1718
        # Information about channel a, b, c, for my own convenience.
 
1719
        # Nice string for channel_a:
 
1720
        # ((8(13),12(-14)>8(-24),id:1001), (7(11),11(-12)>7(-24),id:43),
 
1721
        #  (5(-5),10(-24)>5(-6),id:54),(4(5),9(24)>4(6),id:33),
 
1722
        #  (2(-5),7(-24),8(-24)>2(-6),id:901),
 
1723
        #  (2(-6),3(6),4(6),5(-6),6(6)>1(25),id:800)) ()
 
1724
        
 
1725
        # Nice string of channel_b:
 
1726
        #((10(11),12(-12)>10(-24),id:43),(9(13),11(-14)>9(-24),id:1001),
 
1727
        # (6(-5),9(-24),10(-24)>6(-6),id:901),(3(5),8(24)>3(6),id:33),
 
1728
        # (2(-5),7(-24)>2(-6),id:54),
 
1729
        # (2(-6),3(6),4(6),5(-6),6(-6)>2(25),id:800)) ()
 
1730
 
 
1731
        # Nice string of channel_c:
 
1732
        #((10(13),12(-14)>10(-24),id:1001),(9(13),11(-14)>9(-24),id:1001),
 
1733
        # (6(-5),9(-24),10(-24)>6(-6),id:901),(3(5),8(24)>3(6),id:33),
 
1734
        # (2(-5),7(-24)>2(-6),id:54),
 
1735
        # (2(-6),3(6),4(6),5(-6),6(-6)>2(25),id:800)) ()
 
1736
        """ 
 
1737
        Nice string of channel_a,b,c (new leg ordering):
 
1738
        h--t (2) 
 
1739
          \t~(3) > b~(3) w-(7)          w-(8)
 
1740
                         \ mu(7) vm(11) \ mu(8) vm(12)
 
1741
          \t (4) > w+ (4) b(9)
 
1742
          \t~(5) > b~ (5) w-(10)
 
1743
          \t~(6) 
 
1744
        Channel_a:
 
1745
        ((8(13),12(-14)>8(-24),id:44),(7(13),11(-14)>7(-24),id:44),
 
1746
        (5(-5),10(-24)>5(-6),id:54),(4(24),9(5)>4(6),id:33),
 
1747
        (3(-5),7(-24),8(-24)>3(-6),id:901),
 
1748
        (2(6),3(-6),4(6),5(-6),6(-6),1(25),id:800)) (QED=4) 
 
1749
        (est. further width = 0.000e+00)
 
1750
        --------------
 
1751
        Compare to channel a:  t(2) <> t(4); t~(6) <> t~(3)
 
1752
        h--t (2) > w+ (4) b(9)
 
1753
          \t~(3) > b~ (5) w-(10)
 
1754
          \t (4) 
 
1755
          \t~(5) 
 
1756
          \t~(6) > b~(6) w-(9)          w-(10)
 
1757
                         \ mu(9) vm(11) \ mu(10) vm(12)
 
1758
        Channel_b:
 
1759
        ((10(13),12(-14)>10(-24),id:44),(9(13),11(-14)>9(-24),id:44),
 
1760
        (6(-5),9(-24),10(-24)>6(-6),id:901),(3(-5),8(-24)>3(-6),id:54),
 
1761
        (2(24),7(5)>2(6),id:33),
 
1762
        (2(6),3(-6),4(6),5(-6),6(-6),1(25),id:800)) 
 
1763
        (QED=4) (est. further width = 0.000e+00)
 
1764
        --------------
 
1765
        h--t (2) 
 
1766
          \t~(3) > b~(3) w-(7)          w-(8)
 
1767
                         \ mu(7) vm(11) 
 
1768
          \t (4) > w+ (4) b(9)
 
1769
          \t~(5) > b~ (5) w-(10)
 
1770
                          \ mu(10) vm(12)
 
1771
          \t~(6) 
 
1772
        Channel_c:
 
1773
        ((10(13),12(-14)>10(-24),id:44),(7(13),11(-14)>7(-24),id:44),
 
1774
        (5(-5),10(-24)>5(-6),id:54),(4(24),9(5)>4(6),id:33),
 
1775
        (3(-5),7(-24),8(-24)>3(-6),id:901),
 
1776
        (2(6),3(-6),4(6),5(-6),6(-6),1(25),id:800)) 
 
1777
        (QED=4) (est. further width = 0.000e+00)
 
1778
        """
 
1779
 
 
1780
 
 
1781
        # Initiate channel_a
 
1782
        # h > t t~ t t~ t~
 
1783
        channel_a = decay_objects.Channel({'vertices': base_objects.VertexList(\
 
1784
                    [vert_1_id])})
 
1785
 
 
1786
        # Add t~ > b~ w- w- to first t~
 
1787
        channel_a = higgs.connect_channel_vertex(channel_a, 1,
 
1788
                                                 vert_2_id,
 
1789
                                                 self.my_testmodel)
 
1790
        #print channel_a.nice_string()
 
1791
 
 
1792
        # Add t > b w+ to 2nd t
 
1793
        channel_a = higgs.connect_channel_vertex(channel_a, 4,
 
1794
                                                 vertexlist[2],
 
1795
                                                 self.my_testmodel)
 
1796
        #print channel_a.nice_string()
 
1797
 
 
1798
        # Add t~ > b~ w- to 2nd t~
 
1799
        channel_a = higgs.connect_channel_vertex(channel_a, 6,
 
1800
                                                 vertexlist[2],
 
1801
                                                 self.my_testmodel)
 
1802
        #print channel_a.nice_string()
 
1803
 
 
1804
        # Add w- > mu- vm~ to first w- in t~ > b~ w- w- decay chain
 
1805
        channel_a = higgs.connect_channel_vertex(channel_a, 5,
 
1806
                                                 w_muvm,
 
1807
                                                 self.my_testmodel)
 
1808
        #print channel_a.nice_string()
 
1809
 
 
1810
        # Add w- > mu vm~ to 2nd w- in t~ > b~ w- w- decay chain
 
1811
        channel_a = higgs.connect_channel_vertex(channel_a, 7,
 
1812
                                                 w_muvm,
 
1813
                                                 self.my_testmodel)
 
1814
        #print 'Channel_a:\n', channel_a.nice_string()
 
1815
        
 
1816
 
 
1817
 
 
1818
        # Initiate channel_b
 
1819
        # h > t t~ t t~ t~
 
1820
        channel_b = decay_objects.Channel({'vertices': base_objects.VertexList(\
 
1821
                    [vert_1_id])})
 
1822
 
 
1823
        # Add t > b w+ to 1st t
 
1824
        channel_b = higgs.connect_channel_vertex(channel_b, 0,
 
1825
                                                 vertexlist[2],
 
1826
                                                 self.my_testmodel)
 
1827
        #print '\n', channel_b.nice_string()
 
1828
 
 
1829
        # Add t~ > b~ w- to 1st t~
 
1830
        channel_b = higgs.connect_channel_vertex(channel_b, 2,
 
1831
                                                 vertexlist[2],
 
1832
                                                 self.my_testmodel)
 
1833
        #print channel_b.nice_string()
 
1834
 
 
1835
        # Add t~ > b~ w- w- to final t~
 
1836
        channel_b = higgs.connect_channel_vertex(channel_b, 6,
 
1837
                                                 vert_2_id,
 
1838
                                                 self.my_testmodel)
 
1839
        #print channel_b.nice_string()
 
1840
 
 
1841
        # Add w- > mu vm~ to 1st w- in t~ decay chain
 
1842
        channel_b = higgs.connect_channel_vertex(channel_b, 1,
 
1843
                                                 w_muvm,
 
1844
                                                 self.my_testmodel)
 
1845
        #print channel_b.nice_string()
 
1846
 
 
1847
        # Add w- > mu vm~ to 2nd w- in t~ decay chain
 
1848
        channel_b = higgs.connect_channel_vertex(channel_b, 3,
 
1849
                                                 w_muvm,
 
1850
                                                 self.my_testmodel)
 
1851
        #print 'Channel_b:\n', channel_b.nice_string()
 
1852
 
 
1853
 
 
1854
        # Initiate channel_c
 
1855
        # h > t t~ t t~ t~
 
1856
        channel_c = decay_objects.Channel({'vertices': base_objects.VertexList(\
 
1857
                    [vert_1_id])})
 
1858
 
 
1859
        # Add t~ > b~ w- w- to first t~
 
1860
        channel_c = higgs.connect_channel_vertex(channel_c, 1,
 
1861
                                                 vert_2_id,
 
1862
                                                 self.my_testmodel)
 
1863
        #print channel_a.nice_string()
 
1864
 
 
1865
        # Add t > b w+ to 2nd t
 
1866
        channel_c = higgs.connect_channel_vertex(channel_c, 4,
 
1867
                                                 vertexlist[2],
 
1868
                                                 self.my_testmodel)
 
1869
        #print channel_a.nice_string()
 
1870
 
 
1871
        # Add t~ > b~ w- to 2nd t~
 
1872
        channel_c = higgs.connect_channel_vertex(channel_c, 6,
 
1873
                                                 vertexlist[2],
 
1874
                                                 self.my_testmodel)
 
1875
 
 
1876
        # Add w- > mu vm~ to 1st w- in t~ decay chain
 
1877
        channel_c = higgs.connect_channel_vertex(channel_c, 5,
 
1878
                                                 w_muvm,
 
1879
                                                 self.my_testmodel)
 
1880
        #print channel_c.nice_string()
 
1881
 
 
1882
        # Add w- > mu vm~ to 2nd w- in t~ decay chain
 
1883
        channel_c = higgs.connect_channel_vertex(channel_c, 3,
 
1884
                                                 w_muvm,
 
1885
                                                 self.my_testmodel)
 
1886
        #print 'Channel_c:\n', channel_c.nice_string()
 
1887
 
 
1888
 
 
1889
        """ Test for using DiagramTag """
 
1890
        Tag_a = diagram_generation.DiagramTag(channel_a)
 
1891
        Tag_b = diagram_generation.DiagramTag(channel_b)
 
1892
        Tag_c = diagram_generation.DiagramTag(channel_c)
 
1893
        #print 'Channel_a:\n', channel_a.nice_string(), '\n', Tag_a
 
1894
        #print 'Channel_b:\n', channel_b.nice_string(), '\n', Tag_b
 
1895
        #print 'Channel_c:\n', channel_c.nice_string(), '\n', Tag_c
 
1896
 
 
1897
        # Test the get function
 
1898
        self.assertFalse(channel_a['tag'])
 
1899
        self.assertEqual(channel_a.get('tag'), Tag_a)
 
1900
        self.assertTrue(Tag_a == Tag_b)
 
1901
        self.assertFalse(Tag_a == Tag_c)
 
1902
 
 
1903
        # Test the check_channels_equiv function
 
1904
        self.assertTrue(decay_objects.Channel.check_channels_equiv(channel_a,
 
1905
                                                                   channel_b))
 
1906
        self.assertFalse(decay_objects.Channel.check_channels_equiv(channel_a,
 
1907
                                                                    channel_c))
 
1908
        self.assertFalse(decay_objects.Channel.check_channels_equiv(channel_a,
 
1909
                                                                    self.h_tt_bbmmvv))
 
1910
 
 
1911
 
 
1912
        """ Test of check_channels_equiv for fermionic mother """
 
1913
        # Change initial pid
 
1914
        vert_2_id['legs'][-1]['id'] = -6
 
1915
        # Correct the leg numbers
 
1916
        vert_2_id['legs'][0]['number'] = 2
 
1917
        vert_2_id['legs'][1]['number'] = 3
 
1918
        vert_2_id['legs'][2]['number'] = 4
 
1919
        vert_2_id['legs'][-1]['number'] = 1
 
1920
 
 
1921
        # t > w+ b~ w+, w+ > mu+ vm, w+ > e+ ve
 
1922
        channel_d = decay_objects.Channel({'vertices': base_objects.VertexList(\
 
1923
                    [vert_2_id])})
 
1924
        channel_d = higgs.connect_channel_vertex(channel_d, 0,
 
1925
                                                 w_muvm,
 
1926
                                                 self.my_testmodel)        
 
1927
        channel_d = higgs.connect_channel_vertex(channel_d, 3,
 
1928
                                                 vertexlist[0],
 
1929
                                                 self.my_testmodel)        
 
1930
        # t > w+ b~ w+, w+ > e+ ve, w+ > mu+ vm 
 
1931
        channel_e = decay_objects.Channel({'vertices': base_objects.VertexList(\
 
1932
                    [vert_2_id])})
 
1933
        channel_e = higgs.connect_channel_vertex(channel_e, 0,
 
1934
                                                 vertexlist[0],
 
1935
                                                 self.my_testmodel)        
 
1936
        channel_e = higgs.connect_channel_vertex(channel_e, 3,
 
1937
                                                 w_muvm,
 
1938
                                                 self.my_testmodel)
 
1939
        # t > w+ b~ w+, w+ > e+ vm, w+ > mu+ ve
 
1940
        channel_f = decay_objects.Channel({'vertices': base_objects.VertexList(\
 
1941
                    [vert_2_id])})
 
1942
        channel_f = higgs.connect_channel_vertex(channel_f, 0,
 
1943
                                                 w_evmu,
 
1944
                                                 self.my_testmodel)        
 
1945
        channel_f = higgs.connect_channel_vertex(channel_f, 3,
 
1946
                                                 w_muve,
 
1947
                                                 self.my_testmodel)
 
1948
        
 
1949
        """print "channel_d:", channel_d.nice_string(), '\n',\
 
1950
            "channel_e:", channel_e.nice_string(), '\n',\
 
1951
            "channel_f:", channel_f.nice_string()"""
 
1952
        self.assertTrue(decay_objects.Channel.check_channels_equiv(channel_d,
 
1953
                                                                   channel_e))
 
1954
        self.assertFalse(decay_objects.Channel.check_channels_equiv(channel_d,
 
1955
                                                                   channel_f))
 
1956
        self.assertFalse(decay_objects.Channel.check_channels_equiv(channel_e,
 
1957
                                                                   channel_f))
 
1958
        Tag_d = diagram_generation.DiagramTag(channel_d)
 
1959
        Tag_e = diagram_generation.DiagramTag(channel_e)
 
1960
        Tag_f = diagram_generation.DiagramTag(channel_f)
 
1961
 
 
1962
        self.assertTrue(Tag_d == Tag_e)
 
1963
        self.assertFalse(Tag_d == Tag_f)
 
1964
        self.assertFalse(Tag_e == Tag_f)
 
1965
 
 
1966
        
 
1967
    def test_findchannels(self):
 
1968
        """ Test of the find_channels functions."""
 
1969
 
 
1970
        higgs = self.my_testmodel.get_particle(25)
 
1971
        t = self.my_testmodel.get_particle(6)
 
1972
        t.find_channels(2, self.my_testmodel)
 
1973
        #print t.get_channels(2, True).nice_string()
 
1974
 
 
1975
        # Test exceptions of find_channels
 
1976
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
1977
                          higgs.find_channels,
 
1978
                          'non_int', self.my_testmodel)
 
1979
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
1980
                          higgs.find_channels,
 
1981
                          4, higgs)
 
1982
        non_sm = copy.copy(higgs)
 
1983
        non_sm.set('pdg_code', 800)
 
1984
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
1985
                          higgs.find_channels,
 
1986
                          non_sm, self.my_testmodel)
 
1987
 
 
1988
        # Test for initial id: anti id and state
 
1989
        for c in t.get_channels(2, True):
 
1990
            self.assertEqual(c['vertices'][-1]['legs'][-1]['id'], -6)
 
1991
            self.assertEqual(c['vertices'][-1]['legs'][-1]['state'], False)
 
1992
            self.assertEqual(c.get_initial_id(), 6)
 
1993
 
 
1994
 
 
1995
        # Create two equivalent channels,
 
1996
        # check whether if only one of them is found.
 
1997
 
 
1998
        # Get vertices
 
1999
        for index, vertex in import_vertexlist.full_vertexlist.items():
 
2000
            legs_set = set([l['id'] for l in vertex['legs']])
 
2001
            if legs_set == set([23, 23, 25]) :
 
2002
                # h > z z
 
2003
                vert_1 = copy.deepcopy(vertex)
 
2004
            elif legs_set == set([11, -11, 23]):
 
2005
                # z > e- e+ 
 
2006
                vert_2 = copy.deepcopy(vertex)
 
2007
            elif legs_set == set([24, -24, 25]):
 
2008
                # h > w+ w-
 
2009
                vert_3 = copy.deepcopy(vertex)
 
2010
            elif legs_set == set([-11, 12, 24]):
 
2011
                # w+ > e+ ve
 
2012
                vert_4 = copy.deepcopy(vertex)
 
2013
 
 
2014
        # 1.) h > z (z > e e~)
 
2015
        vert_0 = self.h_tt_bbmmvv.get('vertices')[-1] 
 
2016
        vert_1['legs'][0]['number'] = 2
 
2017
        vert_1['legs'][1]['number'] = 3
 
2018
        vert_1['legs'][2]['number'] = 1
 
2019
        vert_1['legs'][2]['state'] = False
 
2020
        #print vert_1, vert_2
 
2021
 
 
2022
        channel_a = decay_objects.Channel({'vertices': base_objects.VertexList(\
 
2023
                    [vert_1])})
 
2024
        channel_b = decay_objects.Channel({'vertices': base_objects.VertexList(\
 
2025
                    [vert_1])})        
 
2026
        channel_a = higgs.connect_channel_vertex(channel_a, 0, vert_2,
 
2027
                                                self.my_testmodel)
 
2028
        channel_b = higgs.connect_channel_vertex(channel_b, 1, vert_2,
 
2029
                                                self.my_testmodel)
 
2030
        channel_a.get_apx_decaywidth(self.my_testmodel)
 
2031
        channel_b.get_apx_decaywidth(self.my_testmodel)
 
2032
        channel_a.get('tag')
 
2033
        channel_b.get('tag')
 
2034
        channel_a['has_idpart']=True
 
2035
        channel_b['has_idpart']=True
 
2036
        #print channel_a.nice_string(), '\n', channel_b.nice_string()
 
2037
 
 
2038
 
 
2039
        # 2.) h > w+ w- > e+ ve~ e- ve
 
2040
        vert_3['legs'][0]['number'] = 2
 
2041
        vert_3['legs'][1]['number'] = 3
 
2042
        vert_3['legs'][2]['number'] = 1
 
2043
        vert_3['legs'][2]['state'] = False
 
2044
 
 
2045
        channel_c = decay_objects.Channel({'vertices': base_objects.VertexList(\
 
2046
                    [vert_3])})
 
2047
        channel_d = decay_objects.Channel({'vertices': base_objects.VertexList(\
 
2048
                    [vert_3])})        
 
2049
        channel_c = higgs.connect_channel_vertex(channel_c, 1, vert_4,
 
2050
                                                self.my_testmodel)
 
2051
        channel_c = higgs.connect_channel_vertex(channel_c, 2, vert_4,
 
2052
                                                self.my_testmodel)
 
2053
        channel_d = higgs.connect_channel_vertex(channel_d, 0, vert_4,
 
2054
                                                self.my_testmodel)
 
2055
        channel_d = higgs.connect_channel_vertex(channel_d, 2, vert_4,
 
2056
                                                self.my_testmodel)
 
2057
        #print channel_c.nice_string(), '\n', channel_d.nice_string()
 
2058
        channel_c.calculate_orders(self.my_testmodel)
 
2059
        channel_d.calculate_orders(self.my_testmodel)
 
2060
        channel_c.get('tag')
 
2061
        channel_d.get('tag')
 
2062
        
 
2063
 
 
2064
        # Test of find_channels
 
2065
        # The program should run find_stable_particles automatically.
 
2066
        self.my_testmodel.find_channels(self.my_testmodel.get_particle(5), 3)
 
2067
        self.assertFalse(self.my_testmodel.get_particle(5).get_channels(3, True))
 
2068
        self.assertTrue(self.my_testmodel['stable_particles'])
 
2069
 
 
2070
 
 
2071
        higgs.find_channels(3, self.my_testmodel)
 
2072
        higgs.find_channels_nextlevel(self.my_testmodel)
 
2073
        result1 = higgs.get_channels(3, True)
 
2074
        #print result1.nice_string()
 
2075
 
 
2076
 
 
2077
        # Test if the equivalent channels appear only once.
 
2078
        # For both has_idpart and not has_idpart channels
 
2079
        self.assertEqual((result1.count(channel_b)+ result1.count(channel_a)),1)
 
2080
 
 
2081
        # Test if the 2bodymassdiff calculated during
 
2082
        # find_channels_nextlevel is right.
 
2083
        self.assertAlmostEqual(higgs.get('2body_massdiff'), decay_objects.MH-2*decay_objects.MB)
 
2084
 
 
2085
 
 
2086
        # Set MH < MW to get the desire channels.
 
2087
        decay_objects.MH = 50
 
2088
        channel_c.get_apx_decaywidth(self.my_testmodel)
 
2089
        channel_d.get_apx_decaywidth(self.my_testmodel)
 
2090
        higgs['decay_channels'] = {}
 
2091
        higgs.find_channels(3, self.my_testmodel)
 
2092
        higgs.find_channels_nextlevel(self.my_testmodel)
 
2093
        result2 = higgs.get_channels(4, True)
 
2094
        #print result2.nice_string()
 
2095
        self.assertEqual((result2.count(channel_c)+ result2.count(channel_d)),1)
 
2096
 
 
2097
 
 
2098
        """ Test on MSSM, to get a feeling on the execution time. """        
 
2099
        mssm = import_ufo.import_model('mssm')
 
2100
        param_path = os.path.join(_file_path,'../../models/mssm/restrict_default.dat')
 
2101
        decay_mssm = decay_objects.DecayModel(mssm, force=True)
 
2102
        decay_mssm.read_param_card(param_path)
 
2103
        
 
2104
        susy_higgs = decay_mssm.get_particle(25)
 
2105
        susy_higgs.find_channels(3, decay_mssm)
 
2106
        #susy_higgs.find_channels_nextlevel(decay_mssm)
 
2107
        #decay_mssm.find_all_channels(3)
 
2108
 
 
2109
        # Test the calculation of branching ratios
 
2110
        # The total br should be unity
 
2111
        total_br = sum([amp['apx_br'] for amp in susy_higgs.get_amplitudes(2)])
 
2112
        total_br += sum([amp['apx_br'] for amp in susy_higgs.get_amplitudes(3)])
 
2113
        self.assertAlmostEqual(total_br, 1.)
 
2114
 
 
2115
        # Test if the err is from the off-shell 3-body channels
 
2116
        err = sum([c['apx_decaywidth_nextlevel'] \
 
2117
                       for c in susy_higgs.get_channels(3, False)])
 
2118
        self.assertAlmostEqual(err, susy_higgs['apx_decaywidth_err'])
 
2119
 
 
2120
                                           
 
2121
    def test_apx_decaywidth(self):
 
2122
        """ Test for the approximation of decay rate, including
 
2123
            get_apx_fnrule, get_apx_matrixelement_sq (onshell and offshell),
 
2124
            get_apx_psarea"""
 
2125
 
 
2126
        # Import SM
 
2127
 
 
2128
        full_sm_base = import_ufo.import_model('sm')
 
2129
        full_sm = decay_objects.DecayModel(full_sm_base, True)
 
2130
        # save interaction/particle content into input_files/sm
 
2131
        #save_model.save_model(os.path.join(MG5DIR, 'tests/input_files', 
 
2132
        #                                   full_sm['name']), full_sm)
 
2133
 
 
2134
        higgs = self.my_testmodel.get_particle(25)
 
2135
 
 
2136
        # Set the higgs mass < Z-boson mass so that identicle particles appear
 
2137
        # in final state
 
2138
        MH_new = 91
 
2139
        decay_objects.MH = MH_new
 
2140
        higgs.find_channels(4, self.my_testmodel)
 
2141
 
 
2142
 
 
2143
        # Choose h > w- e+ ve
 
2144
 
 
2145
        for c in higgs.get_channels(3, True):
 
2146
            final_ids = set([l.get('id') for l in c.get_final_legs()])
 
2147
            if final_ids == set([-11, 12, -24]):
 
2148
                h_ww_weve = c
 
2149
        #print h_ww_weve.nice_string()
 
2150
 
 
2151
        # Choose h > w- e+ ve, h > w- t b~
 
2152
 
 
2153
        for c in higgs.get_channels(3, False):
 
2154
            final_ids = set([l.get('id') for l in c.get_final_legs()])
 
2155
            if final_ids == set([-5, 5, 23]):
 
2156
                h_zz_zbb = c
 
2157
            # Distiguish h > ww > w t b~ from h > tt~ > t w b~
 
2158
            if final_ids == set([-5, 6, -24]) \
 
2159
                    and abs(c['vertices'][-1]['legs'][0]['id']) == 24:
 
2160
                h_ww_wtb = c
 
2161
        #print h_ww_wtb.nice_string(), h_zz_zbb.nice_string()
 
2162
 
 
2163
 
 
2164
        # Choose h > 2e+ 2e-, h > vt~ vt b b~
 
2165
        """print higgs.get_channels(3, True).nice_string()
 
2166
        print higgs.get_channels(4, True).nice_string()"""
 
2167
 
 
2168
        for c in higgs.get_channels(4, True):
 
2169
            final_ids = set([l.get('id') for l in c.get_final_legs()])
 
2170
            #print final_ids
 
2171
            if final_ids == set([11, -11, 11, -11]):
 
2172
                h_zz_2epairs = c
 
2173
            if final_ids == set([14, -14, 5, -5]):
 
2174
                h_zz_bbvtvt  = c
 
2175
        #print h_zz_bbvtvt.nice_string(), h_zz_2epairs.nice_string()
 
2176
        
 
2177
        
 
2178
        # Test of the symmetric factor
 
2179
 
 
2180
        h_zz_2epairs.get_apx_psarea(self.my_testmodel)
 
2181
        h_zz_bbvtvt.get_apx_psarea(self.my_testmodel)
 
2182
        self.assertEqual(4, h_zz_2epairs['s_factor'])
 
2183
        self.assertEqual(1, h_zz_bbvtvt['s_factor'])
 
2184
 
 
2185
 
 
2186
        # Test of the get_apx_fnrule
 
2187
 
 
2188
        MW = h_ww_weve.get('final_mass_list')[-1]
 
2189
        #print h_ww_weve.get('final_mass_list')
 
2190
        #print h_ww_weve.get_apx_matrixelement_sq(self.my_testmodel)
 
2191
        #print 'Vertor boson, onshell:', \
 
2192
        #    h_ww_weve.get_apx_fnrule(24, 0.5,
 
2193
        #                            False, self.my_testmodel)
 
2194
        q_offshell = 10
 
2195
        q_offshell_2 = 88
 
2196
        q_onshell = 200
 
2197
 
 
2198
        # Spin 1
 
2199
        self.assertAlmostEqual(h_ww_weve.get_apx_fnrule(24, q_onshell, 
 
2200
                                                  True, self.my_testmodel),
 
2201
                         (1+1/(MW ** 2)*q_onshell **2))
 
2202
        self.assertAlmostEqual(h_ww_weve.get_apx_fnrule(24, q_offshell, 
 
2203
                                                  False, self.my_testmodel),
 
2204
                          ((1-2*((q_offshell/MW) ** 2)+(q_offshell/MW) ** 4)/ \
 
2205
                               ((q_offshell**2-MW **2)**2)))
 
2206
        # Fermion
 
2207
        self.assertEqual(h_ww_weve.get_apx_fnrule(11, q_onshell, 
 
2208
                                                  True, self.my_testmodel),
 
2209
                         q_onshell*2)
 
2210
 
 
2211
        self.assertEqual(h_ww_weve.get_apx_fnrule(6, q_onshell, 
 
2212
                                                  True, self.my_testmodel),
 
2213
                         q_onshell*2)
 
2214
        self.assertAlmostEqual(h_ww_weve.get_apx_fnrule(6, q_offshell, 
 
2215
                                                  False, self.my_testmodel),
 
2216
                         q_offshell**2/(q_offshell ** 2 - decay_objects.MT **2)\
 
2217
                             ** 2)
 
2218
        # Scalar
 
2219
        self.assertEqual(h_ww_weve.get_apx_fnrule(25, q_onshell, 
 
2220
                                                  True, self.my_testmodel),
 
2221
                         1)
 
2222
 
 
2223
        self.assertAlmostEqual(h_ww_weve.get_apx_fnrule(25, q_offshell_2, 
 
2224
                                                  False, self.my_testmodel),
 
2225
                               1/(q_offshell_2 ** 2 - MH_new ** 2)**2, 5)
 
2226
 
 
2227
 
 
2228
 
 
2229
        # Test of matrix element square calculation
 
2230
 
 
2231
        E_mean = (MH_new-MW)/3
 
2232
 
 
2233
        """
 
2234
        print h_ww_weve.get_apx_fnrule(-24, 2*E_mean, False, full_sm)
 
2235
        print self.my_testmodel.get_interaction(7), 
 
2236
        self.my_testmodel.get_interaction(66)
 
2237
        print h_ww_weve.get_apx_fnrule(24, E_mean+MW, True, self.my_testmodel)
 
2238
        """
 
2239
 
 
2240
        g_hww, g_wev = None, None
 
2241
        for candidate in full_sm_base['interactions']:
 
2242
            if [p['pdg_code'] for p in candidate['particles']] == [24, 24, 25]:
 
2243
                g_hww = candidate['couplings'][(0,0)]
 
2244
            if [p['pdg_code'] for p in candidate['particles']] == [12,11,24]:
 
2245
                g_wev = candidate['couplings'][(0,0)]
 
2246
 
 
2247
        
 
2248
 
 
2249
        #print self.my_testmodel.get_interaction(7), 
 
2250
        #self.my_testmodel.get_interaction(63)
 
2251
        self.assertAlmostEqual(
 
2252
            h_ww_weve.get_apx_matrixelement_sq(self.my_testmodel),
 
2253
            ((E_mean**2*4*(1-2*(2*E_mean/MW)**2+(2*E_mean/MW)**4)\
 
2254
                  /(((2*E_mean)**2-MW **2)**2))*\
 
2255
                 (1+(1/MW*(E_mean+MW))**2)*\
 
2256
                 abs(getattr(decay_objects,g_wev)) **2*\
 
2257
                 abs(getattr(decay_objects,g_hww)) **2)
 
2258
            )
 
2259
 
 
2260
 
 
2261
        # Test color multiplicity
 
2262
 
 
2263
        # Choose hadronic decay and leptonic decay of tau
 
2264
        tau = full_sm.get_particle(15)
 
2265
        MTAU = abs(eval('decay_objects.' + tau.get('mass')))
 
2266
        tau.find_channels(3, full_sm)
 
2267
        for c in tau.get_channels(3, True):
 
2268
            final_ids = set([l.get('id') for l in c.get_final_legs()])
 
2269
            #print final_ids
 
2270
            if final_ids == set([-2, 1, 16]):
 
2271
                tau_qdecay = c
 
2272
            if final_ids == set([-14, 13, 16]):
 
2273
                tau_ldecay = c
 
2274
        """print tau_qdecay.nice_string(), tau_ldecay.nice_string()"""
 
2275
 
 
2276
 
 
2277
        # width of hadronic decay ~ 3* leptonic
 
2278
        self.assertAlmostEqual(tau_qdecay.get_apx_decaywidth(full_sm),
 
2279
                               3*tau_ldecay.get_apx_decaywidth(full_sm))
 
2280
 
 
2281
        # Using the coupling constant of w > e ve
 
2282
        self.assertAlmostEqual(tau_qdecay.get_apx_matrixelement_sq(full_sm),
 
2283
                               ((MTAU/3) **3 *8*3*MTAU*
 
2284
                                (1-2*(2*MTAU/(3*MW))**2 +(2*MTAU/(3*MW))**4)/ \
 
2285
                                    ((2*MTAU/3) ** 2 - MW **2) **2 *\
 
2286
                                    abs(getattr(decay_objects,g_wev)) **4))
 
2287
 
 
2288
 
 
2289
 
 
2290
        # Test for off-shell estimation of matrix element
 
2291
 
 
2292
        # matrix element: replacing some of the energy into E_est_mean
 
2293
        E_est_mean = MH_new/4
 
2294
        self.assertAlmostEqual(h_ww_wtb.get_apx_matrixelement_sq(full_sm),
 
2295
                               h_ww_wtb.get_apx_fnrule(5, E_est_mean,
 
2296
                                                       True, full_sm)* \
 
2297
                                   h_ww_wtb.get_apx_fnrule(-6, E_est_mean,
 
2298
                                                      True, full_sm)* \
 
2299
                                   h_ww_wtb.get_apx_fnrule(24, E_est_mean,
 
2300
                                                           True, full_sm)* \
 
2301
                                   h_ww_wtb.get_apx_fnrule(24, MH_new,
 
2302
                                                           False, full_sm, \
 
2303
                                                               est = True)*\
 
2304
                                   h_ww_wtb.get_apx_fnrule(25, MH_new,
 
2305
                                                           True, full_sm)*\
 
2306
                                   abs(getattr(decay_objects,g_hww)) **2 *\
 
2307
                                   abs(getattr(decay_objects,g_wev)) **2)
 
2308
 
 
2309
        
 
2310
 
 
2311
        # Test of phase space area calculation
 
2312
 
 
2313
        # On-shell phase space area
 
2314
        #print 'Tau decay ps_area', tau_qdecay.get_apx_psarea(full_sm)
 
2315
        self.assertAlmostEqual(tau_qdecay.calculate_apx_psarea(1.777, [0,0]),
 
2316
                                1/(8*math.pi))
 
2317
        self.assertAlmostEqual(tau_qdecay.calculate_apx_psarea(1.777, [0,0,0]),
 
2318
                               0.000477383, 5)
 
2319
        self.assertAlmostEqual(h_ww_weve.get_apx_psarea(full_sm), 
 
2320
                               0.0042786859,5)
 
2321
 
 
2322
        # Estimation for off-shell case
 
2323
        self.assertAlmostEqual(h_ww_wtb.get_apx_psarea(full_sm),
 
2324
                                1/512/math.pi **3 * 0.8 * MH_new **2)
 
2325
 
 
2326
 
 
2327
 
 
2328
        # Test of decay width = matrix element * phase space are
 
2329
 
 
2330
        # tau hadronic decay
 
2331
        self.assertAlmostEqual(tau_qdecay.get_apx_decaywidth(full_sm),
 
2332
                               (0.5/1.777)*tau_qdecay.get_apx_psarea(full_sm)*\
 
2333
                                   tau_qdecay.get_apx_matrixelement_sq(full_sm))
 
2334
 
 
2335
 
 
2336
 
 
2337
        # Test of the estimated higher level width for off-shell channel
 
2338
 
 
2339
        # Channels with no next-level decay
 
2340
        self.assertEqual(h_ww_wtb.get_apx_decaywidth_nextlevel(full_sm), 0.)
 
2341
 
 
2342
        # find all channels, width of particles are updated automatically,
 
2343
        # couplings are also ran according to mother particles.
 
2344
        full_sm.find_all_channels(3)
 
2345
        WZ = full_sm.get_particle(23).get('apx_decaywidth')
 
2346
        WW = full_sm.get_particle(24).get('apx_decaywidth')
 
2347
 
 
2348
        # Offshell case: Brett-Wigner correction of propagator
 
2349
 
 
2350
        self.assertAlmostEqual(h_ww_weve.get_apx_fnrule(24, q_offshell, 
 
2351
                                                 False, full_sm),
 
2352
                               ((1-2*((q_offshell/MW) ** 2)+(q_offshell/MW) ** 4)/ \
 
2353
                             (((q_offshell**2-MW **2)**2+MW**2*WW**2)))
 
2354
                               )
 
2355
 
 
2356
        # Channels with next-level decay: h > z b b~
 
2357
        ratio = (1+ WZ*abs(decay_objects.MZ)/MH_new*\
 
2358
                     (1/4/math.pi)*MH_new **3 *0.8/ \
 
2359
                     h_zz_zbb.get_apx_fnrule(23, decay_objects.MZ,
 
2360
                                             True, full_sm)/ \
 
2361
                     h_zz_zbb.get_apx_fnrule(23, 0.5*MH_new,
 
2362
                                             True, full_sm)* \
 
2363
                     h_zz_zbb.get_apx_fnrule(23, MH_new,
 
2364
                                             False, full_sm, True))
 
2365
 
 
2366
        self.assertAlmostEqual(\
 
2367
            h_zz_zbb.get_apx_decaywidth(full_sm)*(ratio-1),
 
2368
            h_zz_zbb.get_apx_decaywidth_nextlevel(full_sm), 5)
 
2369
 
 
2370
 
 
2371
 
 
2372
 
 
2373
        """ MSSM test
 
2374
        model_base = import_ufo.import_model('mssm')
 
2375
        param_path = os.path.join(_file_path,
 
2376
        '../input_files/param_card_mssm.dat')
 
2377
        model = decay_objects.DecayModel(model_base, force=True)
 
2378
        model.read_param_card(param_path)
 
2379
        model.find_all_channels(2)
 
2380
 
 
2381
        channel_2 = copy.deepcopy(h_ww_weve)
 
2382
        channel_2['vertices'][0]['legs'][0]['id'] = -1000024
 
2383
        channel_2['vertices'][0]['legs'][1]['id'] = 1000024
 
2384
        channel_2['vertices'][0]['legs'][2]['id'] = 23
 
2385
        channel_2['vertices'][0]['id'] = 128
 
2386
        channel_2['vertices'][1]['legs'][0]['id'] = 1000035
 
2387
        channel_2['vertices'][1]['legs'][1]['id'] = 23
 
2388
        channel_2['vertices'][1]['legs'][2]['id'] = 1000025
 
2389
        channel_2['vertices'][1]['id'] = 516
 
2390
        channel_2['onshell'] = 0
 
2391
        #print channel_2.get_onshell(model)
 
2392
        channel_2.get_final_legs()
 
2393
        #print channel_2.nice_string()
 
2394
        #print channel_2.get_apx_matrixelement_sq(model)
 
2395
        #print channel_2.get_apx_psarea(model)
 
2396
        #print channel_2.get_apx_decaywidth(model)
 
2397
        #print channel_2.get_apx_decaywidth_nextlevel(model)
 
2398
        #print channel_2.nice_string()
 
2399
        """
 
2400
 
 
2401
    def test_colormultiplicity(self):
 
2402
        """ Test the color_multiplicity_def of the DecayModel object and
 
2403
            the get_color_multiplicity function of Channel object. """
 
2404
 
 
2405
        # Test for exception 
 
2406
        self.assertRaises(decay_objects.DecayModel.PhysicsObjectError,
 
2407
                          self.my_testmodel.color_multiplicity_def,
 
2408
                          'a')
 
2409
 
 
2410
        self.assertRaises(decay_objects.DecayModel.PhysicsObjectError,
 
2411
                          self.my_testmodel.color_multiplicity_def,
 
2412
                          [1, 'a'])
 
2413
 
 
2414
        # Test the color_multiplicity_def
 
2415
        self.assertEqual(self.my_testmodel.color_multiplicity_def([6,3]),
 
2416
                         [(3, 2), (8, 3./4)])
 
2417
        
 
2418
        # Test the get_color_multiplicity
 
2419
        # Two-body decay
 
2420
        self.assertEqual(self.h_tt_bbmmvv.get_color_multiplicity(\
 
2421
                8, [3,3], self.my_testmodel, True),
 
2422
                         0.5)
 
2423
 
 
2424
 
 
2425
    def no_test_apx_decaywidth_full_read_MG4_paramcard(self):
 
2426
        """ The test to show the estimation of decay width.
 
2427
            and also read the param_card of MG4. 
 
2428
            Also test the find_all_channels including:
 
2429
            1. unity of total branching ratio
 
2430
            2. the apx_decaywidth_nextlevel comes from the right level."""
 
2431
 
 
2432
        # Options for the test
 
2433
        model_name = 'mssm'
 
2434
        test_param_card = False
 
2435
        test_param_card_suffix = 'test1'
 
2436
        smart_find = False
 
2437
        prec = 1E-2
 
2438
        channel_number = 3
 
2439
 
 
2440
 
 
2441
        # Read model_name
 
2442
        model_base = import_ufo.import_model(model_name)
 
2443
        model = decay_objects.DecayModel(model_base, True)
 
2444
 
 
2445
        # Read MG5 param_card
 
2446
        if test_param_card:
 
2447
            MG5_param_path = os.path.join(_file_path,
 
2448
                                        '../input_files/param_card_'\
 
2449
                                            +model_name \
 
2450
                                            +'_'\
 
2451
                                            +test_param_card_suffix
 
2452
                                            +'.dat')
 
2453
        else:
 
2454
            MG5_param_path = os.path.join(_file_path,
 
2455
                                        '../input_files/param_card_'\
 
2456
                                            +model_name \
 
2457
                                            +'.dat')
 
2458
        model.read_param_card(MG5_param_path)
 
2459
 
 
2460
        # Find channels before read MG4 param_card (option for use smart find)
 
2461
        if smart_find:
 
2462
            model.find_all_channels_smart(prec)
 
2463
        else:
 
2464
            model.find_all_channels(channel_number)
 
2465
 
 
2466
        # Read MG4 param_card
 
2467
        if model_name == "mssm":
 
2468
            if test_param_card:
 
2469
                MG4_param_path = os.path.join(_file_path,
 
2470
                                              '../input_files/param_card_test1.dat')
 
2471
            else:
 
2472
                MG4_param_path = os.path.join(_file_path,
 
2473
                                              '../input_files/param_card_0.dat')
 
2474
 
 
2475
            model.read_MG4_param_card_decay(MG4_param_path)
 
2476
 
 
2477
        # Write the decay summary, decay table, and helas collection
 
2478
        if test_param_card:
 
2479
            model.write_summary_decay_table(model_name\
 
2480
                                                + '_decay_summary_'\
 
2481
                                                + test_param_card_suffix \
 
2482
                                                + '.dat')
 
2483
            model.write_decay_table(MG5_param_path, 'cmp', 
 
2484
                                    model_name\
 
2485
                                        + '_decaytable_'\
 
2486
                                        + test_param_card_suffix \
 
2487
                                        + '.dat')
 
2488
            model.write_helas_collection(model_name\
 
2489
                                             + '_helascollection_'\
 
2490
                                             + test_param_card_suffix \
 
2491
                                             + '.dat')
 
2492
 
 
2493
        else:
 
2494
            # file name 1: default name
 
2495
            model.write_summary_decay_table()
 
2496
            model.write_decay_table(MG5_param_path, 'cmp')
 
2497
            model.write_helas_collection()
 
2498
 
 
2499
        """ Actual Test """
 
2500
        # Test the sum of branching ratios is unity
 
2501
        part = model.get_particle(25)
 
2502
        total_br = sum([sum([amp['apx_br'] for amp in part.get_amplitudes(i)]) \
 
2503
                            for i in range(2, part.get_max_level()+1)])
 
2504
        self.assertAlmostEqual(total_br, 1.)
 
2505
 
 
2506
        # Test if the err is from the off-shell 3-body channels
 
2507
        if part.get_max_level() > 2:
 
2508
            err = sum([c['apx_decaywidth_nextlevel'] \
 
2509
                           for c in part.get_channels(part.get_max_level(), 
 
2510
                                                      False)])
 
2511
            self.assertAlmostEqual(err/part.get('apx_decaywidth'), part['apx_decaywidth_err'])
 
2512
 
 
2513
 
 
2514
 
 
2515
        """ End of test """
 
2516
 
 
2517
        """# Test if the channels are find wisely
 
2518
        for part in model['particles']:
 
2519
            self.assertTrue(prec > part.get('apx_decaywidth_err'))
 
2520
 
 
2521
        # Check if the max_level is expected.
 
2522
        self.assertEqual(model.get_particle(1000016).get_max_level(), 2)
 
2523
        self.assertEqual(model.get_particle(25).get_max_level(), 3)"""
 
2524
 
 
2525
        # Miscellaneous
 
2526
        # Test if the calculated ratio is float or None
 
2527
        """for part in model.get('particles'):
 
2528
            print part.get_pdg_code(), part.get('2body_massdiff')
 
2529
            #n_max = len(part.decay_amplitudes.keys())
 
2530
            for n in range(2,n_max+2):
 
2531
                for amp in part.get_amplitudes(n):
 
2532
                    self.assertTrue(isinstance(amp['exa_decaywidth'], bool) or \
 
2533
                                        isinstance(amp['exa_decaywidth'], float))
 
2534
        """
 
2535
        #particle = model.get_particle(1000035)
 
2536
        #particleb = model.get_particle(2000001)
 
2537
        #channels = particleb.get_channels(3, False)
 
2538
        #channels.sort(decay_objects.channelcmp_width)
 
2539
        #print [c.nice_string() for c in channels[:100]]
 
2540
        #a = particle.get_amplitudes(3)
 
2541
        #print a
 
2542
        #print a.nice_string()#, a['diagrams']
 
2543
        #print model.get_interaction(166)
 
2544
        #print decay_objects.GC_415, decay_objects.GC_681
 
2545
 
 
2546
        """
 
2547
        #print particle.get_amplitude([-11, 2000011])['diagrams'].nice_string()
 
2548
        print model.get_interaction(390)
 
2549
        print decay_objects.GC_780, decay_objects.GC_752
 
2550
        print model.get_interaction(388)
 
2551
        print decay_objects.GC_757, decay_objects.GC_785
 
2552
        
 
2553
        print particle.get_amplitude([2000001, -1]).decaytable_string()
 
2554
        print particle.get_amplitude([2000001, -1])['diagrams'][0].get_apx_psarea(model)
 
2555
        print particle.get_amplitude([2000001, -1])['apx_decaywidth']
 
2556
        print particle.get_amplitude([2000001, -1])['exa_decaywidth']
 
2557
 
 
2558
        #print particleb.get_amplitude([-1000024, 16, 22]).nice_string()
 
2559
 
 
2560
        #print particle.get_amplitude([-11, 2000011])['diagrams'][0].get_apx_matrixelement_sq(model)        
 
2561
        #print particle.get_amplitude([-11, 2000011])['exa_decaywidth']
 
2562
 
 
2563
        #for part in model.get('particles'):
 
2564
        #    print part['pdg_code'], part['decay_width']
 
2565
 
 
2566
        #particle.calculate_branch_ratio()
 
2567
        #print decay_objects.MT, decay_objects.MW
 
2568
        #print decay_objects.GC_857, decay_objects.GC_733, decay_objects.GC_437, decay_objects.GC_665
 
2569
        #print particle.estimate_width_error()
 
2570
        #print len(particle.get_channels(3, False))
 
2571
        #print particle.get_channels(3, False)[0].nice_string(),\
 
2572
        #    particle.get('apx_decaywidth')
 
2573
        """
 
2574
 
 
2575
#===============================================================================
 
2576
# Test_IdentifyHelasTag
 
2577
#===============================================================================
 
2578
class Test_IdentifyHelasTag(unittest.TestCase):
 
2579
    """ Test for IdentifyHelasTag, and all the related functions. """
 
2580
 
 
2581
    my_testmodel_base = import_ufo.import_model('sm')
 
2582
    my_channel = decay_objects.Channel()
 
2583
    h_tt_bbmmvv = decay_objects.Channel()
 
2584
 
 
2585
    def setUp(self):
 
2586
        """ Set up necessary objects for the test"""
 
2587
 
 
2588
        #Import a model from my_testmodel
 
2589
        self.my_testmodel = decay_objects.DecayModel(self.my_testmodel_base, True)
 
2590
        param_path = os.path.join(_file_path,'../input_files/param_card_sm.dat')
 
2591
        self.my_testmodel.read_param_card(param_path)
 
2592
 
 
2593
        # Simplify the model
 
2594
        particles = self.my_testmodel.get('particles')
 
2595
        interactions = self.my_testmodel.get('interactions')
 
2596
        inter_list = copy.copy(interactions)
 
2597
 
 
2598
        # Pids that will be removed
 
2599
        no_want_pid = [1, 2, 15, 16, 21]
 
2600
        for pid in no_want_pid:
 
2601
            particles.remove(self.my_testmodel.get_particle(pid))
 
2602
 
 
2603
        for inter in inter_list:
 
2604
            if any([p.get('pdg_code') in no_want_pid for p in \
 
2605
                        inter.get('particles')]):
 
2606
                interactions.remove(inter)
 
2607
 
 
2608
        # Set a new name
 
2609
        self.my_testmodel.set('name', 'my_smallsm')
 
2610
        self.my_testmodel.set('particles', particles)
 
2611
        self.my_testmodel.set('interactions', interactions)
 
2612
        
 
2613
        # Set up vertexlist for my_testmodel
 
2614
        self.my_testmodel.find_vertexlist()
 
2615
    
 
2616
 
 
2617
    def test_helas_comparison(self):
 
2618
        """Test the ability to identify Helas calls."""
 
2619
 
 
2620
        # Turn higgs decay into 4-body
 
2621
        decay_objects.MH = 80
 
2622
 
 
2623
        h = self.my_testmodel.get_particle(25)
 
2624
        t = self.my_testmodel.get_particle(6)
 
2625
 
 
2626
        h.find_channels(4, self.my_testmodel)
 
2627
        
 
2628
        #print h.get_amplitudes(4).nice_string()
 
2629
 
 
2630
        for c in h.get_channels(4, True):
 
2631
            pids = set([l['id'] for l in c.get_final_legs()])
 
2632
            # Z mediated hadronic decays
 
2633
            if pids == set([3, -3, 4, -4]) and \
 
2634
                    c['vertices'][0]['legs'][-1]['id'] == 23:
 
2635
                h_zz_sscc = c
 
2636
                h_zz_sscc_Tag = decay_objects.IdentifyHelasTag(h_zz_sscc, 
 
2637
                                                               self.my_testmodel)
 
2638
            if pids == set([3, -3, 3, -3]):
 
2639
                h_zz_ssss = c
 
2640
                h_zz_ssss_Tag = decay_objects.IdentifyHelasTag(h_zz_ssss, 
 
2641
                                                               self.my_testmodel)
 
2642
 
 
2643
            if pids == set([5, -5, 5, -5]):
 
2644
                h_zz_bbbb = c
 
2645
                h_zz_bbbb_Tag = decay_objects.IdentifyHelasTag(h_zz_bbbb, 
 
2646
                                                               self.my_testmodel)
 
2647
            # Z mediated leptonic decays
 
2648
            if pids == set([11, -11, 12, -12]) and \
 
2649
                    c['vertices'][0]['legs'][-1]['id'] == 23:
 
2650
                h_zz_eeveve = c
 
2651
                h_zz_eeveve_Tag = decay_objects.IdentifyHelasTag(h_zz_eeveve, 
 
2652
                                                                 self.my_testmodel)
 
2653
            if pids == set([11, -11, 13, -13]):
 
2654
                h_zz_eemumu = c
 
2655
                h_zz_eemumu_Tag = decay_objects.IdentifyHelasTag(h_zz_eemumu, 
 
2656
                                                                 self.my_testmodel)
 
2657
            if pids == set([13, -13, 13, -13]):
 
2658
                h_zz_mumumumu = c
 
2659
                h_zz_mumumumu_Tag = decay_objects.IdentifyHelasTag(h_zz_mumumumu, 
 
2660
                                                                 self.my_testmodel)
 
2661
 
 
2662
            # W mediated hadronic decays
 
2663
            if pids == set([3, -3, 4, -4]) and \
 
2664
                    abs(c['vertices'][0]['legs'][-1]['id']) == 24:
 
2665
                h_ww_sscc = c
 
2666
                h_ww_sscc_Tag = decay_objects.IdentifyHelasTag(h_ww_sscc, 
 
2667
                                                               self.my_testmodel)
 
2668
            # W mediated leptonic decays
 
2669
            if pids == set([11, -11, 12, -12]) and \
 
2670
                    abs(c['vertices'][0]['legs'][-1]['id']) == 24:
 
2671
                h_ww_eeveve = c
 
2672
                h_ww_eeveve_Tag = decay_objects.IdentifyHelasTag(h_ww_eeveve, 
 
2673
                                                               self.my_testmodel)
 
2674
 
 
2675
 
 
2676
        #print h_zz_eemumu_Tag, h_zz_eeveve_Tag
 
2677
        self.assertTrue(h_zz_ssss_Tag == h_zz_bbbb_Tag)
 
2678
        self.assertTrue(h_zz_eemumu_Tag == h_zz_mumumumu_Tag)
 
2679
        # Lorentz structure of z > e e~ != z > ve ve~
 
2680
        self.assertFalse(h_zz_eemumu_Tag == h_zz_eeveve_Tag)
 
2681
        # Lorentz structure of z > up-type up-type != z > down-type down-type
 
2682
        self.assertFalse(h_zz_ssss_Tag == h_zz_sscc_Tag)
 
2683
        self.assertFalse(h_zz_sscc_Tag == h_ww_sscc_Tag)
 
2684
        self.assertFalse(h_zz_sscc_Tag == h_zz_eemumu_Tag)
 
2685
 
 
2686
        # Test diagram_from_tag
 
2687
        new_diagram = h_zz_ssss_Tag.diagram_from_tag(self.my_testmodel)
 
2688
        #print new_diagram.nice_string()
 
2689
        new_tag = diagram_generation.DiagramTag(new_diagram)
 
2690
        old_tag = diagram_generation.DiagramTag(h_zz_ssss)
 
2691
        #print new_tag, old_tag
 
2692
        self.assertTrue(new_tag == old_tag)
 
2693
 
 
2694
    def test_helas_helpers(self):
 
2695
        """ Test related helpers in DecayModel, Channel, DecayParticles. """
 
2696
 
 
2697
        t = self.my_testmodel.get_particle(6)
 
2698
        t.find_channels(2, self.my_testmodel)
 
2699
        t_bw = t.get_channels(2, True)[0]
 
2700
        
 
2701
        # Test get in Channel        
 
2702
        self.assertFalse(t_bw['helastag'] and t_bw['std_diagram'])
 
2703
        t_bw.get('helastag', self.my_testmodel)
 
2704
        self.assertTrue(t_bw['helastag'])
 
2705
        t_bw['helastag'] = None
 
2706
 
 
2707
 
 
2708
        # Test get_helas_properties in Channel
 
2709
        self.assertFalse(t_bw['helastag'] and t_bw['std_diagram'])
 
2710
        tag, std_diagram = t_bw.get_helas_properties(self.my_testmodel)
 
2711
        self.assertTrue(t_bw['helastag'] and t_bw['std_diagram'])
 
2712
        self.assertTrue(isinstance(tag, decay_objects.IdentifyHelasTag))
 
2713
        self.assertTrue(isinstance(std_diagram, base_objects.Diagram))
 
2714
        t_bw['helastag'] = None
 
2715
        t_bw['std_diagram'] = None
 
2716
 
 
2717
 
 
2718
        # Test get/add helascalls in DecayModel
 
2719
        self.assertEqual(self.my_testmodel.get_helascalls(4), [])
 
2720
        #     Error when input diagram has no helastag.
 
2721
        self.assertRaises(AssertionError, 
 
2722
                          self.my_testmodel.add_helascalls, 2, t_bw)
 
2723
        tag, std_diagram = t_bw.get_helas_properties(self.my_testmodel)
 
2724
        self.my_testmodel.add_helascalls(2, t_bw)
 
2725
        self.assertTrue(len(self.my_testmodel.get_helascalls(2))==1)
 
2726
        self.assertTrue(self.my_testmodel.get_helascalls(2)[0]['helastag'])
 
2727
 
 
2728
        #     The helas calls should be a deepcopy of the argument.
 
2729
        std_diagram['vertices'] = []
 
2730
        self.assertTrue(self.my_testmodel.get_helascalls(2)[0]['vertices'])
 
2731
 
 
2732
 
 
2733
    
 
2734
 
 
2735
 
 
2736
    def test_collect_helascalls(self):
 
2737
        """ Test the collect_helascalls. """
 
2738
 
 
2739
        # Turn top decay into 3-body
 
2740
        decay_objects.MT = 60
 
2741
 
 
2742
        h = self.my_testmodel.get_particle(25)
 
2743
        t = self.my_testmodel.get_particle(6)
 
2744
 
 
2745
        t.find_channels(3, self.my_testmodel)
 
2746
        #print t.get_channels(3, True).nice_string()
 
2747
 
 
2748
        # Test exceptions
 
2749
        self.assertRaises(decay_objects.DecayModel.PhysicsObjectError,
 
2750
                          self.my_testmodel.collect_helascalls, t, 'non-int')
 
2751
        self.assertRaises(decay_objects.DecayModel.PhysicsObjectError,
 
2752
                          self.my_testmodel.collect_helascalls, 'non-part', 3)
 
2753
 
 
2754
 
 
2755
        self.my_testmodel.collect_helascalls(t, 3)
 
2756
        # helas calls contain leptonic and hadronic top decay
 
2757
        self.assertEqual(len(self.my_testmodel.get_helascalls(3)), 2)
 
2758
 
 
2759
        # Test the properties set during collect_helascalls
 
2760
        t_beve = t.get_amplitude([-11, 12, 5])['diagrams'][0]
 
2761
        t_bmuvm = t.get_amplitude([-13, 14, 5])['diagrams'][0]
 
2762
        std_diagram = self.my_testmodel.get_helascalls(3)[t_beve['helas_number']]
 
2763
        self.assertEqual(t_beve['helas_number'],
 
2764
                         t_bmuvm['helas_number'])
 
2765
        self.assertTrue(std_diagram['helastag'] == t_bmuvm['helastag'])
 
2766
        
 
2767
 
 
2768
#===============================================================================
 
2769
# Test_DecayAmplitude
 
2770
#===============================================================================
 
2771
class Test_DecayAmplitude(unittest.TestCase):
 
2772
    """ Test for the DecayAmplitude and DecayAmplitudeList object."""
 
2773
 
 
2774
    def setUp(self):
 
2775
        """ Set up necessary objects for the test"""
 
2776
        self.my_testmodel_base = import_ufo.import_model('sm')
 
2777
        #Import a model from my_testmodel
 
2778
        self.my_testmodel = decay_objects.DecayModel(self.my_testmodel_base, True)
 
2779
        param_path = os.path.join(_file_path,'../input_files/param_card_sm.dat')
 
2780
        self.my_testmodel.read_param_card(param_path)
 
2781
 
 
2782
        my_channel = decay_objects.Channel()
 
2783
        h_tt_bbmmvv = decay_objects.Channel()
 
2784
 
 
2785
        # Simplify the model
 
2786
        particles = self.my_testmodel.get('particles')
 
2787
        interactions = self.my_testmodel.get('interactions')
 
2788
        inter_list = copy.copy(interactions)
 
2789
        # Pids that will be removed
 
2790
        no_want_pid = [1, 2, 3, 4, 15, 16, 21]
 
2791
        for pid in no_want_pid:
 
2792
            particles.remove(self.my_testmodel.get_particle(pid))
 
2793
 
 
2794
        for inter in inter_list:
 
2795
            if any([p.get('pdg_code') in no_want_pid for p in \
 
2796
                        inter.get('particles')]):
 
2797
                interactions.remove(inter)
 
2798
 
 
2799
        # Set a new name
 
2800
        self.my_testmodel.set('name', 'my_smallsm')
 
2801
        self.my_testmodel.set('particles', particles)
 
2802
        self.my_testmodel.set('interactions', interactions)
 
2803
 
 
2804
        #Setup the vertexlist for my_testmodel and save this model (optional)
 
2805
        import_vertexlist.make_vertexlist(self.my_testmodel)
 
2806
        #save_model.save_model(os.path.join(MG5DIR, 'tests/input_files', 
 
2807
        #self.my_testmodel['name']), self.my_testmodel)
 
2808
    
 
2809
 
 
2810
    def test_init_setget(self):
 
2811
        """ Test the set and get function of DecayAmplitude. """
 
2812
 
 
2813
        # Setup higgs and lower its mass
 
2814
        higgs = self.my_testmodel.get_particle(25)
 
2815
        t = self.my_testmodel.get_particle(6)
 
2816
        decay_objects.MH = 50
 
2817
 
 
2818
        # Set channels
 
2819
        self.my_testmodel.find_all_channels(4)
 
2820
        #print higgs.get_channels(4, True).nice_string()
 
2821
        #print t.get_channels(2,True)[0]
 
2822
 
 
2823
        # Get channels: h > ww > e+ e- ve ve~
 
2824
        #               h > zz > e+ e- e+ e-
 
2825
        for c in higgs.get_channels(4, True):
 
2826
            final_ids = set([l['id'] for l in c.get_final_legs()])
 
2827
            if final_ids == set([11, -11, 12, -12]):
 
2828
                # w mediated, not z
 
2829
                if abs(c['vertices'][0]['legs'][-1]['id']) == 24:
 
2830
                    h_mmvv_1 = c
 
2831
            if final_ids == set([11, -11, 11, -11]):
 
2832
                # z mediated
 
2833
                if abs(c['vertices'][0]['legs'][-1]['id']) == 23:
 
2834
                    h_mmvv_2 = c
 
2835
 
 
2836
 
 
2837
        # Test the initialization
 
2838
        amplt_h_mmvv = decay_objects.DecayAmplitude(h_mmvv_1,
 
2839
                                                    self.my_testmodel)
 
2840
        amplt_t_bw = decay_objects.DecayAmplitude(t.get_channels(2,True)[0],
 
2841
                                                  self.my_testmodel)
 
2842
        #print amplt_t_bw.nice_string()
 
2843
 
 
2844
        # goal id list for legs in process
 
2845
        # The legs are in the order of numbers!
 
2846
        # Subtest: test decay_objects.legcmp_bynumber
 
2847
        l1 = base_objects.Leg({'number':1, 'id':10})
 
2848
        l2 = base_objects.Leg({'number':2, 'id': 8})
 
2849
        sorted_legs = sorted([l2, l1], decay_objects.legcmp_bynumber)
 
2850
        self.assertEqual([l1, l2], sorted_legs)
 
2851
 
 
2852
        # Note: initial id in process should be POSITIVE
 
2853
        goal_id_list_h = set([25, 11, -11, -12, 12])
 
2854
        goal_id_list_t = set([6, 5, 24])
 
2855
        self.assertEqual(set([l.get('id') \
 
2856
                              for l in amplt_h_mmvv.get('process').get('legs')]),
 
2857
                         goal_id_list_h)
 
2858
        self.assertEqual(set([l.get('id') \
 
2859
                              for l in amplt_t_bw.get('process').get('legs')]),
 
2860
                         goal_id_list_t)
 
2861
 
 
2862
        # Check the legs in process (id, number, and state)
 
2863
        final_numbers = [2,3,4,5]
 
2864
        for l in amplt_h_mmvv.get('process').get('legs'):
 
2865
            if l.get('id') != 25:
 
2866
                self.assertTrue(l.get('state'))
 
2867
                final_numbers.remove(l.get('number'))
 
2868
            else:
 
2869
                self.assertFalse(l.get('state'))
 
2870
                self.assertEqual(1, l.get('number'))
 
2871
        self.assertFalse(final_numbers)
 
2872
 
 
2873
 
 
2874
        # Test the set and get in Amplitude
 
2875
        goal_width = h_mmvv_1.get('apx_decaywidth')
 
2876
        self.assertEqual(amplt_h_mmvv.get('apx_decaywidth'), goal_width)
 
2877
        goal_width += h_mmvv_2.get('apx_decaywidth')
 
2878
        amplt_h_mmvv.get('diagrams').append(h_mmvv_2)
 
2879
 
 
2880
 
 
2881
        # Test if the reset works.
 
2882
        amplt_h_mmvv.reset_width_br()
 
2883
        self.assertEqual(amplt_h_mmvv['apx_decaywidth'], 0.)
 
2884
        # Test the get for decaywidth and branch ratio.
 
2885
        self.assertEqual(amplt_h_mmvv.get('apx_decaywidth'), goal_width)
 
2886
        self.assertEqual(amplt_h_mmvv.get('apx_br'),
 
2887
                         goal_width/higgs.get('apx_decaywidth'))
 
2888
        WH = higgs['apx_decaywidth']
 
2889
        higgs['apx_decaywidth'] = 0.
 
2890
        amplt_h_mmvv.reset_width_br()
 
2891
        # WARNING should show for getting br from zero-width particle.
 
2892
        amplt_h_mmvv.get('apx_br')
 
2893
 
 
2894
 
 
2895
        # Test for non-self-conjugate mother
 
2896
        self.assertEqual(amplt_t_bw.get('apx_br'), 1.)
 
2897
        
 
2898
 
 
2899
        # Test for exceptions in set get of Amplitude
 
2900
        wrong_prop_list = {'process': [1, 'a', base_objects.Diagram()],
 
2901
                           'diagrams': [1, 'a', base_objects.Diagram()]}
 
2902
        for key, proplist in wrong_prop_list.items():
 
2903
            for prop in proplist:
 
2904
                self.assertRaises(decay_objects.DecayAmplitude.PhysicsObjectError,
 
2905
                                  amplt_h_mmvv.filter,
 
2906
                                  key, prop)
 
2907
 
 
2908
 
 
2909
        # Test for set and get in DecayParticle
 
2910
        my_amplist = decay_objects.DecayAmplitudeList([amplt_h_mmvv])
 
2911
        higgs.set('decay_amplitudes', {4: my_amplist})
 
2912
        self.assertEqual(higgs.get('decay_amplitudes'), {4: my_amplist})
 
2913
 
 
2914
 
 
2915
        # Test for exceptions
 
2916
        valuelist = ['nondict', {'a': my_amplist}, {4: base_objects.Process()}]
 
2917
        for value in valuelist:
 
2918
            self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
2919
                              higgs.filter,
 
2920
                              'decay_amplitudes', value)            
 
2921
 
 
2922
 
 
2923
        # Test for set_amplitudes and get_amplitudes
 
2924
        higgs.set_amplitudes(4, decay_objects.DecayAmplitudeList())
 
2925
        self.assertEqual(higgs.get_amplitudes(4), 
 
2926
                         decay_objects.DecayAmplitudeList())
 
2927
 
 
2928
 
 
2929
        # Test the set from normal list of Amplitude
 
2930
        higgs.set_amplitudes(4, [amplt_h_mmvv])
 
2931
        self.assertEqual(higgs.get_amplitudes(4), my_amplist)
 
2932
        self.assertEqual(higgs.get_amplitudes(6), [])
 
2933
 
 
2934
 
 
2935
        # Test for exceptions
 
2936
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
2937
                          higgs.get_amplitudes, 'a')
 
2938
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
2939
                          higgs.set_amplitudes,
 
2940
                          'a', decay_objects.DecayAmplitudeList())
 
2941
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
2942
                          higgs.set_amplitudes,
 
2943
                          4, decay_objects.DecayAmplitude())
 
2944
 
 
2945
 
 
2946
        # Test for get_amplitude
 
2947
        #higgs.get_amplitudes(4)
 
2948
        self.assertEqual(higgs.get_amplitude([12, -11, -12, 11]),
 
2949
                         amplt_h_mmvv)
 
2950
        self.assertEqual(higgs.get_amplitude([-12, -12, 11, 12]),
 
2951
                         None)
 
2952
 
 
2953
 
 
2954
        # Test for exception
 
2955
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
2956
                          higgs.get_amplitude,
 
2957
                          'Non-list')
 
2958
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
2959
                          higgs.get_amplitude,
 
2960
                          ['a', 1.23])
 
2961
 
 
2962
 
 
2963
    def test_group_channels2amplitudes(self):
 
2964
        """ Test the group_channels_2_amplitudes function."""
 
2965
 
 
2966
        # Setup higgs and lower its mass
 
2967
        higgs = self.my_testmodel.get_particle(25)
 
2968
        decay_objects.MH = 50
 
2969
 
 
2970
        # Set channels and amplitude
 
2971
 
 
2972
        self.my_testmodel.find_all_channels(4)
 
2973
        #print higgs.get_channels(4, True).nice_string()
 
2974
 
 
2975
        # Get channels: h > ww > e+ e- ve ve~
 
2976
        #               h > zz > e+ e- ve ve~
 
2977
        #               h > zz > e+ e- e+ e-
 
2978
        for c in higgs.get_channels(4, True):
 
2979
            final_ids = set([l['id'] for l in c.get_final_legs()])
 
2980
            if final_ids == set([13, -13, 14, -14]):
 
2981
                # w mediated, not z
 
2982
                if abs(c['vertices'][0]['legs'][-1]['id']) == 24:
 
2983
                    h_mmvv_1 = c
 
2984
                # z mediated
 
2985
                elif abs(c['vertices'][0]['legs'][-1]['id']) == 23:
 
2986
                    h_mmvv_2 = c
 
2987
            if final_ids == set([11, -11, 11, -11]):
 
2988
                h_epairs = c
 
2989
 
 
2990
 
 
2991
        #print h_mmvv_1.nice_string(), h_mmvv_2.nice_string()
 
2992
        amplt_h_mmvv = decay_objects.DecayAmplitude(h_mmvv_1, self.my_testmodel)
 
2993
        amplt_h_mmvv.add_std_diagram(h_mmvv_2)
 
2994
        amplt_h_epairs = decay_objects.DecayAmplitude(h_epairs, 
 
2995
                                                      self.my_testmodel)
 
2996
 
 
2997
 
 
2998
        # Test for exceptions
 
2999
 
 
3000
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
3001
                          higgs.group_channels_2_amplitudes,
 
3002
                          'a', self.my_testmodel)
 
3003
        self.assertRaises(decay_objects.DecayParticle.PhysicsObjectError,
 
3004
                          higgs.group_channels_2_amplitudes,
 
3005
                          3, 'a')
 
3006
 
 
3007
 
 
3008
        # Test if the group works
 
3009
 
 
3010
        higgs.find_channels(4, self.my_testmodel)
 
3011
        higgs.group_channels_2_amplitudes(4, self.my_testmodel)
 
3012
        #print higgs.get_amplitudes(4).nice_string()
 
3013
 
 
3014
        # Group function will calculate the decaywidth but not the apx_br
 
3015
        # fix it manually to compare with automatic result
 
3016
        amplt_h_mmvv.get('apx_decaywidth')
 
3017
        amplt_h_epairs.get('apx_decaywidth')
 
3018
 
 
3019
        self.assertTrue(amplt_h_mmvv in higgs.get_amplitudes(4))
 
3020
        self.assertTrue(amplt_h_epairs in higgs.get_amplitudes(4))
 
3021
 
 
3022
        # Test that all the final_legs has the same number
 
3023
        for amp in higgs.get_amplitudes(4):
 
3024
            final_list = set([(l['number'], l['id']) for l in amp['process']['legs'][1:]])
 
3025
            for dia in amp['diagrams']:
 
3026
                self.assertEqual(set([(l['number'], l['id']) for l in dia.get_final_legs()]), final_list)
 
3027
 
 
3028
 
 
3029
    def test_decaytable_string(self):
 
3030
        """ Test the decaytable_string """
 
3031
 
 
3032
        # Setup higgs and lower its mass
 
3033
        higgs = self.my_testmodel.get_particle(25)
 
3034
        decay_objects.MH = 50
 
3035
 
 
3036
        # Set channels and amplitude
 
3037
        self.my_testmodel.find_all_channels(4)
 
3038
        amp_list = higgs.get_amplitudes(4)
 
3039
        # Test for exception
 
3040
        self.assertRaises(decay_objects.DecayAmplitude.PhysicsObjectError,
 
3041
                          amp_list.decaytable_string,'wrongformat')
 
3042
        # Test for type
 
3043
        self.assertTrue(isinstance(amp_list.decaytable_string('full'), str))
 
3044
        self.assertTrue(isinstance(amp_list.decaytable_string(), str))
 
3045
 
 
3046
        # Test for decaytable_string from DecayParticle
 
3047
        self.assertTrue(isinstance(higgs.decaytable_string(), str))
 
3048
 
 
3049
        self.my_testmodel.write_decay_table(os.path.join(_file_path,'../input_files/param_card_sm.dat'), 
 
3050
                                            'full', 'mysmallmodel')
 
3051
 
 
3052
        #print self.my_testmodel['parameters'], '\n',\
 
3053
        #    self.my_testmodel['functions']
 
3054
 
 
3055
    def test_get_amplitude_givenfinal(self):
 
3056
        """ Test the get_amplitude function of DecayAmplitudeList. """
 
3057
 
 
3058
        # Setup higgs and lower its mass
 
3059
        higgs = self.my_testmodel.get_particle(25)
 
3060
        decay_objects.MH = 50
 
3061
 
 
3062
        # Set channels and amplitude
 
3063
        higgs.find_channels(4, self.my_testmodel)
 
3064
        amp_list = higgs.get_amplitudes(4)
 
3065
 
 
3066
        input_pids = [l.get('id') for l in amp_list[0]['process']['legs']\
 
3067
                          if l.get('state')]
 
3068
        self.assertEqual(amp_list.get_amplitude(input_pids),
 
3069
                         amp_list[0])
 
3070
        self.assertFalse(amp_list.get_amplitude([-11, 11, 6, -6]))
 
3071
 
 
3072
 
 
3073
    def test_add_std_diagram(self):
 
3074
        """ Test the add_std_diagram of DecayAmplitude."""
 
3075
 
 
3076
        # Setup higgs and lower its mass
 
3077
        higgs = self.my_testmodel.get_particle(25)
 
3078
        decay_objects.MH = 50
 
3079
 
 
3080
        # Set channels and amplitude
 
3081
        self.my_testmodel.find_all_channels(4)
 
3082
        h_eeveve =  higgs.get_amplitude([-11, 11, -12, 12])
 
3083
        
 
3084
        # Construct new amplitude
 
3085
        # diagram 0: h > ww > e+ e- ve ve~
 
3086
        std_amp = decay_objects.DecayAmplitude(h_eeveve['diagrams'][0], 
 
3087
                                          self.my_testmodel)
 
3088
 
 
3089
        # Modify the number of legs in second diagram
 
3090
        # diagram 1: h > zz > e+ e- ve ve~
 
3091
        new_diagram = copy.deepcopy(h_eeveve['diagrams'][1])
 
3092
        
 
3093
        # Make sure which z goes to which vertex
 
3094
        z_1 = new_diagram['vertices'][-1]['legs'][0]
 
3095
        z_2 = new_diagram['vertices'][-1]['legs'][0]
 
3096
 
 
3097
        if z_1 == new_diagram['vertices'][0]['legs'][-1]:
 
3098
            first_to_first = True
 
3099
        else:
 
3100
            first_to_first = False
 
3101
 
 
3102
        # Change the numbers
 
3103
        for vert in new_diagram['vertices']:
 
3104
            for leg in vert['legs']:
 
3105
                # change number 2 -> 5, 5-> 4, 4->2
 
3106
                old_number = leg['number']
 
3107
                if old_number == 2:
 
3108
                    leg['number'] = 5
 
3109
                if old_number == 5:
 
3110
                    leg['number'] = 4
 
3111
                if old_number == 4:
 
3112
                    leg['number'] = 2
 
3113
 
 
3114
        new_diagram.initial_setups(self.my_testmodel, True)
 
3115
 
 
3116
        #print 'dia1:', h_eeveve['diagrams'][0].nice_string(),\
 
3117
        #    '\ndia2(old):', h_eeveve['diagrams'][1].nice_string()
 
3118
        #print 'dia2 new:', new_diagram.nice_string()
 
3119
 
 
3120
 
 
3121
        # Add diagram (with a different number assignment) to std_amp
 
3122
        std_amp.add_std_diagram(new_diagram)
 
3123
        #print std_amp['diagrams'][0].nice_string()
 
3124
        #print std_amp['diagrams'][1].nice_string()
 
3125
 
 
3126
        # Test final legs
 
3127
        list_a = sorted(std_amp['diagrams'][0].get_final_legs(), 
 
3128
                        key=lambda l: l['id'])
 
3129
        number_a = [(l['number'], l['id']) for l in list_a]
 
3130
        list_b = sorted(std_amp['diagrams'][1].get_final_legs(), 
 
3131
                        key=lambda l: l['id'])
 
3132
        number_b = [(l['number'], l['id']) for l in list_b]
 
3133
        self.assertEqual(number_a, number_b)
 
3134
 
 
3135
 
 
3136
 
 
3137
        # Test all legs,
 
3138
        # numbers must be consistent.
 
3139
        #print std_amp.nice_string()
 
3140
        for dia in std_amp['diagrams']:
 
3141
            for i, vert in enumerate(dia['vertices'][:-1]):
 
3142
                # Check final (mother) leg of each vertices is the minimum
 
3143
                # among its children
 
3144
                previous_number = [l['number'] for l in vert['legs'][:-1]]
 
3145
                self.assertEqual(vert['legs'][-1]['number'],
 
3146
                                min(previous_number))
 
3147
                
 
3148
                # Check intermediate legs
 
3149
                l=vert['legs'][-1]
 
3150
                for vert2 in dia['vertices'][i+1:]:
 
3151
                    for l2 in vert2['legs'][:-1]:
 
3152
                        if l['number'] == l2['number']:
 
3153
                            self.assertEqual(l['id'], l2['id'])
 
3154
 
 
3155
                # Check the particles are being placed before antiparticles
 
3156
                ids = [l['id'] for l in vert['legs'][:-1]]
 
3157
                goal_ids = sorted(ids, reverse=True)
 
3158
                self.assertEqual(goal_ids, ids)
 
3159
                    
 
3160
 
 
3161
        # Test if the two z bosons still connect to the correct vertex
 
3162
        if first_to_first:
 
3163
            # 1st vertex from 1st z
 
3164
            self.assertEqual(\
 
3165
                std_amp['diagrams'][1]['vertices'][0]['legs'][-1]['number'],
 
3166
                std_amp['diagrams'][1]['vertices'][-1]['legs'][0]['number'])
 
3167
            # 2nd vertex from 2nd z
 
3168
            self.assertEqual(\
 
3169
                std_amp['diagrams'][1]['vertices'][1]['legs'][-1]['number'],
 
3170
                std_amp['diagrams'][1]['vertices'][-1]['legs'][1]['number'])
 
3171
 
 
3172
        # and vice versa
 
3173
        else:
 
3174
            self.assertEqual(\
 
3175
                std_amp['diagrams'][1]['vertices'][0]['legs'][-1]['number'],
 
3176
                std_amp['diagrams'][1]['vertices'][-1]['legs'][1]['number'])
 
3177
            self.assertEqual(\
 
3178
                std_amp['diagrams'][1]['vertices'][1]['legs'][-1]['number'],
 
3179
                std_amp['diagrams'][1]['vertices'][-1]['legs'][0]['number'])
 
3180
            
 
3181
 
 
3182
#===============================================================================
 
3183
# Test_AbstractModel
 
3184
#===============================================================================
 
3185
class Test_AbstractModel(unittest.TestCase):
 
3186
    """ Test for the AbstractModel object."""
 
3187
 
 
3188
    def setUp(self):
 
3189
        """ Set up necessary objects for the test"""
 
3190
        self.my_testmodel_base = import_ufo.import_model('sm')
 
3191
        #Import a model from my_testmodel
 
3192
        self.my_testmodel = decay_objects.DecayModel(self.my_testmodel_base,
 
3193
                                                     force=True)
 
3194
        param_path = os.path.join(_file_path,'../input_files/param_card_sm.dat')
 
3195
        self.my_testmodel.read_param_card(param_path)
 
3196
 
 
3197
        my_channel = decay_objects.Channel()
 
3198
        h_tt_bbmmvv = decay_objects.Channel()
 
3199
 
 
3200
        # Simplify the model
 
3201
        particles = self.my_testmodel.get('particles')
 
3202
        interactions = self.my_testmodel.get('interactions')
 
3203
        inter_list = copy.copy(interactions)
 
3204
        # Pids that will be removed
 
3205
        no_want_pid = [1, 2, 3, 4, 21]
 
3206
        for pid in no_want_pid:
 
3207
            particles.remove(self.my_testmodel.get_particle(pid))
 
3208
 
 
3209
        for inter in inter_list:
 
3210
            if any([p.get('pdg_code') in no_want_pid for p in \
 
3211
                        inter.get('particles')]):
 
3212
                interactions.remove(inter)
 
3213
 
 
3214
        # Set a new name
 
3215
        self.my_testmodel.set('name', 'my_smallsm')
 
3216
        self.my_testmodel.set('particles', particles)
 
3217
        self.my_testmodel.set('interactions', interactions)
 
3218
 
 
3219
    def test_get_particles_type(self):
 
3220
        """Test the set_new_particle. """
 
3221
        
 
3222
        ab_model = self.my_testmodel['ab_model']
 
3223
 
 
3224
        # Set electron as Majorana fermion
 
3225
        majorana_e = self.my_testmodel.get_particle(11)
 
3226
        majorana_e['self_antipart'] = True
 
3227
        # get_particles_type can be run only after generate the abstract_model
 
3228
        self.my_testmodel.generate_abstract_model()
 
3229
 
 
3230
        # Test of particle input
 
3231
        self.assertEqual((2, 3, False),
 
3232
                         ab_model.get_particle_type(\
 
3233
                self.my_testmodel.get_particle(6)))
 
3234
        self.assertEqual(9902300,
 
3235
                         ab_model.get_particle_type(\
 
3236
                self.my_testmodel.get_particle(6), get_code=True))
 
3237
        # Test of pdg_code input
 
3238
        self.assertEqual((2, 3, False),
 
3239
                         ab_model.get_particle_type(6))
 
3240
        self.assertEqual(9902300,
 
3241
                         ab_model.get_particle_type(6, get_code=True))
 
3242
        # Test of get_particle_type_code
 
3243
        self.assertEqual(-9902300,
 
3244
                         ab_model.get_particle_type_code(self.my_testmodel.get_particle(-6)))
 
3245
        self.assertEqual(9912100,
 
3246
                         ab_model.get_particle_type_code(majorana_e))
 
3247
 
 
3248
 
 
3249
        input_list = [6,5,24,-6,11,22]
 
3250
        reorder_list = [5,6,24,11,-6,22]
 
3251
 
 
3252
        goal_list_noignore = \
 
3253
            [9902300, 9902301, 9903101, -9902300, 9912100, 9903100]
 
3254
        goal_list_default = \
 
3255
            [9902302, 9902301, 9903101, -9902300, 9912100, 9903100]
 
3256
        goal_list_nonzerostart = \
 
3257
            [9902302, 9902303, 9903101, -9902302, 9912101, 9903100]
 
3258
 
 
3259
        goal_serial_dict_noignore = {(2, 3, False): 2, 
 
3260
                                     (3, 1, True): 2, 
 
3261
                                     (2, 1, True): 1}
 
3262
        goal_serial_dict_default = {(2, 3, False): 3, 
 
3263
                                    (3, 1, True): 2, 
 
3264
                                    (2, 1, True): 1}
 
3265
        goal_serial_dict_nonzerostart = {(2, 3, False): 4, 
 
3266
                                         (3, 1, True): 2, 
 
3267
                                         (2, 1, True): 2}
 
3268
        #print ab_model.get_particlelist_type.serial_number_dict
 
3269
        self.assertEqual(ab_model.get_particlelist_type(input_list, False)[0],
 
3270
                         goal_list_noignore)
 
3271
        self.assertEqual(ab_model.get_particlelist_type(input_list, False)[1],
 
3272
                         goal_serial_dict_noignore)
 
3273
        self.assertEqual(ab_model.get_particlelist_type(input_list)[0],
 
3274
                         goal_list_default)
 
3275
        self.assertEqual(ab_model.get_particlelist_type(input_list)[1],
 
3276
                         goal_serial_dict_default)
 
3277
 
 
3278
        self.assertEqual(ab_model.get_particlelist_type(input_list, False,
 
3279
                                                        sn_dict={\
 
3280
                    (2, 3, False):2, (2, 1, True):1})[0],
 
3281
                         goal_list_nonzerostart)
 
3282
        self.assertEqual(ab_model.get_particlelist_type(input_list, False,
 
3283
                                                        sn_dict={\
 
3284
                    (2, 3, False):2, (2, 1, True):1})[1],
 
3285
                         goal_serial_dict_nonzerostart)
 
3286
 
 
3287
        # Test if the output is regardless of input order
 
3288
        self.assertEqual(sorted(ab_model.get_particlelist_type(input_list)[0]),
 
3289
                         sorted(ab_model.get_particlelist_type(reorder_list)[0]))
 
3290
 
 
3291
 
 
3292
    def test_add_ab_particle(self):
 
3293
        """Test the set_new_particle. """
 
3294
        
 
3295
        ab_model = self.my_testmodel['ab_model']
 
3296
        # Set electron to be Majorana
 
3297
        majorana_e = self.my_testmodel.get_particle(11)
 
3298
        majorana_e['self_antipart'] = True
 
3299
        # Add new particle
 
3300
        newpart = decay_objects.DecayParticle({'pdg_code': 999,
 
3301
                                               'mass':'ZERO',
 
3302
                                               'spin':4, 'color':6, 
 
3303
                                               'self_antipart': True})
 
3304
        self.my_testmodel['particles'].append(newpart)
 
3305
        self.my_testmodel.reset_dictionaries()
 
3306
        self.my_testmodel.generate_abstract_model()
 
3307
 
 
3308
        # Test for exceptions
 
3309
        self.assertRaises(decay_objects.AbstractModel.PhysicsObjectError,
 
3310
                          ab_model.add_ab_particle,
 
3311
                          'NoneParticle')
 
3312
 
 
3313
        # Test get_particle_type
 
3314
        self.assertEqual(ab_model.get_particle_type(22),
 
3315
                         (3,1, True))
 
3316
        self.assertEqual(ab_model.get_particle_type(11),
 
3317
                         (2,1, True))
 
3318
 
 
3319
 
 
3320
        # Test for setting existing particle
 
3321
        ab_model.add_ab_particle(-6)
 
3322
        # Check if new particle is in paticles, abstract_particles_dict,
 
3323
        # and particle_dict of Model
 
3324
        self.assertEqual(len(ab_model['abstract_particles_dict'][(2,3,False)]),
 
3325
                         3)
 
3326
        newpart = ab_model['abstract_particles_dict'][(2,3, False)][-1]
 
3327
        self.assertTrue(newpart in ab_model.get('particles'))
 
3328
        # Test if the get_particle can find the correct anti-particle
 
3329
        self.assertEqual(ab_model.get_particle(-newpart.get_pdg_code())['is_part'],
 
3330
                        False)
 
3331
 
 
3332
        # Check properties
 
3333
        self.assertEqual(newpart.get('name'), 'F3_02')
 
3334
        self.assertEqual(newpart.get('antiname'), 'F3_02~')
 
3335
        self.assertEqual(newpart.get('mass'), 'MNF3_02')
 
3336
        self.assertEqual(newpart.get('pdg_code'), 9902302)
 
3337
 
 
3338
        # Test for setting new particle
 
3339
        ab_model.add_ab_particle(999)
 
3340
 
 
3341
        # Check if new particle is in paticles and abstract_particles_dict
 
3342
        self.assertTrue((4,6, True) in \
 
3343
                            ab_model['abstract_particles_dict'].keys())
 
3344
        self.assertTrue(ab_model.get_particle(9914601))
 
3345
 
 
3346
        # Check properties
 
3347
        newpart = ab_model['abstract_particles_dict'][(4,6, True)][-1]
 
3348
        self.assertEqual(newpart.get('name'), 'P6_01')
 
3349
        self.assertEqual(newpart.get('antiname'), 'none')
 
3350
        self.assertEqual(newpart.get('mass'), 'MSP6_01')
 
3351
        self.assertEqual(newpart.get('pdg_code'), 9914601)
 
3352
 
 
3353
    def test_setup_particles(self):
 
3354
        """ Test the add_particles from the generate_abstract_model."""
 
3355
        
 
3356
        # The gen_abtmodel should automatically generate abstract particles
 
3357
 
 
3358
        sm_path = import_ufo.find_ufo_path('sm')
 
3359
        full_sm_base = import_ufo.import_full_model(sm_path)
 
3360
        self.my_testmodel_new = decay_objects.DecayModel(full_sm_base, True)
 
3361
        param_path = os.path.join(_file_path,
 
3362
                                  '../input_files/param_card_full_sm.dat')
 
3363
        self.my_testmodel_new.read_param_card(param_path)
 
3364
        
 
3365
 
 
3366
        ab_model = self.my_testmodel_new['ab_model']
 
3367
        ab_model.setup_particles(self.my_testmodel_new['particles'])
 
3368
        goal_abpart_keys = set([(1,1,True), (2,1,False), (2,3,False), 
 
3369
                                (3,1,True), (3,8,True)])
 
3370
        goal_abpart_prop = {(1,1,True):('S1_00', 'none', 
 
3371
                                         'MSS1_00', 9901100),
 
3372
                            (2,1,False):('F1_00', 'F1_00~', 
 
3373
                                         'MNF1_00', 9902100),
 
3374
                            (2,3,False):('F3_00', 'F3_00~', 
 
3375
                                         'MNF3_00', 9902300),
 
3376
                            (3,1,True):('V1_00', 'none',
 
3377
                                         'MSV1_00', 9903100),
 
3378
                            (3,8,True):('V8_00', 'none', 
 
3379
                                         'MSV8_00', 9903800)}
 
3380
 
 
3381
        # Check keys in abstract_particles_dict
 
3382
        self.assertEqual(set(ab_model.get('abstract_particles_dict').keys()),
 
3383
                         goal_abpart_keys)
 
3384
        for plist in ab_model.get('abstract_particles_dict').values():
 
3385
            self.assertEqual(len(plist), 1)
 
3386
            self.assertEqual(tuple([plist[0].get('name'),
 
3387
                                    plist[0].get('antiname'),
 
3388
                                    plist[0].get('mass'),
 
3389
                                    plist[0].get('pdg_code')]),
 
3390
                             goal_abpart_prop[\
 
3391
                    ab_model.get_particle_type(plist[0])])
 
3392
 
 
3393
        # Check for exceptions
 
3394
        self.assertRaises(decay_objects.AbstractModel.PhysicsObjectError,
 
3395
                          ab_model.setup_particles,
 
3396
                          'NoneParticleList')
 
3397
        self.assertRaises(decay_objects.AbstractModel.PhysicsObjectError,
 
3398
                          ab_model.setup_particle,
 
3399
                          'NoneParticle')
 
3400
 
 
3401
    def test_setup_interactions(self):
 
3402
        """ Test the sets_interactions and get_interaction_type. """
 
3403
 
 
3404
        normal_sm_base = import_ufo.import_model('sm-full')
 
3405
        normal_sm = decay_objects.DecayModel(normal_sm_base,
 
3406
                                             force=True)
 
3407
        param_path = os.path.join(_file_path,'../../models/sm/restrict_ckm.dat')
 
3408
        normal_sm.read_param_card(param_path)
 
3409
        #print normal_sm.get_interaction(59)
 
3410
        normal_sm.generate_abstract_model()
 
3411
        ab_model = normal_sm['ab_model']
 
3412
 
 
3413
        # Test if the interaction_type_dict/ abstract_interactions_dict
 
3414
        # has been constructed properly.
 
3415
        for inter in normal_sm.get('interactions'):
 
3416
            try:
 
3417
                inter_type = ab_model['interaction_type_dict'][inter['id']]
 
3418
            except KeyError:
 
3419
                continue
 
3420
 
 
3421
 
 
3422
            # Test the particlelist type
 
3423
            # Note: the particlelist type in inter_type has transformed into
 
3424
            # tuple
 
3425
            parttype, sndict =ab_model.get_particlelist_type([p.get_pdg_code() for p in inter['particles']])
 
3426
            self.assertEqual(set(inter_type[0]), set(parttype))
 
3427
 
 
3428
            # Test the lorentz type is the superset of real lorentz
 
3429
            self.assertTrue(set(inter_type[1]).issuperset(inter['lorentz']))
 
3430
 
 
3431
            # Test the color is the superset of real lorentz
 
3432
            self.assertTrue(set(inter_type[2]).issuperset([str(c) for c in \
 
3433
                                                               inter['color']]))
 
3434
 
 
3435
            # Test if the abstract_interactions_dict has been established
 
3436
            self.assertTrue(inter_type in \
 
3437
                                ab_model['abstract_interactions_dict'].keys())
 
3438
 
 
3439
            # Test if the coupling_dict has been established
 
3440
            for ab_key, real_coup in \
 
3441
                    ab_model['interaction_coupling_dict'][inter['id']].items():
 
3442
 
 
3443
                ab_inter = ab_model['abstract_interactions_dict'][inter_type][0]
 
3444
                lorentz = ab_inter['lorentz'][ab_key[1]]
 
3445
                color = ab_inter['color'][ab_key[0]]
 
3446
                try:
 
3447
                    real_key = [0,0]
 
3448
                    real_key[0] = inter['color'].index(color)
 
3449
                    real_key[1] = inter['lorentz'].index(lorentz)
 
3450
                    self.assertEqual(inter['couplings'][tuple(real_key)],
 
3451
                                     real_coup)
 
3452
                except (ValueError, KeyError):
 
3453
                    self.assertEqual(real_coup, 'ZERO')
 
3454
 
 
3455
            # Test the interaction_dict
 
3456
            ab_inter = ab_model['abstract_interactions_dict'][inter_type][0]
 
3457
            self.assertTrue(ab_model.get_interaction(ab_inter['id']))
 
3458
 
 
3459
            # Test anti interaction
 
3460
            _has_anti = False
 
3461
            if inter['id'] in normal_sm['conj_int_dict'].keys():
 
3462
                _has_anti = True
 
3463
            if _has_anti:
 
3464
                self.assertEqual(ab_model['interaction_type_dict'][normal_sm['conj_int_dict'][inter['id']]],
 
3465
                                 ab_model['interaction_type_dict'][-inter['id']])
 
3466
                self.assertEqual(ab_model['interaction_coupling_dict'][normal_sm['conj_int_dict'][inter['id']]], 
 
3467
                                 ab_model['interaction_coupling_dict'][-inter['id']])
 
3468
 
 
3469
        # Test for non-repeated interaction id
 
3470
        # Test if all interactions in ab_model is up-to-date
 
3471
        id_list = [i['id'] for i in ab_model['interactions']]
 
3472
        new_inter_list = [intlist[0] for key, intlist \
 
3473
                              in ab_model['abstract_interactions_dict'].items()]
 
3474
 
 
3475
        for ab_inter in ab_model['interactions']:
 
3476
            self.assertEqual(id_list.count(ab_inter['id']), 1)
 
3477
            self.assertTrue(ab_inter in new_inter_list)
 
3478
 
 
3479
 
 
3480
        # Test if the lorentz and color types have no intersection
 
3481
        # if the particles type are the same            
 
3482
        def lorentzcmp(x,y):
 
3483
            return cmp(x[1],y[1])
 
3484
        keylist = sorted(ab_model['abstract_interactions_dict'].keys(), lorentzcmp)
 
3485
        #print normal_sm.get_particle(1000021)['self_antipart']
 
3486
        #print "Interaction type (%d):" % len(keylist)
 
3487
        #for k in keylist:
 
3488
        #    print k
 
3489
 
 
3490
        def check_keys(keylist):
 
3491
            try:
 
3492
                key = keylist.pop()
 
3493
            except:
 
3494
                return True
 
3495
            for other in keylist:
 
3496
                if other[0] == key[0]:
 
3497
                    self.assertTrue(set(other[1]).isdisjoint(key[1]))
 
3498
                    self.assertTrue(set(other[2]).isdisjoint(key[2]))
 
3499
 
 
3500
            return check_keys(keylist)
 
3501
 
 
3502
        self.assertTrue(check_keys(keylist))
 
3503
 
 
3504
 
 
3505
        # Test for exceptions
 
3506
        # Set an interaction that will not be included in AbstractModel
 
3507
        # (The [22, 6, -6] only works for testmodel = standard model)
 
3508
        for inter in normal_sm['interactions']:
 
3509
            if set([21, 6, -6]) == \
 
3510
                    set([part.get_pdg_code() for part in inter['particles']]):
 
3511
                I_rad = inter
 
3512
            if [25, 25, 25] == \
 
3513
                    [part.get_pdg_code() for part in inter['particles']]:
 
3514
                I_hhh = inter
 
3515
 
 
3516
        self.assertRaises(decay_objects.AbstractModel.PhysicsObjectError,
 
3517
                          ab_model.get_interaction_type,
 
3518
                          I_rad.get('id'))
 
3519
        self.assertRaises(decay_objects.AbstractModel.PhysicsObjectError,
 
3520
                          ab_model.get_interaction_type,
 
3521
                          I_hhh.get('id'))
 
3522
 
 
3523
 
 
3524
    def test_add_ab_interaction(self):
 
3525
        """Test the add_ab_interaction. """
 
3526
        
 
3527
        ab_model = self.my_testmodel['ab_model']
 
3528
        self.my_testmodel.generate_abstract_model()
 
3529
 
 
3530
        # Test for exceptions
 
3531
        self.assertRaises(decay_objects.AbstractModel.PhysicsObjectError,
 
3532
                          ab_model.add_ab_interaction,
 
3533
                          'NoneParticle')
 
3534
 
 
3535
        # Test for setting existing particle
 
3536
        for inter in self.my_testmodel['interactions']:
 
3537
            if set([24, -6, 5]) == \
 
3538
                    set([part.get_pdg_code() for part in inter['particles']]):
 
3539
                I_wtb = inter
 
3540
 
 
3541
        #print I_wtb
 
3542
        ab_model.add_ab_interaction(I_wtb.get('id'))
 
3543
        inter_type = ab_model.get_interaction_type(I_wtb.get('id'))
 
3544
        # The real inter_type is strictly sorted tuple
 
3545
        inter_type = (inter_type[0], inter_type[1], inter_type[2])
 
3546
 
 
3547
        # Check if new interaction is in paticles and abstract_particles_dict
 
3548
        self.assertEqual(len(ab_model['abstract_interactions_dict'][inter_type]),  
 
3549
                         2)
 
3550
        new_inter = ab_model['abstract_interactions_dict'][inter_type][-1]
 
3551
        self.assertTrue(new_inter in ab_model.get('interactions'))
 
3552
 
 
3553
        # Check properties
 
3554
        type_sn = int(new_inter.get('id') /1000)
 
3555
        #print type_sn, inter_type[1], inter_type[2], new_inter['couplings']
 
3556
        self.assertEqual(new_inter.get('id') % 10, 1)
 
3557
        self.assertEqual(new_inter.get('color'), I_wtb['color'])
 
3558
        self.assertEqual(new_inter.get('lorentz'), ['FFV2', 'FFV3', 'FFV5'])
 
3559
        self.assertEqual(new_inter.get('couplings'), 
 
3560
                         {(0,0):'G%03d0001' %type_sn,
 
3561
                          (0,1):'G%03d0101' %type_sn,
 
3562
                          (0,2):'G%03d0201' %type_sn})
 
3563
 
 
3564
 
 
3565
    def test_get_interactions_type(self):
 
3566
        """Test the set_new_particle. """
 
3567
        
 
3568
        ab_model = self.my_testmodel['ab_model']
 
3569
        self.my_testmodel.generate_abstract_model()
 
3570
 
 
3571
        higgs = self.my_testmodel.get_particle(25)
 
3572
        tau = self.my_testmodel.get_particle(15)
 
3573
 
 
3574
        higgs.find_channels(3, self.my_testmodel)
 
3575
        tau.find_channels(3, self.my_testmodel)
 
3576
        #for c in higgs.get_channels(3, True):
 
3577
        #    print c.nice_string()
 
3578
        #for c in tau.get_channels(3, True):
 
3579
        #    print c.nice_string()
 
3580
 
 
3581
        h_zz_zbb = higgs.get_channels(3, True)[0]
 
3582
        h_zz_zee = higgs.get_channels(3, True)[1]
 
3583
        h_zz_zmm = higgs.get_channels(3, True)[4]
 
3584
        #print h_zz_zbb.nice_string(), h_zz_zee.nice_string()
 
3585
        tau_wvt_vteve = tau.get_channels(3, True)[0]
 
3586
        tau_wvt_vtmvm = tau.get_channels(3, True)[1]
 
3587
        #print tau_wvt_vteve.nice_string(), tau_wvt_vtmvm.nice_string()
 
3588
 
 
3589
        
 
3590
        # Test get_interactionlist_type
 
3591
        input_list_1 = [v.get('id') for v in h_zz_zbb['vertices']]
 
3592
        input_list_3 = [v.get('id') for v in tau_wvt_vteve['vertices']]
 
3593
        input_list_3.append(input_list_3[0])
 
3594
        FFV_type = ab_model['interaction_type_dict'][abs(input_list_3[0])]
 
3595
        FFV_type_sn = \
 
3596
            int(ab_model['abstract_interactions_dict'][FFV_type][0]['id']/1000)
 
3597
 
 
3598
        goal_list3_default = \
 
3599
            [FFV_type_sn*1000, FFV_type_sn*1000+1, FFV_type_sn*1000]
 
3600
        goal_list3_ignoredup = \
 
3601
            [FFV_type_sn*1000, FFV_type_sn*1000+1, FFV_type_sn*1000+2]
 
3602
        goal_list_nonzerostart = \
 
3603
            [FFV_type_sn*1000+3, FFV_type_sn*1000+4, FFV_type_sn*1000+3]
 
3604
 
 
3605
        goal_serial_dict3_default = {FFV_type: 2}
 
3606
        goal_serial_dict3_ignoredup = {FFV_type: 3}
 
3607
        goal_serial_dict_nonzerostart = {FFV_type: 5} 
 
3608
 
 
3609
        #print ab_model.get_particlelist_type.serial_number_dict
 
3610
        self.assertEqual(ab_model.get_interactionlist_type(input_list_3)[0],
 
3611
                         goal_list3_default)
 
3612
        self.assertEqual(ab_model.get_interactionlist_type(input_list_3, 
 
3613
                                                           True)[0],
 
3614
                         goal_list3_ignoredup)
 
3615
        self.assertEqual(ab_model.get_interactionlist_type(input_list_3)[1],
 
3616
                         goal_serial_dict3_default)
 
3617
        self.assertEqual(ab_model.get_interactionlist_type(input_list_3,
 
3618
                                                           True)[1],
 
3619
                         goal_serial_dict3_ignoredup)
 
3620
        self.assertEqual(ab_model.get_interactionlist_type(input_list_3,
 
3621
                                                           sn_dict=goal_serial_dict3_ignoredup)[0],
 
3622
                         goal_list_nonzerostart)
 
3623
        self.assertEqual(ab_model.get_interactionlist_type(input_list_3,
 
3624
                                                           sn_dict=goal_serial_dict3_ignoredup)[1],
 
3625
                         goal_serial_dict_nonzerostart)
 
3626
 
 
3627
 
 
3628
 
 
3629
    def test_help_generate_ab_amplitude(self):
 
3630
        """ Test helper functions for generate abstract amplitude,
 
3631
        including compare_diagrams, add_ab_diagrams, 
 
3632
        generate_variables_dicts, and set_final_legs_dict."""
 
3633
 
 
3634
        ab_model = self.my_testmodel['ab_model']
 
3635
        self.my_testmodel.generate_abstract_model()
 
3636
 
 
3637
        #----------------------
 
3638
        # Set amplitudes
 
3639
        #----------------------
 
3640
        # Setup higgs, lower the mass,  and find amplitudes
 
3641
        #decay_objects.MH = 50
 
3642
 
 
3643
        higgs = self.my_testmodel.get_particle(25)
 
3644
        tau = self.my_testmodel.get_particle(15)
 
3645
        zboson = self.my_testmodel.get_particle(23)
 
3646
 
 
3647
        higgs.find_channels(3, self.my_testmodel)
 
3648
        tau.find_channels(3, self.my_testmodel)
 
3649
        #for c in higgs.get_channels(3, True):
 
3650
        #    print c.nice_string()
 
3651
        #for c in tau.get_channels(3, True):
 
3652
        #    print c.nice_string()
 
3653
 
 
3654
        h_zz_zbb = higgs.get_channels(3, True)[0]
 
3655
        h_zz_zee = higgs.get_channels(3, True)[4]
 
3656
        h_ww_weve = higgs.get_channels(3, True)[2]
 
3657
        h_zz_zmm = higgs.get_channels(3, True)[5]
 
3658
        #print h_zz_zbb.nice_string(), h_zz_zee.nice_string(), h_ww_weve.nice_string()
 
3659
        tau_wvt_vteve = tau.get_channels(3, True)[0]
 
3660
        tau_wvt_vtmvm = tau.get_channels(3, True)[1]
 
3661
        #print tau_wvt_vteve.nice_string(), tau_wvt_vtmvm.nice_string()
 
3662
 
 
3663
 
 
3664
        #----------------------
 
3665
        # Test set_final_legs_dict
 
3666
        #----------------------
 
3667
        ab2realdict = decay_objects.Ab2RealDict()
 
3668
        ab2realdict['process'] = tau.get_amplitude([16, 11, -12])['process']
 
3669
        ab_process = copy.deepcopy(tau.get_amplitude([16, 11, -12])['process'])
 
3670
        ab_process['model'] = ab_model
 
3671
        ab2realdict['ab_process'] = ab_process
 
3672
        real_pids = [l['id'] for l in tau_wvt_vteve.get_final_legs()]
 
3673
        ab2realdict.set_final_legs_dict()
 
3674
        self.assertEqual({-9902101:-12, 9902102:11, 9902103:16},
 
3675
                         ab2realdict['final_legs_dict'])
 
3676
        ab2realdict.set_final_legs_dict(real_dia=tau_wvt_vteve)
 
3677
        self.assertEqual({-9902101:-12, 9902102:11, 9902103:16},
 
3678
                         ab2realdict['final_legs_dict'])
 
3679
        ab2realdict.set_final_legs_dict(ab_object=ab_model)
 
3680
        self.assertEqual({-9902101:-12, 9902102:11, 9902103:16},
 
3681
                         ab2realdict['final_legs_dict'])
 
3682
        self.assertRaises(decay_objects.Ab2RealDict.PhysicsObjectError,
 
3683
                          ab2realdict.set_final_legs_dict,
 
3684
                          'Non-model')
 
3685
        
 
3686
 
 
3687
        #----------------------
 
3688
        # Test add abstract diagram
 
3689
        #----------------------
 
3690
        ab_amp = decay_objects.DecayAmplitude()
 
3691
        # test i/o
 
3692
        self.assertRaises(decay_objects.AbstractModel.PhysicsObjectError,
 
3693
                          ab_model.add_ab_diagram,
 
3694
                          h_zz_zee, h_zz_zbb)
 
3695
        self.assertRaises(decay_objects.AbstractModel.PhysicsObjectError,
 
3696
                          ab_model.add_ab_diagram,
 
3697
                          ab_amp, ab_amp)
 
3698
 
 
3699
        # Set the initial dict, including the final and initial legs
 
3700
        ab_amp['part_sn_dict'] = {(1,1,True): 1, (2,3,False):2, (3,1,True):1}
 
3701
        ab_amp['ab2real_dicts'].append(decay_objects.Ab2RealDict())
 
3702
 
 
3703
        ab_model.add_ab_diagram(ab_amp, h_zz_zbb)
 
3704
        ab_dia = ab_amp['diagrams'][0]
 
3705
        #print ab_dia.nice_string(), h_zz_zbb.nice_string()
 
3706
        goal_pids = set([-9902300, 9902301, 9903101, 9903101, 9903100, 
 
3707
                      9901100])
 
3708
        result_pids = []
 
3709
        [result_pids.extend([l.get('id') for l in v.get('legs')]) for v in ab_dia['vertices']]
 
3710
        self.assertEqual(set(result_pids), goal_pids)
 
3711
        real_iid_1 = zboson['decay_vertexlist'][(2, True)][0]['id']
 
3712
        real_iid_2 = higgs['decay_vertexlist'][(2, False)][1]['id']
 
3713
        goal_interids = [ab_model['abstract_interactions_dict'][\
 
3714
                ab_model.get_interaction_type(real_iid_1)][0]['id'],
 
3715
                         ab_model['abstract_interactions_dict'][\
 
3716
                ab_model.get_interaction_type(real_iid_2)][0]['id']]
 
3717
 
 
3718
        self.assertEqual([v.get('id') for v in ab_dia['vertices']], 
 
3719
                          goal_interids)
 
3720
        self.assertEqual(ab_amp['diagrams'][-1]['abstract_type'],
 
3721
                         [goal_interids, [9903100], [9902301, -9902300, 9903100]])
 
3722
        # Add the second diagram
 
3723
        # Abstract type should remain the same
 
3724
        ab_model.add_ab_diagram(ab_amp, h_zz_zbb)
 
3725
        self.assertEqual(ab_amp['diagrams'][-1]['abstract_type'], 
 
3726
                         ab_amp['diagrams'][-2]['abstract_type'])
 
3727
 
 
3728
 
 
3729
        #----------------------
 
3730
        # Test compare_diagrams
 
3731
        #----------------------
 
3732
        # test i/o
 
3733
        self.assertRaises(decay_objects.AbstractModel.PhysicsObjectError,
 
3734
                          ab_model.compare_diagrams,
 
3735
                          ab_amp, h_zz_zbb)
 
3736
        self.assertRaises(decay_objects.AbstractModel.PhysicsObjectError,
 
3737
                          ab_model.compare_diagrams,
 
3738
                          ab_dia, h_zz_zbb, {})
 
3739
 
 
3740
        ab_amp['ab2real_dicts'][-1].set_final_legs_dict(ab_model, h_zz_zbb)
 
3741
        #print ab_dia.nice_string(), h_zz_zbb.nice_string()
 
3742
        self.assertTrue(ab_model.compare_diagrams(ab_dia, h_zz_zbb,
 
3743
                                                  ab_amp['ab2real_dicts'][-1]))
 
3744
        # Test if no Ab2RealDict
 
3745
        self.assertTrue(ab_model.compare_diagrams(ab_dia, h_zz_zbb))
 
3746
        self.assertEqual(ab_amp['ab2real_dicts'][-1]['final_legs_dict'],
 
3747
                         {-9902300:-5, 9902301:5, 9903100:23})
 
3748
        # Set a wrong interaction id
 
3749
        h_zz_zbb_wrong = copy.deepcopy(h_zz_zbb)
 
3750
        h_zz_zbb_wrong['vertices'][1]['id'] = real_iid_1
 
3751
        self.assertFalse(ab_model.compare_diagrams(ab_dia, h_zz_zbb_wrong))
 
3752
 
 
3753
        # h_zz_zmm != h_zz_zbb
 
3754
        # no final_legs_dict available
 
3755
        self.assertFalse(ab_model.compare_diagrams(ab_dia, h_zz_zmm,
 
3756
                                                   ab_amp['ab2real_dicts'][-1]))
 
3757
 
 
3758
        # Construct abstract diagram for h_zz_zmm
 
3759
        ab_amp_2 = decay_objects.DecayAmplitude()
 
3760
        ab_amp_2['part_sn_dict'] = {(1,1,True): 1, (2,1,False):2, (3,1,True):1}
 
3761
        ab_amp_2['ab2real_dicts'].append(decay_objects.Ab2RealDict())
 
3762
        ab_model.add_ab_diagram(ab_amp_2, h_zz_zmm)
 
3763
        ab_dia_2 = ab_amp_2['diagrams'][0]
 
3764
 
 
3765
        # h_zz_zmm == h_zz_zee
 
3766
        ab_amp_2['ab2real_dicts'][-1].set_final_legs_dict(ab_model, h_zz_zmm)
 
3767
        self.assertTrue(ab_model.compare_diagrams(ab_dia_2, h_zz_zmm,
 
3768
                                                  ab_amp_2['ab2real_dicts'][-1]))
 
3769
        self.assertTrue(ab_model.compare_diagrams(ab_dia_2, h_zz_zmm))
 
3770
 
 
3771
        # Same final state, but not consistent as final_legs_dict
 
3772
        ab_amp_2['ab2real_dicts'].append(decay_objects.Ab2RealDict())
 
3773
        ab_amp_2['ab2real_dicts'][-1]['final_legs_dict'] = {-9902100:11, 9902101:-11, 9903100:23}
 
3774
        self.assertFalse(ab_model.compare_diagrams(ab_dia_2, h_zz_zee,
 
3775
                                                   ab_amp_2['ab2real_dicts'][-1]))
 
3776
        ab_amp_2['ab2real_dicts'].append(decay_objects.Ab2RealDict())
 
3777
        
 
3778
        #self.assertTrue(ab_model.compare_diagrams(ab_dia_2, h_ww_weve,
 
3779
        #                                          ab_amp_2['ab2real_dicts'][-1]))
 
3780
        #self.assertFalse(ab_model.compare_diagrams(ab_dia_2, h_zz_zbb,
 
3781
        #                                           ab_amp_2['ab2real_dicts'][-1]))
 
3782
 
 
3783
 
 
3784
 
 
3785
    def test_help_generate_ab_amplitude_2(self):
 
3786
        """ Test helper functions for generate abstract amplitude,
 
3787
        including compare_diagrams, add_ab_diagrams, and ."""
 
3788
 
 
3789
        ab_model = self.my_testmodel['ab_model']
 
3790
        self.my_testmodel.generate_abstract_model()
 
3791
 
 
3792
        # Setup higgs, lower the mass,  and find amplitudes
 
3793
        decay_objects.MH = 50
 
3794
 
 
3795
        higgs = self.my_testmodel.get_particle(25)
 
3796
        tau = self.my_testmodel.get_particle(15)
 
3797
        zboson = self.my_testmodel.get_particle(23)
 
3798
        wboson = self.my_testmodel.get_particle(24)
 
3799
 
 
3800
        higgs.find_channels(4, self.my_testmodel)
 
3801
        wboson.find_channels(2, self.my_testmodel)
 
3802
        h_zz_bbbb = None
 
3803
        # ta = tau
 
3804
        h_zz_tatabb = None
 
3805
        h_zz_tatatata = None
 
3806
        h_zz_tataee = None
 
3807
        h_zz_tatamm = None
 
3808
        for i,c in enumerate(higgs.get_channels(4, True)):
 
3809
            tag = [l['id'] for l in c.get_final_legs()]
 
3810
            if tag == [5, -5, 5, -5] and not h_zz_bbbb:
 
3811
                h_zz_bbbb = c
 
3812
            elif tag == [15, -15, 5, -5] and not h_zz_tatabb:
 
3813
                h_zz_tatabb = c
 
3814
            elif tag == [15, -15, 15, -15] and not h_zz_tatatata:
 
3815
                h_zz_tatatata = c
 
3816
            elif tag == [15, -15, 13, -13] and not h_zz_tatamm:
 
3817
                 h_zz_tatamm = c
 
3818
            elif tag == [15, -15, 11, -11]and not h_zz_tataee:
 
3819
                 h_zz_tataee = c         
 
3820
                        
 
3821
        for i,c in enumerate(wboson.get_channels(2, True)):
 
3822
            tag = [l['id'] for l in c.get_final_legs()]
 
3823
            if tag == [16, -15]:
 
3824
                w_lvl = c
 
3825
                break
 
3826
 
 
3827
        #print h_zz_bbbb.nice_string(), h_zz_tatabb.nice_string(), \
 
3828
        #    h_zz_tatatata.nice_string(), h_zz_tataee.nice_string()
 
3829
 
 
3830
        # Amplitude
 
3831
        h_zz_eevv = higgs.get_amplitude([-11, 11, 12, -12])
 
3832
        #print h_zz_eevv.nice_string()
 
3833
 
 
3834
        #----------------------
 
3835
        # Test add abstract diagram
 
3836
        #----------------------
 
3837
        ab_amp = decay_objects.DecayAmplitude()
 
3838
        # Set the initial dict, including the final and initial legs
 
3839
        ab_amp['part_sn_dict'] = {(1,1,True): 1, 
 
3840
                                  (2,3,False):4}
 
3841
        ab_amp['ab2real_dicts'].append(decay_objects.Ab2RealDict())
 
3842
        ab_model.add_ab_diagram(ab_amp, h_zz_bbbb)
 
3843
        ab_dia = ab_amp['diagrams'][0]
 
3844
        #print ab_dia.nice_string(), h_zz_bbbb.nice_string()
 
3845
        goal_pids = set([-9902300, 9902302, 9903100, 
 
3846
                      -9902301, 9902303, 9903101,
 
3847
                      9903101, 9903100, 9901100])
 
3848
 
 
3849
        result_pids = []
 
3850
        [result_pids.extend([l.get('id') for l in v.get('legs')]) for v in ab_dia['vertices']]
 
3851
        self.assertEqual(set(result_pids), goal_pids)
 
3852
 
 
3853
        real_iid_1 = zboson['decay_vertexlist'][(2, True)][0]['id']
 
3854
        real_iid_2 = higgs['decay_vertexlist'][(2, False)][1]['id']
 
3855
        ab_iid_1 = ab_model['abstract_interactions_dict'][\
 
3856
                ab_model.get_interaction_type(real_iid_1)][0]['id']
 
3857
        ab_iid_2 = ab_model['abstract_interactions_dict'][\
 
3858
                ab_model.get_interaction_type(real_iid_2)][0]['id']
 
3859
        goal_interids = [ab_iid_1, ab_iid_1, ab_iid_2]
 
3860
        self.assertEqual([v.get('id') for v in ab_dia['vertices']], 
 
3861
                          goal_interids)
 
3862
 
 
3863
        self.assertEqual(ab_amp['diagrams'][-1]['abstract_type'],
 
3864
                         [goal_interids, [9903100, 9903101], 
 
3865
                          [9902302, -9902300, 9902303, -9902301]])
 
3866
 
 
3867
        #----------------------
 
3868
        # Test compare diagrams, set_final_legs_dict, again
 
3869
        #----------------------
 
3870
        # Test set final_legs_dict
 
3871
        ab_amp['ab2real_dicts'][-1].set_final_legs_dict(ab_model, h_zz_bbbb)
 
3872
        self.assertTrue(ab_model.compare_diagrams(ab_dia, h_zz_bbbb,
 
3873
                                                  ab_amp['ab2real_dicts'][-1]))
 
3874
        #print ab_amp['ab2real_dicts'][-1]['final_legs_dict']
 
3875
        self.assertEqual(ab_amp['ab2real_dicts'][-1]['final_legs_dict'],
 
3876
                         {-9902300:-5,-9902301:-5, 9902302:5, 9902303:5})
 
3877
 
 
3878
 
 
3879
        # Test if the diagrams in one amplitude will share the same
 
3880
        # final particle correspondence
 
3881
        ab_amp = decay_objects.DecayAmplitude()
 
3882
        ab_amp['part_sn_dict'] = {(1,1,True): 1, 
 
3883
                                  (2,1,False):4}
 
3884
        ab_amp['ab2real_dicts'].append(decay_objects.Ab2RealDict())
 
3885
        ab_model.add_ab_diagram(ab_amp, h_zz_eevv['diagrams'][0])
 
3886
        ab_dia_1 = ab_amp['diagrams'][0]
 
3887
        ab_amp['ab2real_dicts'][-1].set_final_legs_dict(ab_model,
 
3888
                                                        h_zz_eevv['diagrams'][0])
 
3889
        ab_model.add_ab_diagram(ab_amp, h_zz_eevv['diagrams'][1])
 
3890
        ab_dia_2 = ab_amp['diagrams'][-1]
 
3891
        #print ab_amp.nice_string(), h_zz_eevv.nice_string()
 
3892
        # ab_dia_1: -e, ve, e, -ve; ab_dia_2: -ve, ve, -e, e
 
3893
        self.assertEqual(set([l['id'] for l in ab_dia_2.get_final_legs()]),
 
3894
                         set([9902103, -9902100, 9902102, -9902101]))
 
3895
 
 
3896
        #----------------------
 
3897
        # Test compare diagrams, set_final_legs_dict, again
 
3898
        #----------------------
 
3899
        # h_zz_zmm != h_zz_zbb
 
3900
        # no final_legs_dict available
 
3901
        self.assertFalse(ab_model.compare_diagrams(ab_dia, h_zz_tataee,
 
3902
                                                   ab_amp['ab2real_dicts'][-1]))
 
3903
 
 
3904
        # Construct abstract diagram for h_zz_zmm
 
3905
        ab_amp_2 = decay_objects.DecayAmplitude()
 
3906
        ab_amp_2['part_sn_dict'] = {(1,1,True): 1, 
 
3907
                                    (2,1,False):4}
 
3908
        ab_amp_2['ab2real_dicts'].append(decay_objects.Ab2RealDict())
 
3909
        ab_model.add_ab_diagram(ab_amp_2, h_zz_tataee)
 
3910
        ab_dia_2 = ab_amp_2['diagrams'][0]
 
3911
        # h_zz_tau tau e e == h_zz_tau tau tau tau
 
3912
        ab_amp_2['ab2real_dicts'][-1].set_final_legs_dict(ab_model,
 
3913
                                                          h_zz_tataee)
 
3914
        self.assertTrue(ab_model.compare_diagrams(\
 
3915
                ab_dia_2, 
 
3916
                h_zz_tataee,
 
3917
                ab_amp_2['ab2real_dicts'][-1]))
 
3918
        # For other diagram, add new Ab2RealDict
 
3919
        ab_amp_2['ab2real_dicts'].append(decay_objects.Ab2RealDict())
 
3920
        # ZZ > tau tau tau tau != ZZ > tau tau e e
 
3921
        # because the former use the same interaction but not the latter
 
3922
        self.assertFalse(ab_model.compare_diagrams(\
 
3923
                ab_dia_2, 
 
3924
                h_zz_tatatata,
 
3925
                ab_amp_2['ab2real_dicts'][-1]))
 
3926
 
 
3927
        # tau tau e e == tau tau mu mu
 
3928
        ab_amp_2['ab2real_dicts'].append(decay_objects.Ab2RealDict())
 
3929
        self.assertTrue(ab_model.compare_diagrams(\
 
3930
                ab_dia_2, 
 
3931
                h_zz_tatamm,
 
3932
                ab_amp_2['ab2real_dicts'][-1]))
 
3933
        # Set final legs, the ab_dia_2 should be the same type
 
3934
        ab_amp_2['ab2real_dicts'][-1].set_final_legs_dict(ab_model, 
 
3935
                                                          h_zz_tatamm)
 
3936
        self.assertTrue(ab_model.compare_diagrams(\
 
3937
                ab_dia_2, 
 
3938
                h_zz_tatamm,
 
3939
                ab_amp_2['ab2real_dicts'][-1]))
 
3940
 
 
3941
        #----------------------
 
3942
        # Test generate_variables_dicts
 
3943
        #----------------------       
 
3944
        ab_amp['process']['legs'] = base_objects.LegList([\
 
3945
                base_objects.Leg({'id':9901100, 'state':False}),
 
3946
                base_objects.Leg({'id':-9902100, 'state':True}),
 
3947
                base_objects.Leg({'id':9902101, 'state':True}),
 
3948
                base_objects.Leg({'id':-9902102, 'state':True}),
 
3949
                base_objects.Leg({'id':9902103, 'state':True})])
 
3950
        ab_amp['process']['model'] = ab_model
 
3951
        ab_amp['ab2real_dicts'][-1]['dia_sn_dict'] = {0:0, 1:1}
 
3952
        ab_amp.generate_variables_dicts(h_zz_eevv)
 
3953
 
 
3954
        # Test for initial particle
 
3955
        #print ab_amp.nice_string(), h_zz_eevv.nice_string()
 
3956
        #print ab_amp['ab2real_dicts'][-1]['mass_dict']
 
3957
        #print ab_amp['ab2real_dicts'][-1]['coup_dict']
 
3958
        #print ab_model['interaction_coupling_dict'][54], self.my_testmodel.get_interaction(54), ab_model['abstract_interactions_dict'][ab_model['interaction_type_dict'][54]]
 
3959
 
 
3960
        self.assertEqual(ab_amp['ab2real_dicts'][-1]['mass_dict'],
 
3961
                         {'MSS1_00':'MH',
 
3962
                          # lepton mass
 
3963
                          'MNF1_00':'ZERO', 'MNF1_01':'ZERO',
 
3964
                          'MNF1_02':'ZERO', 'MNF1_03':'ZERO',
 
3965
                          # w, z boson
 
3966
                          'MSV1_00':'MW', 'MSV1_01':'MW',
 
3967
                          'MSV1_02':'MZ', 'MSV1_03':'MZ'})
 
3968
        
 
3969
        coupling = {}
 
3970
        for candidate in self.my_testmodel['interactions']:
 
3971
            coupling[tuple([p['pdg_code'] for p in candidate['particles']])] = \
 
3972
                                candidate['couplings'][(0,0)]
 
3973
            if [p['pdg_code'] for p in candidate['particles']] == [11,11,23]:
 
3974
                second_z = candidate['couplings'][(0,1)]
 
3975
      
 
3976
        self.assertEqual(ab_amp['ab2real_dicts'][-1]['coup_dict'],
 
3977
                         {# type: s > vv
 
3978
                          #         : h > w+ w-
 
3979
                          'G0010000': coupling[(24,24,25)],
 
3980
                          #         : h > z z
 
3981
                          'G0010001': coupling[(23,23,25)],
 
3982
                          # type: v > ff
 
3983
                          #         : w+ > e+ ve
 
3984
                          'G0060000': coupling[(12,11,24)],
 
3985
                          #         : w- > e- ve~
 
3986
                          'G0060001': coupling[(12,11,24)],
 
3987
                          #         : z > ve ve~
 
3988
                          'G0060002': coupling[(12,12,23)],
 
3989
                          #         : z > e+ e- (lorentz=0)
 
3990
                          'G0060003': coupling[(11,11,23)],
 
3991
                          #         : z > e+ e- (lorentz=1)
 
3992
                          'G0060103':second_z
 
3993
                          })
 
3994
 
 
3995
        #----------------------
 
3996
        # Test generate_ab_amplitude
 
3997
        #----------------------       
 
3998
        ab_model.generate_ab_amplitudes(higgs.get_amplitudes(4))
 
3999
        #print higgs.get_amplitudes(4).nice_string()
 
4000
        # amp of abstract higgs:
 
4001
        # l l~ l' l'~, q q~ l l~, q q~ q q~
 
4002
        amplist = ab_model.get_particle(9901100).get_amplitudes(4)
 
4003
        self.assertEqual(len(amplist), 3)
 
4004
        #print amplist.abstract_nice_string()
 
4005
 
 
4006
        wboson.find_channels(3, self.my_testmodel)
 
4007
        ab_model.generate_ab_amplitudes(wboson.get_amplitudes(2))
 
4008
        amplist2 = ab_model.get_particle(9903100).get_amplitudes(2)
 
4009
        # w+ > l vl~ (3 families)
 
4010
        self.assertEqual(len(amplist2), 1)
 
4011
        #print amplist2.abstract_nice_string()
 
4012
 
 
4013
        tau.find_channels(3, self.my_testmodel)
 
4014
        ab_model.generate_ab_amplitudes(tau.get_amplitudes(3))
 
4015
        amplist3 = ab_model.get_particle(9902100).get_amplitudes(3)
 
4016
        # tt > vt l vl~ (3 families)
 
4017
        self.assertEqual(len(amplist3), 1)
 
4018
        #print amplist3.abstract_nice_string()
 
4019
 
 
4020
 
 
4021
        # Test if the diagrams in each amp have consistent final leg numbers.
 
4022
        for testlist in [amplist, amplist2, amplist3]:
 
4023
            for amp in testlist:
 
4024
                number_to_id_dict = dict([(l['number'],l['id']) \
 
4025
                                              for l in amp.get('process')['legs']])
 
4026
                #print number_to_id_dict
 
4027
                for channel in amp['diagrams']:
 
4028
                    for l in channel.get_final_legs():
 
4029
                        self.assertEqual(l['id'],
 
4030
                                         number_to_id_dict[l['number']])
 
4031
                        
 
4032
                        # ini_leg has anti pid compared to process
 
4033
                        ini_leg = channel['vertices'][-1]['legs'][-1]
 
4034
                        anti_pid = ab_model.get_particle(ini_leg['id']).get_anti_pdg_code()
 
4035
                        self.assertEqual(anti_pid,
 
4036
                                         number_to_id_dict[ini_leg['number']])
 
4037
 
 
4038
 
 
4039
        # Test the ab2real_dicts for mother which is not self-conjugate.
 
4040
        # Note that w boson is viewed as self-conjugate in abstract model.
 
4041
        #print amplist2[0].abstract_nice_string()
 
4042
        self.assertEqual(amplist2[0]['ab2real_dicts'][-1]['mass_dict'],
 
4043
                         {'MSV1_00':'MW',
 
4044
                          # lepton mass
 
4045
                          'MNF1_00':'MTA', 'MNF1_01':'ZERO'})
 
4046
        self.assertEqual(amplist3[0]['ab2real_dicts'][-1]['mass_dict'],
 
4047
                         {'MNF1_00':'MTA',
 
4048
                          # wboson mass
 
4049
                          'MSV1_00':'MW',
 
4050
                          # lepton mass
 
4051
                          'MNF1_01':'ZERO', 'MNF1_02':'ZERO', 'MNF1_03':'ZERO'})
 
4052
 
 
4053
 
 
4054
    def test_generate_ab_amplitudes(self):
 
4055
        """ Test generate the abstract amplitudes, matrixelement. """
 
4056
 
 
4057
        model_type = 'sm'
 
4058
        if model_type == 'full_sm':
 
4059
            sm_path = import_ufo.find_ufo_path('sm')
 
4060
            full_sm_base = import_ufo.import_full_model(sm_path)
 
4061
            full_sm = decay_objects.DecayModel(full_sm_base, True)
 
4062
            param_path = os.path.join(_file_path,
 
4063
                                      '../input_files/param_card_full_sm.dat')
 
4064
            full_sm.read_param_card(param_path)
 
4065
        else:
 
4066
            
 
4067
            normal_sm_base = import_ufo.import_model(\
 
4068
                model_type)
 
4069
            normal_sm = decay_objects.DecayModel(normal_sm_base,
 
4070
                                                 force=True)
 
4071
            param_path = os.path.join(_file_path,
 
4072
                                      '../input_files/param_card_%s.dat'%model_type)
 
4073
 
 
4074
        normal_sm.read_param_card(param_path)
 
4075
        #print normal_sm.get_interaction(59)        
 
4076
        #normal_sm.generate_abstract_model()
 
4077
        normal_sm.find_all_channels(3, generate_abstract=True)
 
4078
        #normal_sm.find_vertexlist()
 
4079
        #normal_sm.generate_abstract_amplitudes(3)
 
4080
        ab_model = normal_sm['ab_model']
 
4081
        if model_type == 'sm' or model_type == 'full_sm':
 
4082
            #print ab_model.get_particle(9901100).get_amplitudes(3).abstract_nice_string()
 
4083
            # h > q q'~ (w,z),  h > l l'~ (w,z)
 
4084
            self.assertEqual(len(ab_model.get_particle(9901100).get_amplitudes(3)), 
 
4085
                             2)
 
4086
 
 
4087
            # tau > q q'~ vt,  h > l l'~ vt
 
4088
            self.assertEqual(len(ab_model.get_particle(9902100).get_amplitudes(3)), 
 
4089
                             2)
 
4090
            # tau > q q'~ vt,  h > l l'~ vt
 
4091
            self.assertEqual(len(ab_model.get_particle(9902300).get_amplitudes(2)), 
 
4092
                             1)
 
4093
            # The CP is not invariant for w > q q'~
 
4094
            if model_type == 'full_sm':
 
4095
                self.assertTrue(ab_model.get_particle(9901100).get_amplitudes(3)[0]['ab2real_dicts'][0]['coup_dict'] != \
 
4096
                                    ab_model.get_particle(9901100).get_amplitudes(3)[0]['ab2real_dicts'][1]['coup_dict'])
 
4097
 
 
4098
 
 
4099
        ab_model.generate_ab_matrixelements_all()
 
4100
 
 
4101
if __name__ == '__main__':
 
4102
    unittest.unittest.main()