def setUp(self):

        # Set up model

        mypartlist = base_objects.ParticleList()
        myinterlist = base_objects.InteractionList()

        # A gluon
        mypartlist.append(
            base_objects.Particle({
                'name': 'g',
                'antiname': 'g',
                'spin': 3,
                'color': 8,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'g',
                'antitexname': 'g',
                'line': 'curly',
                'charge': 0.,
                'pdg_code': 21,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))

        g = mypartlist[len(mypartlist) - 1]

        # A quark U and its antiparticle
        mypartlist.append(
            base_objects.Particle({
                'name': 'u',
                'antiname': 'u~',
                'spin': 2,
                'color': 3,
                'mass': 'mu',
                'width': 'zero',
                'texname': 'u',
                'antitexname': '\bar u',
                'line': 'straight',
                'charge': 2. / 3.,
                'pdg_code': 2,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        u = mypartlist[len(mypartlist) - 1]
        antiu = copy.copy(u)
        antiu.set('is_part', False)

        # A quark D and its antiparticle
        mypartlist.append(
            base_objects.Particle({
                'name': 'd',
                'antiname': 'd~',
                'spin': 2,
                'color': 3,
                'mass': 'mu',
                'width': 'zero',
                'texname': 'd',
                'antitexname': '\bar d',
                'line': 'straight',
                'charge': -1. / 3.,
                'pdg_code': 1,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        d = mypartlist[len(mypartlist) - 1]
        antid = copy.copy(d)
        antid.set('is_part', False)

        # A electron and positron
        mypartlist.append(
            base_objects.Particle({
                'name': 'e+',
                'antiname': 'e-',
                'spin': 2,
                'color': 1,
                'mass': 'me',
                'width': 'zero',
                'texname': 'e^+',
                'antitexname': 'e^-',
                'line': 'straight',
                'charge': -1.,
                'pdg_code': 11,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        eminus = mypartlist[len(mypartlist) - 1]
        eplus = copy.copy(eminus)
        eplus.set('is_part', False)

        # A photon
        mypartlist.append(
            base_objects.Particle({
                'name': 'a',
                'antiname': 'a',
                'spin': 3,
                'color': 1,
                'mass': 'zero',
                'width': 'zero',
                'texname': '\gamma',
                'antitexname': '\gamma',
                'line': 'wavy',
                'charge': 0.,
                'pdg_code': 22,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))
        a = mypartlist[len(mypartlist) - 1]

        # A T particle
        mypartlist.append(
            base_objects.Particle({
                'name': 'T1',
                'antiname': 'T1',
                'spin': 5,
                'color': 1,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'T',
                'antitexname': 'T',
                'line': 'double',
                'charge': 0.,
                'pdg_code': 8000002,
                'propagating': False,
                'is_part': True,
                'self_antipart': True
            }))
        T1 = mypartlist[len(mypartlist) - 1]

        # A U squark and its antiparticle
        mypartlist.append(
            base_objects.Particle({
                'name': 'su',
                'antiname': 'su~',
                'spin': 1,
                'color': 3,
                'mass': 'Musq2',
                'width': 'Wusq2',
                'texname': '\tilde u',
                'antitexname': '\bar {\tilde u}',
                'line': 'dashed',
                'charge': 2. / 3.,
                'pdg_code': 1000002,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        su = mypartlist[len(mypartlist) - 1]
        antisu = copy.copy(su)
        antisu.set('is_part', False)

        # A E slepton and its antiparticle
        mypartlist.append(
            base_objects.Particle({
                'name': 'sl2-',
                'antiname': 'sl2+',
                'spin': 1,
                'color': 1,
                'mass': 'Msl2',
                'width': 'Wsl2',
                'texname': '\tilde e^-',
                'antitexname': '\tilde e^+',
                'line': 'dashed',
                'charge': 1.,
                'pdg_code': 1000011,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        seminus = mypartlist[len(mypartlist) - 1]
        seplus = copy.copy(seminus)
        seplus.set('is_part', False)

        # A neutralino
        mypartlist.append(
            base_objects.Particle({
                'name': 'n1',
                'antiname': 'n1',
                'spin': 2,
                'color': 1,
                'mass': 'Mneu1',
                'width': 'Wneu1',
                'texname': '\chi_0^1',
                'antitexname': '\chi_0^1',
                'line': 'straight',
                'charge': 0.,
                'pdg_code': 1000022,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))
        n1 = mypartlist[len(mypartlist) - 1]

        # W+ and W-
        mypartlist.append(
            base_objects.Particle({
                'name': 'W+',
                'antiname': 'W-',
                'spin': 3,
                'color': 1,
                'mass': 'wmas',
                'width': 'wwid',
                'texname': 'W^+',
                'antitexname': 'W^-',
                'line': 'wavy',
                'charge': 1.,
                'pdg_code': 24,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        wplus = mypartlist[len(mypartlist) - 1]
        wminus = copy.copy(u)
        wminus.set('is_part', False)

        # Z
        mypartlist.append(
            base_objects.Particle({
                'name': 'Z',
                'antiname': 'Z',
                'spin': 3,
                'color': 1,
                'mass': 'zmas',
                'width': 'zwid',
                'texname': 'Z',
                'antitexname': 'Z',
                'line': 'wavy',
                'charge': 1.,
                'pdg_code': 23,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))
        z = mypartlist[len(mypartlist) - 1]

        # Gluon and photon couplings to quarks
        myinterlist.append(base_objects.Interaction({
                      'id': 3,
                      'particles': base_objects.ParticleList(\
                                            [u, \
                                             antiu, \
                                             g]),
                      'color': [],
                      'lorentz':[''],
                      'couplings':{(0, 0):'GG'},
                      'orders':{'QCD':1}}))

        myinterlist.append(base_objects.Interaction({
                      'id': 4,
                      'particles': base_objects.ParticleList(\
                                            [d, \
                                             antid, \
                                             a]),
                      'color': [],
                      'lorentz':[''],
                      'couplings':{(0, 0):'MGVX15'},
                      'orders':{'QED':1}}))

        myinterlist.append(base_objects.Interaction({
                      'id': 10,
                      'particles': base_objects.ParticleList(\
                                            [d, \
                                             antid, \
                                             g]),
                      'color': [],
                      'lorentz':[''],
                      'couplings':{(0, 0):'GG'},
                      'orders':{'QCD':1}}))

        myinterlist.append(base_objects.Interaction({
                      'id': 11,
                      'particles': base_objects.ParticleList(\
                                            [u, \
                                             antiu, \
                                             a]),
                      'color': [],
                      'lorentz':[''],
                      'couplings':{(0, 0):'MGVX15'},
                      'orders':{'QED':1}}))

        # Tgg coupling
        myinterlist.append(base_objects.Interaction({
                      'id': 5,
                      'particles': base_objects.ParticleList(\
                                            [g, \
                                             g, \
                                             T1]),
                      'color': [],
                      'lorentz':['A'],
                      'couplings':{(0, 0):'MGVX2'},
                      'orders':{'QCD':1}}))

        # ggg coupling
        myinterlist.append(base_objects.Interaction({
                      'id': 15,
                      'particles': base_objects.ParticleList(\
                                            [g, \
                                             g, \
                                             g]),
                      'color': [],
                      'lorentz':[''],
                      'couplings':{(0, 0):'MGVX1'},
                      'orders':{'QCD':1}}))

        # Coupling of e to gamma
        myinterlist.append(base_objects.Interaction({
                      'id': 7,
                      'particles': base_objects.ParticleList(\
                                            [eminus, \
                                             eplus, \
                                             a]),
                      'color': [],
                      'lorentz':[''],
                      'couplings':{(0, 0):'MGVX12'},
                      'orders':{'QED':1}}))

        # Gluon coupling to su
        myinterlist.append(base_objects.Interaction({
                      'id': 105,
                      'particles': base_objects.ParticleList(\
                                            [g, \
                                             su, \
                                             antisu]),
                      'color': [],
                      'lorentz':[''],
                      'couplings':{(0, 0):'MGVX74'},
                      'orders':{'QCD':1}}))

        # Coupling of n1 to u and su
        myinterlist.append(base_objects.Interaction({
                      'id': 101,
                      'particles': base_objects.ParticleList(\
                                            [n1, \
                                             u, \
                                             antisu]),
                      'color': [],
                      'lorentz':[''],
                      'couplings':{(0, 0):'MGVX570'},
                      'orders':{'QED':1}}))

        myinterlist.append(base_objects.Interaction({
                      'id': 102,
                      'particles': base_objects.ParticleList(\
                                            [antiu, \
                                             n1, \
                                             su]),
                      'color': [],
                      'lorentz':[''],
                      'couplings':{(0, 0):'MGVX575'},
                      'orders':{'QED':1}}))

        # Coupling of n1 to e and se
        myinterlist.append(base_objects.Interaction({
                      'id': 103,
                      'particles': base_objects.ParticleList(\
                                            [n1, \
                                             eminus, \
                                             seplus]),
                      'color': [],
                      'lorentz':[''],
                      'couplings':{(0, 0):'MGVX350'},
                      'orders':{'QED':1}}))

        myinterlist.append(base_objects.Interaction({
                      'id': 104,
                      'particles': base_objects.ParticleList(\
                                            [eplus, \
                                             n1, \
                                             seminus]),
                      'color': [],
                      'lorentz':[''],
                      'couplings':{(0, 0):'MGVX494'},
                      'orders':{'QED':1}}))

        # Coupling of n1 to z
        myinterlist.append(base_objects.Interaction({
                      'id': 106,
                      'particles': base_objects.ParticleList(\
                                            [n1, \
                                             n1,
                                             z]),
                      'color': [],
                      'lorentz':[''],
                      'couplings':{(0, 0):'GZN11'},
                      'orders':{'QED':1}}))

        # g-gamma-su-subar coupling
        myinterlist.append(base_objects.Interaction({
                      'id': 100,
                      'particles': base_objects.ParticleList(\
                                            [a,
                                             g,
                                             su,
                                             antisu]),
                      'color': [],
                      'lorentz':[''],
                      'couplings':{(0, 0):'MGVX89'},
                      'orders':{'QED':1, 'QCD':1}}))

        # w+w-w+w- coupling
        myinterlist.append(base_objects.Interaction({
                      'id': 8,
                      'particles': base_objects.ParticleList(\
                                            [wplus,
                                             wminus,
                                             wplus,
                                             wminus]),
                      'color': [],
                      'lorentz':[''],
                      'couplings':{(0, 0):'MGVX6'},
                      'orders':{'QED':2}}))

        # w+w-zz coupling
        myinterlist.append(base_objects.Interaction({
                      'id': 9,
                      'particles': base_objects.ParticleList(\
                                            [wplus,
                                             wminus,
                                             z,
                                             z]),
                      'color': [],
                      'lorentz':[''],
                      'couplings':{(0, 0):'MGVX8'},
                      'orders':{'QED':2}}))

        self.mybasemodel.set('particles', mypartlist)
        self.mybasemodel.set('interactions', myinterlist)
        self.mybasemodel.set('name', 'sm')
        self.mymodel.set('model', self.mybasemodel)
    def setUp(self):
        # Set up model

        mypartlist = base_objects.ParticleList()
        myinterlist = base_objects.InteractionList()

        # A photon
        mypartlist.append(
            base_objects.Particle({
                'name': 'a',
                'antiname': 'a',
                'spin': 3,
                'color': 1,
                'mass': 'zero',
                'width': 'zero',
                'texname': '\gamma',
                'antitexname': '\gamma',
                'line': 'wavy',
                'charge': 0.,
                'pdg_code': 22,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))
        a = mypartlist[len(mypartlist) - 1]

        # W+ and W-
        mypartlist.append(
            base_objects.Particle({
                'name': 'w+',
                'antiname': 'w-',
                'spin': 3,
                'color': 1,
                'mass': 'wmas',
                'width': 'wwid',
                'texname': 'W^+',
                'antitexname': 'W^-',
                'line': 'wavy',
                'charge': 1.,
                'pdg_code': 24,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        wplus = mypartlist[len(mypartlist) - 1]
        wminus = copy.copy(wplus)
        wminus.set('is_part', False)

        # Z
        mypartlist.append(
            base_objects.Particle({
                'name': 'z',
                'antiname': 'z',
                'spin': 3,
                'color': 1,
                'mass': 'zmas',
                'width': 'zwid',
                'texname': 'Z',
                'antitexname': 'Z',
                'line': 'wavy',
                'charge': 1.,
                'pdg_code': 23,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))
        z = mypartlist[len(mypartlist) - 1]

        # a-a-w+w- 4-vertex
        myinterlist.append(base_objects.Interaction({
                      'id': 1,
                      'particles': base_objects.ParticleList(\
                                            [a, \
                                             a,
                                             wminus,
                                             wplus]),
                      'color': [],
                      'lorentz':['VVVV1'],
                      'couplings':{(0, 0):'GC_51'},
                      'orders':{'QED':2}}))

        # w+w-z vertex
        myinterlist.append(base_objects.Interaction({
                      'id': 2,
                      'particles': base_objects.ParticleList(\
                                            [wminus,
                                             wplus,
                                             z]),
                      'color': [],
                      'lorentz':['VVV1'],
                      'couplings':{(0, 0):'GC_12'},
                      'orders':{'QED':1}}))

        self.mybasemodel.set('particles', mypartlist)
        self.mybasemodel.set('interactions', myinterlist)
        self.mybasemodel.set('name', 'sm')

        #import madgraph.interface.cmd_interface as cmd
        #CMD = cmd.MadGraphCmdShell()
        #CMD._curr_model = self.mybasemodel
        #CMD._curr_fortran_model = helas_call_writers.FortranUFOHelasCallWriter
        #CMD.do_generate('a w- > w- a z')
        #CMD.do_export('matrix_v4 /tmp/')

        leg1 = base_objects.Leg({'id': 22, 'state': False})
        leg2 = base_objects.Leg({'id': 24, 'state': False})
        leg3 = base_objects.Leg({'id': 22, 'state': True})
        leg4 = base_objects.Leg({'id': 24, 'state': True})
        leg5 = base_objects.Leg({'id': 23, 'state': True})

        legList1 = base_objects.LegList([leg1, leg2, leg3, leg4, leg5])

        myproc = base_objects.Process({
            'legs': legList1,
            'model': self.mybasemodel
        })

        myamplitude = diagram_generation.Amplitude({'process': myproc})

        self.mymatrixelement = helas_objects.HelasMatrixElement(myamplitude)
    def test_export_matrix_element_python_madevent_group(self):
        """Test the result of exporting a subprocess group matrix element"""

        # Setup a model

        mypartlist = base_objects.ParticleList()
        myinterlist = base_objects.InteractionList()

        # A gluon
        mypartlist.append(
            base_objects.Particle({
                'name': 'g',
                'antiname': 'g',
                'spin': 3,
                'color': 8,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'g',
                'antitexname': 'g',
                'line': 'curly',
                'charge': 0.,
                'pdg_code': 21,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))

        g = mypartlist[-1]

        # A quark U and its antiparticle
        mypartlist.append(
            base_objects.Particle({
                'name': 'u',
                'antiname': 'u~',
                'spin': 2,
                'color': 3,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'u',
                'antitexname': '\bar u',
                'line': 'straight',
                'charge': 2. / 3.,
                'pdg_code': 2,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        u = mypartlist[-1]
        antiu = copy.copy(u)
        antiu.set('is_part', False)

        # A quark D and its antiparticle
        mypartlist.append(
            base_objects.Particle({
                'name': 'd',
                'antiname': 'd~',
                'spin': 2,
                'color': 3,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'd',
                'antitexname': '\bar d',
                'line': 'straight',
                'charge': -1. / 3.,
                'pdg_code': 1,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        d = mypartlist[-1]
        antid = copy.copy(d)
        antid.set('is_part', False)

        # A photon
        mypartlist.append(
            base_objects.Particle({
                'name': 'a',
                'antiname': 'a',
                'spin': 3,
                'color': 1,
                'mass': 'zero',
                'width': 'zero',
                'texname': '\gamma',
                'antitexname': '\gamma',
                'line': 'wavy',
                'charge': 0.,
                'pdg_code': 22,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))

        a = mypartlist[-1]

        # A Z
        mypartlist.append(
            base_objects.Particle({
                'name': 'z',
                'antiname': 'z',
                'spin': 3,
                'color': 1,
                'mass': 'MZ',
                'width': 'WZ',
                'texname': 'Z',
                'antitexname': 'Z',
                'line': 'wavy',
                'charge': 0.,
                'pdg_code': 23,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))
        z = mypartlist[-1]

        # Gluon and photon couplings to quarks
        myinterlist.append(base_objects.Interaction({
                      'id': 1,
                      'particles': base_objects.ParticleList(\
                                            [antiu, \
                                             u, \
                                             g]),
                      'color': [],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQQ'},
                      'orders':{'QCD':1}}))

        myinterlist.append(base_objects.Interaction({
                      'id': 2,
                      'particles': base_objects.ParticleList(\
                                            [antiu, \
                                             u, \
                                             a]),
                      'color': [],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQED'},
                      'orders':{'QED':1}}))

        myinterlist.append(base_objects.Interaction({
                      'id': 3,
                      'particles': base_objects.ParticleList(\
                                            [antid, \
                                             d, \
                                             g]),
                      'color': [],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQQ'},
                      'orders':{'QCD':1}}))

        myinterlist.append(base_objects.Interaction({
                      'id': 4,
                      'particles': base_objects.ParticleList(\
                                            [antid, \
                                             d, \
                                             a]),
                      'color': [],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQED'},
                      'orders':{'QED':1}}))

        # 3 gluon vertiex
        myinterlist.append(base_objects.Interaction({
                      'id': 5,
                      'particles': base_objects.ParticleList(\
                                            [g] * 3),
                      'color': [],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'G'},
                      'orders':{'QCD':1}}))

        # Coupling of Z to quarks

        myinterlist.append(base_objects.Interaction({
                      'id': 6,
                      'particles': base_objects.ParticleList(\
                                            [antiu, \
                                             u, \
                                             z]),
                      'color': [],
                      'lorentz':['L1', 'L2'],
                      'couplings':{(0, 0):'GUZ1', (0, 1):'GUZ2'},
                      'orders':{'QED':1}}))

        myinterlist.append(base_objects.Interaction({
                      'id': 7,
                      'particles': base_objects.ParticleList(\
                                            [antid, \
                                             d, \
                                             z]),
                      'color': [],
                      'lorentz':['L1', 'L2'],
                      'couplings':{(0, 0):'GDZ1', (0, 0):'GDZ2'},
                      'orders':{'QED':1}}))

        mymodel = base_objects.Model()
        mymodel.set('particles', mypartlist)
        mymodel.set('interactions', myinterlist)

        procs = [[2, -2, 21, 21], [2, -2, 2, -2]]
        amplitudes = diagram_generation.AmplitudeList()

        for proc in procs:
            # Define the multiprocess
            my_leglist = base_objects.LegList([\
                base_objects.Leg({'id': id, 'state': True}) for id in proc])

            my_leglist[0].set('state', False)
            my_leglist[1].set('state', False)

            my_process = base_objects.Process({
                'legs': my_leglist,
                'model': mymodel
            })
            my_amplitude = diagram_generation.Amplitude(my_process)
            amplitudes.append(my_amplitude)

        # Calculate diagrams for all processes

        subprocess_group = group_subprocs.SubProcessGroup.\
                           group_amplitudes(amplitudes, "madevent")[0]

        # Test amp2 lines
        helas_writer = helas_call_writers.PythonUFOHelasCallWriter(mymodel)
        python_exporter = export_python.ProcessExporterPython(
            subprocess_group, helas_writer)

        amp2_lines = \
                 python_exporter.get_amp2_lines(subprocess_group.\
                                        get('matrix_elements')[0],
                                        subprocess_group.get('diagram_maps')[0])
        self.assertEqual(amp2_lines, [
            'self.amp2[0]+=abs(amp[0]*amp[0].conjugate())',
            'self.amp2[1]+=abs(amp[1]*amp[1].conjugate())',
            'self.amp2[2]+=abs(amp[2]*amp[2].conjugate())'
        ])
    def setUp(self):

        # Set up model
        mypartlist = base_objects.ParticleList()
        myinterlist = base_objects.InteractionList()

        # u and c quarkd and their antiparticles
        mypartlist.append(
            base_objects.Particle({
                'name': 'u',
                'antiname': 'u~',
                'spin': 2,
                'color': 3,
                'mass': 'ZERO',
                'width': 'ZERO',
                'texname': 'u',
                'antitexname': '\bar u',
                'line': 'straight',
                'charge': 2. / 3.,
                'pdg_code': 2,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        u = mypartlist[len(mypartlist) - 1]
        antiu = copy.copy(u)
        antiu.set('is_part', False)

        mypartlist.append(
            base_objects.Particle({
                'name': 'c',
                'antiname': 'c~',
                'spin': 2,
                'color': 3,
                'mass': 'MC',
                'width': 'ZERO',
                'texname': 'c',
                'antitexname': '\bar c',
                'line': 'straight',
                'charge': 2. / 3.,
                'pdg_code': 4,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        c = mypartlist[len(mypartlist) - 1]
        antic = copy.copy(c)
        antic.set('is_part', False)

        # A gluon
        mypartlist.append(
            base_objects.Particle({
                'name': 'g',
                'antiname': 'g',
                'spin': 3,
                'color': 8,
                'mass': 'ZERO',
                'width': 'ZERO',
                'texname': 'g',
                'antitexname': 'g',
                'line': 'curly',
                'charge': 0.,
                'pdg_code': 21,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))

        g = mypartlist[len(mypartlist) - 1]

        # A photon
        mypartlist.append(
            base_objects.Particle({
                'name': 'Z',
                'antiname': 'Z',
                'spin': 3,
                'color': 1,
                'mass': 'MZ',
                'width': 'WZ',
                'texname': 'Z',
                'antitexname': 'Z',
                'line': 'wavy',
                'charge': 0.,
                'pdg_code': 23,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))
        z = mypartlist[len(mypartlist) - 1]

        # Gluon couplings to quarks
        myinterlist.append(base_objects.Interaction({
                      'id': 1,
                      'particles': base_objects.ParticleList(\
                                            [antiu, \
                                             u, \
                                             g]),
                      'color': [color.ColorString([color.T(2, 1, 0)])],
                      'lorentz':['FFV1'],
                      'couplings':{(0, 0):'GC_10'},
                      'orders':{'QCD':1}}))

        # Gamma couplings to quarks
        myinterlist.append(base_objects.Interaction({
                      'id': 2,
                      'particles': base_objects.ParticleList(\
                                            [antiu, \
                                             u, \
                                             z]),
                      'color': [color.ColorString([color.T(1, 0)])],
                      'lorentz':['FFV2', 'FFV5'],
                      'couplings':{(0,0): 'GC_35', (0,1): 'GC_47'},
                      'orders':{'QED':1}}))

        self.mymodel.set('particles', mypartlist)
        self.mymodel.set('interactions', myinterlist)
        self.mymodel.set('name', 'sm')

        self.mypythonmodel = helas_call_writers.PythonUFOHelasCallWriter(
            self.mymodel)

        myleglist = base_objects.LegList()

        myleglist.append(base_objects.Leg({'id': 2, 'state': False}))
        myleglist.append(base_objects.Leg({'id': -2, 'state': False}))
        myleglist.append(base_objects.Leg({'id': 2, 'state': True}))
        myleglist.append(base_objects.Leg({'id': -2, 'state': True}))

        myproc = base_objects.Process({
            'legs': myleglist,
            'model': self.mymodel
        })

        myamplitude = diagram_generation.Amplitude({'process': myproc})

        self.mymatrixelement = helas_objects.HelasMultiProcess(myamplitude)

        myleglist = base_objects.LegList()

        myleglist.append(
            base_objects.Leg({
                'id': 4,
                'state': False,
                'number': 1
            }))
        myleglist.append(
            base_objects.Leg({
                'id': -4,
                'state': False,
                'number': 2
            }))
        myleglist.append(
            base_objects.Leg({
                'id': 4,
                'state': True,
                'number': 3
            }))
        myleglist.append(
            base_objects.Leg({
                'id': -4,
                'state': True,
                'number': 4
            }))

        myproc = base_objects.Process({
            'legs': myleglist,
            'model': self.mymodel
        })

        self.mymatrixelement.get('matrix_elements')[0].\
                                               get('processes').append(myproc)

        self.exporter = export_python.ProcessExporterPython(\
            self.mymatrixelement, self.mypythonmodel)
Example #5
0
class ColorSquareTest(unittest.TestCase):
    """Test class for the color_amp module"""

    mypartlist = base_objects.ParticleList()
    myinterlist = base_objects.InteractionList()
    mymodel = base_objects.Model()

    def setUp(self):
        # A gluon
        self.mypartlist.append(
            base_objects.Particle({
                'name': 'g',
                'antiname': 'g',
                'spin': 3,
                'color': 8,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'g',
                'antitexname': 'g',
                'line': 'curly',
                'charge': 0.,
                'pdg_code': 21,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))

        # A quark U and its antiparticle
        self.mypartlist.append(
            base_objects.Particle({
                'name': 'u',
                'antiname': 'u~',
                'spin': 2,
                'color': 3,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'u',
                'antitexname': '\bar u',
                'line': 'straight',
                'charge': 2. / 3.,
                'pdg_code': 2,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        antiu = copy.copy(self.mypartlist[1])
        antiu.set('is_part', False)

        # A quark D and its antiparticle
        self.mypartlist.append(
            base_objects.Particle({
                'name': 'd',
                'antiname': 'd~',
                'spin': 2,
                'color': 3,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'u',
                'antitexname': '\bar u',
                'line': 'straight',
                'charge': -1. / 3.,
                'pdg_code': 1,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        antid = copy.copy(self.mypartlist[2])
        antid.set('is_part', False)

        # 3 gluon vertiex
        self.myinterlist.append(base_objects.Interaction({
                      'id': 1,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[0]] * 3),
                      'color': [color.ColorString([color.f(0, 1, 2)])],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'G'},
                      'orders':{'QCD':1}}))

        # 4 gluon vertex
        self.myinterlist.append(base_objects.Interaction({
                      'id': 2,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[0]] * 4),
                      'color': [color.ColorString([color.f(-1, 0, 2),
                                                   color.f(-1, 1, 3)]),
                                color.ColorString([color.f(-1, 0, 3),
                                                   color.f(-1, 1, 2)]),
                                color.ColorString([color.f(-1, 0, 1),
                                                   color.f(-1, 2, 3)])],
                      'lorentz':['L(p1,p2,p3)', 'L(p2,p3,p1)', 'L3'],
                      'couplings':{(0, 0):'G^2',
                                   (1, 1):'G^2',
                                   (2, 2):'G^2'},
                      'orders':{'QCD':2}}))

        # Gluon couplings to up and down quarks
        self.myinterlist.append(base_objects.Interaction({
                      'id': 3,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[1], \
                                             antiu, \
                                             self.mypartlist[0]]),
                      'color': [color.ColorString([color.T(2, 0, 1)])],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQQ'},
                      'orders':{'QCD':1}}))

        self.myinterlist.append(base_objects.Interaction({
                      'id': 4,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[2], \
                                             antid, \
                                             self.mypartlist[0]]),
                      'color': [color.ColorString([color.T(2, 0, 1)])],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQQ'},
                      'orders':{'QCD':1}}))

        self.mymodel.set('particles', self.mypartlist)
        self.mymodel.set('interactions', self.myinterlist)

    def test_color_matrix_multi_gluons(self):
        """Test the color matrix building for gg > n*g with n up to 3"""

        goal = [
            fractions.Fraction(7, 3),
            fractions.Fraction(19, 6),
            fractions.Fraction(455, 108),
            fractions.Fraction(3641, 648)
        ]

        goal_line1 = [(fractions.Fraction(7, 3), fractions.Fraction(-2, 3)),
                      (fractions.Fraction(19, 6), fractions.Fraction(-1, 3),
                       fractions.Fraction(-1, 3), fractions.Fraction(-1, 3),
                       fractions.Fraction(-1, 3), fractions.Fraction(2, 3)),
                      (fractions.Fraction(455,
                                          108), fractions.Fraction(-29, 54),
                       fractions.Fraction(-29, 54), fractions.Fraction(7, 54),
                       fractions.Fraction(7, 54), fractions.Fraction(17, 27),
                       fractions.Fraction(-29, 54), fractions.Fraction(-1, 27),
                       fractions.Fraction(7, 54), fractions.Fraction(-29, 54),
                       fractions.Fraction(5, 108), fractions.Fraction(-1, 27),
                       fractions.Fraction(7, 54), fractions.Fraction(5, 108),
                       fractions.Fraction(17, 27), fractions.Fraction(-1, 27),
                       fractions.Fraction(7, 54), fractions.Fraction(17, 27),
                       fractions.Fraction(-29, 54), fractions.Fraction(-1, 27),
                       fractions.Fraction(-1, 27), fractions.Fraction(17, 27),
                       fractions.Fraction(17, 27), fractions.Fraction(-10,
                                                                      27))]

        for n in range(3):
            myleglist = base_objects.LegList()

            myleglist.append(base_objects.Leg({'id': 21, 'state': False}))
            myleglist.append(base_objects.Leg({'id': 21, 'state': False}))

            myleglist.extend([base_objects.Leg({
                'id': 21,
                'state': True
            })] * (n + 1))

            myprocess = base_objects.Process({
                'legs': myleglist,
                'model': self.mymodel
            })

            myamplitude = diagram_generation.Amplitude()

            myamplitude.set('process', myprocess)

            myamplitude.generate_diagrams()

            col_basis = color_amp.ColorBasis(myamplitude)

            col_matrix = color_amp.ColorMatrix(col_basis, Nc=3)
            # Check diagonal
            for i in range(len(col_basis.items())):
                self.assertEqual(col_matrix.col_matrix_fixed_Nc[(i, i)],
                                 (goal[n], 0))

            # Check first line
            for i in range(len(col_basis.items())):
                self.assertEqual(col_matrix.col_matrix_fixed_Nc[(0, i)],
                                 (goal_line1[n][i], 0))

    def test_color_matrix_multi_quarks(self):
        """Test the color matrix building for qq~ > n*(qq~) with n up to 2"""

        goal = [fractions.Fraction(9, 1), fractions.Fraction(27, 1)]

        goal_line1 = [(fractions.Fraction(9, 1), fractions.Fraction(3, 1)),
                      (fractions.Fraction(27, 1), fractions.Fraction(9, 1),
                       fractions.Fraction(9, 1), fractions.Fraction(3, 1),
                       fractions.Fraction(3, 1), fractions.Fraction(9, 1))]

        goal_den_list = [[1] * 2, [1] * 6]

        goal_first_line_num = [[9, 3], [27, 9, 9, 3, 3, 9]]

        for n in range(2):
            myleglist = base_objects.LegList()

            myleglist.append(base_objects.Leg({'id': -2, 'state': False}))
            myleglist.append(base_objects.Leg({'id': 2, 'state': False}))

            myleglist.extend([
                base_objects.Leg({
                    'id': -2,
                    'state': True
                }),
                base_objects.Leg({
                    'id': 2,
                    'state': True
                })
            ] * (n + 1))

            myprocess = base_objects.Process({
                'legs': myleglist,
                'model': self.mymodel
            })

            myamplitude = diagram_generation.Amplitude()

            myamplitude.set('process', myprocess)

            myamplitude.generate_diagrams()

            col_basis = color_amp.ColorBasis(myamplitude)

            col_matrix = color_amp.ColorMatrix(col_basis, Nc=3)
            # Check diagonal
            for i in range(len(col_basis.items())):
                self.assertEqual(col_matrix.col_matrix_fixed_Nc[(i, i)],
                                 (goal[n], 0))

            # Check first line
            for i in range(len(col_basis.items())):
                self.assertEqual(col_matrix.col_matrix_fixed_Nc[(0, i)],
                                 (goal_line1[n][i], 0))

            self.assertEqual(col_matrix.get_line_denominators(),
                             goal_den_list[n])
            self.assertEqual(
                col_matrix.get_line_numerators(
                    0,
                    col_matrix.get_line_denominators()[0]),
                goal_first_line_num[n])

    def test_color_matrix_Nc_restrictions(self):
        """Test the Nc power restriction during color basis building """

        goal = [
            fractions.Fraction(3, 8),
            fractions.Fraction(-9, 4),
            fractions.Fraction(45, 16)
        ]

        for n in range(3):
            myleglist = base_objects.LegList()

            myleglist.append(base_objects.Leg({'id': 21, 'state': False}))
            myleglist.append(base_objects.Leg({'id': 21, 'state': False}))

            myleglist.extend([base_objects.Leg({'id': 21, 'state': True})] * 2)

            myprocess = base_objects.Process({
                'legs': myleglist,
                'model': self.mymodel
            })

            myamplitude = diagram_generation.Amplitude()

            myamplitude.set('process', myprocess)

            myamplitude.generate_diagrams()

            col_basis = color_amp.ColorBasis(myamplitude)

            col_matrix = color_amp.ColorMatrix(col_basis,
                                               Nc=3,
                                               Nc_power_min=n,
                                               Nc_power_max=2 * n)

            for i in range(len(col_basis.items())):
                self.assertEqual(col_matrix.col_matrix_fixed_Nc[(i, i)],
                                 (goal[n], 0))

    def test_color_matrix_fixed_indices(self):
        """Test index fixing for immutable color string"""

        immutable1 = (('f', (1, -1, -3)), ('T', (-1, -3, 4)))
        immutable2 = (('d', (1, -2, -1)), ('T', (-1, -2, 4)))

        self.assertEqual(
            color_amp.ColorMatrix.fix_summed_indices(immutable1, immutable2),
            (('d', (1, -2, -5)), ('T', (-5, -2, 4))))

    def test_helper_lcm_functions(self):
        """Test the helper functions to derive common denominators for
        lines in the color matrix."""

        self.assertEqual(color_amp.ColorMatrix.lcm(6, 3), 6)
        self.assertEqual(color_amp.ColorMatrix.lcmm(6, 3, 5, 2), 30)
Example #6
0
    def setUp(self):
        # A gluon
        self.mypartlist.append(
            base_objects.Particle({
                'name': 'g',
                'antiname': 'g',
                'spin': 3,
                'color': 8,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'g',
                'antitexname': 'g',
                'line': 'curly',
                'charge': 0.,
                'pdg_code': 21,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))

        # A quark U and its antiparticle
        self.mypartlist.append(
            base_objects.Particle({
                'name': 'u',
                'antiname': 'u~',
                'spin': 2,
                'color': 3,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'u',
                'antitexname': '\bar u',
                'line': 'straight',
                'charge': 2. / 3.,
                'pdg_code': 2,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        antiu = copy.copy(self.mypartlist[1])
        antiu.set('is_part', False)

        # A quark D and its antiparticle
        self.mypartlist.append(
            base_objects.Particle({
                'name': 'd',
                'antiname': 'd~',
                'spin': 2,
                'color': 3,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'u',
                'antitexname': '\bar u',
                'line': 'straight',
                'charge': -1. / 3.,
                'pdg_code': 1,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        antid = copy.copy(self.mypartlist[2])
        antid.set('is_part', False)

        # 3 gluon vertiex
        self.myinterlist.append(base_objects.Interaction({
                      'id': 1,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[0]] * 3),
                      'color': [color.ColorString([color.f(0, 1, 2)])],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'G'},
                      'orders':{'QCD':1}}))

        # 4 gluon vertex
        self.myinterlist.append(base_objects.Interaction({
                      'id': 2,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[0]] * 4),
                      'color': [color.ColorString([color.f(-1, 0, 2),
                                                   color.f(-1, 1, 3)]),
                                color.ColorString([color.f(-1, 0, 3),
                                                   color.f(-1, 1, 2)]),
                                color.ColorString([color.f(-1, 0, 1),
                                                   color.f(-1, 2, 3)])],
                      'lorentz':['L(p1,p2,p3)', 'L(p2,p3,p1)', 'L3'],
                      'couplings':{(0, 0):'G^2',
                                   (1, 1):'G^2',
                                   (2, 2):'G^2'},
                      'orders':{'QCD':2}}))

        # Gluon couplings to up and down quarks
        self.myinterlist.append(base_objects.Interaction({
                      'id': 3,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[1], \
                                             antiu, \
                                             self.mypartlist[0]]),
                      'color': [color.ColorString([color.T(2, 0, 1)])],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQQ'},
                      'orders':{'QCD':1}}))

        self.myinterlist.append(base_objects.Interaction({
                      'id': 4,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[2], \
                                             antid, \
                                             self.mypartlist[0]]),
                      'color': [color.ColorString([color.T(2, 0, 1)])],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQQ'},
                      'orders':{'QCD':1}}))

        self.mymodel.set('particles', self.mypartlist)
        self.mymodel.set('interactions', self.myinterlist)
Example #7
0
class ColorAmpTest(unittest.TestCase):
    """Test class for the color_amp module"""

    mypartlist = base_objects.ParticleList()
    myinterlist = base_objects.InteractionList()
    mymodel = base_objects.Model()

    def setUp(self):
        # A gluon
        self.mypartlist.append(
            base_objects.Particle({
                'name': 'g',
                'antiname': 'g',
                'spin': 3,
                'color': 8,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'g',
                'antitexname': 'g',
                'line': 'curly',
                'charge': 0.,
                'pdg_code': 21,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))

        # A quark U and its antiparticle
        self.mypartlist.append(
            base_objects.Particle({
                'name': 'u',
                'antiname': 'u~',
                'spin': 2,
                'color': 3,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'u',
                'antitexname': '\bar u',
                'line': 'straight',
                'charge': 2. / 3.,
                'pdg_code': 2,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        antiu = copy.copy(self.mypartlist[1])
        antiu.set('is_part', False)

        # A quark D and its antiparticle
        self.mypartlist.append(
            base_objects.Particle({
                'name': 'd',
                'antiname': 'd~',
                'spin': 2,
                'color': 3,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'u',
                'antitexname': '\bar u',
                'line': 'straight',
                'charge': -1. / 3.,
                'pdg_code': 1,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        antid = copy.copy(self.mypartlist[2])
        antid.set('is_part', False)

        # A photon
        self.mypartlist.append(
            base_objects.Particle({
                'name': 'a',
                'antiname': 'a',
                'spin': 3,
                'color': 1,
                'mass': 'zero',
                'width': 'zero',
                'texname': '\gamma',
                'antitexname': '\gamma',
                'line': 'wavy',
                'charge': 0.,
                'pdg_code': 22,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))

        # A Higgs
        self.mypartlist.append(
            base_objects.Particle({
                'name': 'h',
                'antiname': 'h',
                'spin': 1,
                'color': 1,
                'mass': 'mh',
                'width': 'wh',
                'texname': 'h',
                'antitexname': 'h',
                'line': 'dashed',
                'charge': 0.,
                'pdg_code': 25,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))

        # 3 gluon vertiex
        self.myinterlist.append(base_objects.Interaction({
                      'id': 1,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[0]] * 3),
                      'color': [color.ColorString([color.f(0, 1, 2)])],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'G'},
                      'orders':{'QCD':1}}))

        # 4 gluon vertex
        self.myinterlist.append(base_objects.Interaction({
                      'id': 2,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[0]] * 4),
                      'color': [color.ColorString([color.f(-1, 0, 2),
                                                   color.f(-1, 1, 3)]),
                                color.ColorString([color.f(-1, 0, 3),
                                                   color.f(-1, 1, 2)]),
                                color.ColorString([color.f(-1, 0, 1),
                                                   color.f(-1, 2, 3)])],
                      'lorentz':['L(p1,p2,p3)', 'L(p2,p3,p1)', 'L3'],
                      'couplings':{(0, 0):'G^2',
                                   (1, 1):'G^2',
                                   (2, 2):'G^2'},
                      'orders':{'QCD':2}}))

        # Gluon couplings to up and down quarks
        self.myinterlist.append(base_objects.Interaction({
                      'id': 3,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[1], \
                                             antiu, \
                                             self.mypartlist[0]]),
                      'color': [color.ColorString([color.T(2, 0, 1)])],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQQ'},
                      'orders':{'QCD':1}}))

        self.myinterlist.append(base_objects.Interaction({
                      'id': 4,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[2], \
                                             antid, \
                                             self.mypartlist[0]]),
                      'color': [color.ColorString([color.T(2, 0, 1)])],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQQ'},
                      'orders':{'QCD':1}}))

        # Photon coupling to up
        self.myinterlist.append(base_objects.Interaction({
                      'id': 5,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[1], \
                                             antiu, \
                                             self.mypartlist[3]]),
                      'color': [color.ColorString([color.T(0, 1)])],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQED'},
                      'orders':{'QED':1}}))

        self.mymodel.set('particles', self.mypartlist)
        self.mymodel.set('interactions', self.myinterlist)

    def test_colorize_uu_gg(self):
        """Test the colorize function for uu~ > gg"""

        myleglist = base_objects.LegList()

        myleglist.append(base_objects.Leg({'id': -2, 'state': False}))
        myleglist.append(base_objects.Leg({'id': 2, 'state': False}))

        myleglist.extend([base_objects.Leg({'id': 21, 'state': True})] * 2)

        myprocess = base_objects.Process({
            'legs': myleglist,
            'model': self.mymodel
        })

        myamplitude = diagram_generation.Amplitude()

        myamplitude.set('process', myprocess)

        myamplitude.generate_diagrams()

        my_col_basis = color_amp.ColorBasis()

        # S channel
        col_dict = my_col_basis.colorize(myamplitude['diagrams'][0],
                                         self.mymodel)

        goal_dict = {
            (0, 0):
            color.ColorString([color.T(-1000, 1, 2),
                               color.f(3, 4, -1000)])
        }

        self.assertEqual(col_dict, goal_dict)

        # T channel
        col_dict = my_col_basis.colorize(myamplitude['diagrams'][1],
                                         self.mymodel)

        goal_dict = {
            (0, 0):
            color.ColorString([color.T(3, 1, -1000),
                               color.T(4, -1000, 2)])
        }

        self.assertEqual(col_dict, goal_dict)

        # U channel
        col_dict = my_col_basis.colorize(myamplitude['diagrams'][2],
                                         self.mymodel)

        goal_dict = {
            (0, 0):
            color.ColorString([color.T(4, 1, -1000),
                               color.T(3, -1000, 2)])
        }

        self.assertEqual(col_dict, goal_dict)

    def test_colorize_uux_ggg(self):
        """Test the colorize function for uu~ > ggg"""

        myleglist = base_objects.LegList()

        myleglist.append(base_objects.Leg({'id': 2, 'state': False}))
        myleglist.append(base_objects.Leg({'id': -2, 'state': False}))

        myleglist.extend([base_objects.Leg({'id': 21, 'state': True})] * 3)

        myprocess = base_objects.Process({
            'legs': myleglist,
            'model': self.mymodel
        })

        myamplitude = diagram_generation.Amplitude()

        myamplitude.set('process', myprocess)

        myamplitude.generate_diagrams()

        my_col_basis = color_amp.ColorBasis()

        # First diagram with two 3-gluon vertices
        col_dict = my_col_basis.colorize(myamplitude['diagrams'][0],
                                         self.mymodel)
        goal_dict = {
            (0, 0, 0):
            color.ColorString([
                color.T(-1000, 2, 1),
                color.f(-1001, 3, 4),
                color.f(-1000, -1001, 5)
            ])
        }

        self.assertEqual(col_dict, goal_dict)

        # Diagram with one 4-gluon vertex
        col_dict = my_col_basis.colorize(myamplitude['diagrams'][3],
                                         self.mymodel)

        goal_dict = {
            (0, 0):
            color.ColorString([
                color.T(-1000, 2, 1),
                color.f(-1001, 3, 5),
                color.f(-1001, 4, -1000)
            ]),
            (0, 1):
            color.ColorString([
                color.T(-1000, 2, 1),
                color.f(-1002, 3, -1000),
                color.f(-1002, 4, 5)
            ]),
            (0, 2):
            color.ColorString([
                color.T(-1000, 2, 1),
                color.f(-1003, 3, 4),
                color.f(-1003, 5, -1000)
            ])
        }

        self.assertEqual(col_dict, goal_dict)

    def test_colorize_funny_model(self):
        """Test the colorize function for uu~ > ggg"""

        mypartlist = base_objects.ParticleList()
        myinterlist = base_objects.InteractionList()
        mymodel = base_objects.Model()

        # A gluon
        mypartlist.append(
            base_objects.Particle({
                'name': 'g',
                'antiname': 'g',
                'spin': 3,
                'color': 8,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'g',
                'antitexname': 'g',
                'line': 'curly',
                'charge': 0.,
                'pdg_code': 21,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))

        # 3 gluon vertiex
        myinterlist.append(base_objects.Interaction({
                      'id': 1,
                      'particles': base_objects.ParticleList(\
                                            [mypartlist[0]] * 3),
                      'color': [color.ColorString([color.f(0, 1, 2)]),
                                color.ColorString([color.f(0, 2, 1)]),
                                color.ColorString([color.f(1, 2, 0)])],
                      'lorentz':['L1'],
            'couplings':{(0, 0):'G', (2,0):'G'},
                      'orders':{'QCD':1}}))

        mymodel.set('particles', mypartlist)
        mymodel.set('interactions', myinterlist)

        myleglist = base_objects.LegList()

        myleglist.append(base_objects.Leg({'id': 21, 'state': False}))
        myleglist.append(base_objects.Leg({'id': 21, 'state': False}))

        myleglist.extend([base_objects.Leg({'id': 21, 'state': True})] * 2)

        myprocess = base_objects.Process({'legs': myleglist, 'model': mymodel})

        myamplitude = diagram_generation.Amplitude()

        myamplitude.set('process', myprocess)

        myamplitude.generate_diagrams()

        my_col_basis = color_amp.ColorBasis()

        # Check that only the color structures that are actually used are included
        col_dict = my_col_basis.colorize(myamplitude['diagrams'][0], mymodel)
        goal_dict = {
            (2, 0):
            color.ColorString([color.f(1, 2, -1000),
                               color.f(3, 4, -1000)]),
            (0, 0):
            color.ColorString([color.f(-1000, 1, 2),
                               color.f(3, 4, -1000)]),
            (0, 2):
            color.ColorString([color.f(-1000, 1, 2),
                               color.f(4, -1000, 3)]),
            (2, 2):
            color.ColorString([color.f(1, 2, -1000),
                               color.f(4, -1000, 3)])
        }

        self.assertEqual(col_dict, goal_dict)

    def test_color_basis_uux_aggg(self):
        """Test the color basis building for uu~ > aggg (3! elements)"""

        myleglist = base_objects.LegList()

        myleglist.append(base_objects.Leg({'id': -2, 'state': False}))
        myleglist.append(base_objects.Leg({'id': 2, 'state': False}))

        myleglist.append(base_objects.Leg({'id': 22, 'state': True}))
        myleglist.extend([base_objects.Leg({'id': 21, 'state': True})] * 3)

        myprocess = base_objects.Process({
            'legs': myleglist,
            'model': self.mymodel
        })

        myamplitude = diagram_generation.Amplitude()

        myamplitude.set('process', myprocess)

        myamplitude.generate_diagrams()

        new_col_basis = color_amp.ColorBasis(myamplitude)

        self.assertEqual(len(new_col_basis), 6)

        # Test the color flow decomposition
        self.assertEqual(
            new_col_basis.color_flow_decomposition(
                {
                    1: 3,
                    2: -3,
                    3: 1,
                    4: 8,
                    5: 8,
                    6: -8
                }, 2), [{
                    1: [0, 501],
                    2: [504, 0],
                    3: [0, 0],
                    4: [502, 501],
                    5: [503, 502],
                    6: [504, 503]
                }, {
                    1: [0, 501],
                    2: [503, 0],
                    3: [0, 0],
                    4: [502, 501],
                    5: [503, 504],
                    6: [504, 502]
                }, {
                    1: [0, 501],
                    2: [504, 0],
                    3: [0, 0],
                    4: [502, 503],
                    5: [503, 501],
                    6: [504, 502]
                }, {
                    1: [0, 501],
                    2: [502, 0],
                    3: [0, 0],
                    4: [502, 504],
                    5: [503, 501],
                    6: [504, 503]
                }, {
                    1: [0, 501],
                    2: [503, 0],
                    3: [0, 0],
                    4: [502, 504],
                    5: [503, 502],
                    6: [504, 501]
                }, {
                    1: [0, 501],
                    2: [502, 0],
                    3: [0, 0],
                    4: [502, 503],
                    5: [503, 504],
                    6: [504, 501]
                }])

    def test_color_flow_string(self):
        """Test the color flow decomposition of various color strings"""

        # qq~>qq~
        my_cs = color.ColorString([color.T(-1000, 1, 2), color.T(-1000, 3, 4)])

        goal_cs = color.ColorString([color.T(1, 4), color.T(3, 2)])
        goal_cs.coeff = fractions.Fraction(1, 2)

        self.assertEqual(color_amp.ColorBasis.get_color_flow_string(my_cs, []),
                         goal_cs)

        # qq~>qq~g
        my_cs = color.ColorString(
            [color.T(-1000, 1, 2),
             color.T(-1000, 3, 4),
             color.T(5, 4, 6)])
        goal_cs = color.ColorString(
            [color.T(1, 2005),
             color.T(3, 2), color.T(1005, 6)])
        goal_cs.coeff = fractions.Fraction(1, 4)
        self.assertEqual(
            color_amp.ColorBasis.get_color_flow_string(my_cs,
                                                       [(8, 5, 1005, 2005)]),
            goal_cs)

        # gg>gg
        my_cs = color.ColorString(
            [color.Tr(-1000, 1, 2),
             color.Tr(-1000, 3, 4)])

        goal_cs = color.ColorString([
            color.T(1001, 2002),
            color.T(1002, 2003),
            color.T(1003, 2004),
            color.T(1004, 2001)
        ])
        goal_cs.coeff = fractions.Fraction(1, 32)

        self.assertEqual(
            color_amp.ColorBasis.get_color_flow_string(my_cs,
                                                       [(8, 1, 1001, 2001),
                                                        (8, 2, 1002, 2002),
                                                        (8, 3, 1003, 2003),
                                                        (8, 4, 1004, 2004)]),
            goal_cs)

    def test_color_flow_string_epsilon(self):
        """Test the color flow decomposition of strings including Epsilon
        and color sextets"""

        # g q > trip q
        my_cs = color.ColorString(
            [color.Epsilon(-1000, 2, 3),
             color.T(1, 4, -1000)])

        goal_cs = color.ColorString(
            [color.T(4, 2001), color.Epsilon(2, 3, 1001)])
        goal_cs.coeff = fractions.Fraction(1, 2)
        self.assertEqual(
            color_amp.ColorBasis.get_color_flow_string(my_cs,
                                                       [(8, 1, 1001, 2001)]),
            goal_cs)

        # g q > six q~
        my_cs = color.ColorString(
            [color.K6(3, -1000, 4),
             color.T(1, -1000, 2)])
        goal_cs = color.ColorString(
            [color.T(1001, 2),
             color.T(1003, 4),
             color.T(3003, 2001)])
        goal_cs.coeff = fractions.Fraction(1, 4)
        self.assertEqual(
            color_amp.ColorBasis.get_color_flow_string(my_cs,
                                                       [(8, 1, 1001, 2001),
                                                        (6, 3, 1003, 3003)]),
            goal_cs)

        # g q~ > trip > q~ q q~

        my_cs = color.ColorString([
            color.Epsilon(-1000, 2, 4),
            color.EpsilonBar(-1001, 3, 5),
            color.T(1, -1001, -1000)
        ])

        goal_cs = color.ColorString(
            [color.Epsilon(2, 4, 1001),
             color.EpsilonBar(3, 5, 2001)])
        goal_cs.coeff = fractions.Fraction(1, 2)
        self.assertEqual(
            color_amp.ColorBasis.get_color_flow_string(my_cs,
                                                       [(8, 1, 1001, 2001)]),
            goal_cs)
Example #8
0
    def test_colorize_funny_model(self):
        """Test the colorize function for uu~ > ggg"""

        mypartlist = base_objects.ParticleList()
        myinterlist = base_objects.InteractionList()
        mymodel = base_objects.Model()

        # A gluon
        mypartlist.append(
            base_objects.Particle({
                'name': 'g',
                'antiname': 'g',
                'spin': 3,
                'color': 8,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'g',
                'antitexname': 'g',
                'line': 'curly',
                'charge': 0.,
                'pdg_code': 21,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))

        # 3 gluon vertiex
        myinterlist.append(base_objects.Interaction({
                      'id': 1,
                      'particles': base_objects.ParticleList(\
                                            [mypartlist[0]] * 3),
                      'color': [color.ColorString([color.f(0, 1, 2)]),
                                color.ColorString([color.f(0, 2, 1)]),
                                color.ColorString([color.f(1, 2, 0)])],
                      'lorentz':['L1'],
            'couplings':{(0, 0):'G', (2,0):'G'},
                      'orders':{'QCD':1}}))

        mymodel.set('particles', mypartlist)
        mymodel.set('interactions', myinterlist)

        myleglist = base_objects.LegList()

        myleglist.append(base_objects.Leg({'id': 21, 'state': False}))
        myleglist.append(base_objects.Leg({'id': 21, 'state': False}))

        myleglist.extend([base_objects.Leg({'id': 21, 'state': True})] * 2)

        myprocess = base_objects.Process({'legs': myleglist, 'model': mymodel})

        myamplitude = diagram_generation.Amplitude()

        myamplitude.set('process', myprocess)

        myamplitude.generate_diagrams()

        my_col_basis = color_amp.ColorBasis()

        # Check that only the color structures that are actually used are included
        col_dict = my_col_basis.colorize(myamplitude['diagrams'][0], mymodel)
        goal_dict = {
            (2, 0):
            color.ColorString([color.f(1, 2, -1000),
                               color.f(3, 4, -1000)]),
            (0, 0):
            color.ColorString([color.f(-1000, 1, 2),
                               color.f(3, 4, -1000)]),
            (0, 2):
            color.ColorString([color.f(-1000, 1, 2),
                               color.f(4, -1000, 3)]),
            (2, 2):
            color.ColorString([color.f(1, 2, -1000),
                               color.f(4, -1000, 3)])
        }

        self.assertEqual(col_dict, goal_dict)
Example #9
0
    'antiname': 'h',
    'spin': 1,
    'color': 1,
    'mass': 'mh',
    'width': 'wh',
    'texname': 'h',
    'antitexname': 'h',
    'line': 'dashed',
    'charge': 0.,
    'pdg_code': 25,
    'propagating': True,
    'is_part': True,
    'self_antipart': True
})

particles = base_objects.ParticleList([gluon, up, down, photon, higgs])

# Interactions
#=========================================================================================

interactions = base_objects.InteractionList()

# g g g
interactions.append(base_objects.Interaction({
    'id': 1,
    'particles': base_objects.ParticleList([gluon, ] * 3),
    'color': [color.ColorString([color.f(0, 1, 2)])],
    'lorentz': ['L1'],
    'couplings': {(0, 0): 'G'},
    'orders': {'QCD': 1}
}))
Example #10
0
class ColorAmpTest(unittest.TestCase):
    """Test class for the color_amp module"""

    mypartlist = base_objects.ParticleList()
    myinterlist = base_objects.InteractionList()
    mymodel = base_objects.Model()

    def setUp(self):
        # A gluon
        self.mypartlist.append(
            base_objects.Particle({
                'name': 'g',
                'antiname': 'g',
                'spin': 3,
                'color': 8,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'g',
                'antitexname': 'g',
                'line': 'curly',
                'charge': 0.,
                'pdg_code': 21,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))

        # A quark U and its antiparticle
        self.mypartlist.append(
            base_objects.Particle({
                'name': 'u',
                'antiname': 'u~',
                'spin': 2,
                'color': 3,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'u',
                'antitexname': '\bar u',
                'line': 'straight',
                'charge': 2. / 3.,
                'pdg_code': 2,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        antiu = copy.copy(self.mypartlist[1])
        antiu.set('is_part', False)

        # A quark D and its antiparticle
        self.mypartlist.append(
            base_objects.Particle({
                'name': 'd',
                'antiname': 'd~',
                'spin': 2,
                'color': 3,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'u',
                'antitexname': '\bar u',
                'line': 'straight',
                'charge': -1. / 3.,
                'pdg_code': 1,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        antid = copy.copy(self.mypartlist[2])
        antid.set('is_part', False)

        # A photon
        self.mypartlist.append(
            base_objects.Particle({
                'name': 'a',
                'antiname': 'a',
                'spin': 3,
                'color': 1,
                'mass': 'zero',
                'width': 'zero',
                'texname': '\gamma',
                'antitexname': '\gamma',
                'line': 'wavy',
                'charge': 0.,
                'pdg_code': 22,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))

        # A Higgs
        self.mypartlist.append(
            base_objects.Particle({
                'name': 'h',
                'antiname': 'h',
                'spin': 1,
                'color': 1,
                'mass': 'mh',
                'width': 'wh',
                'texname': 'h',
                'antitexname': 'h',
                'line': 'dashed',
                'charge': 0.,
                'pdg_code': 25,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))

        # An electron e- and its antiparticle
        self.mypartlist.append(
            base_objects.Particle({
                'name': 'e-',
                'antiname': 'e+',
                'spin': 2,
                'color': 1,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'e^-',
                'antitexname': 'e^+',
                'line': 'straight',
                'charge': 1,
                'pdg_code': 11,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        antie = copy.copy(self.mypartlist[5])
        antie.set('is_part', False)

        # 3 gluon vertiex
        self.myinterlist.append(base_objects.Interaction({
                      'id': 1,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[0]] * 3),
                      'color': [color.ColorString([color.f(0, 1, 2)])],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'G'},
                      'orders':{'QCD':1}}))

        # 4 gluon vertex
        self.myinterlist.append(base_objects.Interaction({
                      'id': 2,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[0]] * 4),
                      'color': [color.ColorString([color.f(-1, 0, 2),
                                                   color.f(-1, 1, 3)]),
                                color.ColorString([color.f(-1, 0, 3),
                                                   color.f(-1, 1, 2)]),
                                color.ColorString([color.f(-1, 0, 1),
                                                   color.f(-1, 2, 3)])],
                      'lorentz':['L(p1,p2,p3)', 'L(p2,p3,p1)', 'L3'],
                      'couplings':{(0, 0):'G^2',
                                   (1, 1):'G^2',
                                   (2, 2):'G^2'},
                      'orders':{'QCD':2}}))

        # Gluon couplings to up and down quarks
        self.myinterlist.append(base_objects.Interaction({
                      'id': 3,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[1], \
                                             antiu, \
                                             self.mypartlist[0]]),
                      'color': [color.ColorString([color.T(2, 0, 1)])],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQQ'},
                      'orders':{'QCD':1}}))

        self.myinterlist.append(base_objects.Interaction({
                      'id': 4,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[2], \
                                             antid, \
                                             self.mypartlist[0]]),
                      'color': [color.ColorString([color.T(2, 0, 1)])],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQQ'},
                      'orders':{'QCD':1}}))

        # Photon coupling to up
        self.myinterlist.append(base_objects.Interaction({
                      'id': 5,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[1], \
                                             antiu, \
                                             self.mypartlist[3]]),
                      'color': [color.ColorString([color.T(0, 1)])],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQED'},
                      'orders':{'QED':1}}))

        # Photon coupling to down
        self.myinterlist.append(base_objects.Interaction({
                      'id': 6,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[2], \
                                             antid, \
                                             self.mypartlist[3]]),
                      'color': [color.ColorString([color.T(0, 1)])],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQED'},
                      'orders':{'QED':1}}))

        # Photon coupling to electrons
        self.myinterlist.append(base_objects.Interaction({
                      'id': 7,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[5], \
                                             antie, \
                                             self.mypartlist[3]]),
                      'color': [color.ColorString([color.T(0, 1)])],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQED'},
                      'orders':{'QED':1}}))

        self.mymodel.set('particles', self.mypartlist)
        self.mymodel.set('interactions', self.myinterlist)

    def test_colorize_uu_gg(self):
        """Test the colorize function for uu~ > gg"""

        myleglist = base_objects.LegList()

        myleglist.append(base_objects.Leg({'id': -2, 'state': False}))
        myleglist.append(base_objects.Leg({'id': 2, 'state': False}))

        myleglist.extend([base_objects.Leg({'id': 21, 'state': True})] * 2)

        myprocess = base_objects.Process({
            'legs': myleglist,
            'model': self.mymodel
        })

        myamplitude = diagram_generation.Amplitude()

        myamplitude.set('process', myprocess)

        myamplitude.generate_diagrams()

        my_col_basis = color_amp.ColorBasis()

        # S channel
        col_dict = my_col_basis.colorize(myamplitude['diagrams'][0],
                                         self.mymodel)

        goal_dict = {
            (0, 0):
            color.ColorString([color.T(-1000, 1, 2),
                               color.f(3, 4, -1000)])
        }

        self.assertEqual(col_dict, goal_dict)

        # T channel
        col_dict = my_col_basis.colorize(myamplitude['diagrams'][1],
                                         self.mymodel)

        goal_dict = {
            (0, 0):
            color.ColorString([color.T(3, 1, -1000),
                               color.T(4, -1000, 2)])
        }

        self.assertEqual(col_dict, goal_dict)

        # U channel
        col_dict = my_col_basis.colorize(myamplitude['diagrams'][2],
                                         self.mymodel)

        goal_dict = {
            (0, 0):
            color.ColorString([color.T(4, 1, -1000),
                               color.T(3, -1000, 2)])
        }

        self.assertEqual(col_dict, goal_dict)

    def test_colorize_uux_ggg(self):
        """Test the colorize function for uu~ > ggg"""

        myleglist = base_objects.LegList()

        myleglist.append(base_objects.Leg({'id': 2, 'state': False}))
        myleglist.append(base_objects.Leg({'id': -2, 'state': False}))

        myleglist.extend([base_objects.Leg({'id': 21, 'state': True})] * 3)

        myprocess = base_objects.Process({
            'legs': myleglist,
            'model': self.mymodel
        })

        myamplitude = diagram_generation.Amplitude()

        myamplitude.set('process', myprocess)

        myamplitude.generate_diagrams()

        my_col_basis = color_amp.ColorBasis()

        # First diagram with two 3-gluon vertices
        col_dict = my_col_basis.colorize(myamplitude['diagrams'][0],
                                         self.mymodel)
        goal_dict = {
            (0, 0, 0):
            color.ColorString([
                color.T(-1000, 2, 1),
                color.f(-1001, 3, 4),
                color.f(-1000, -1001, 5)
            ])
        }

        self.assertEqual(col_dict, goal_dict)

        # Diagram with one 4-gluon vertex
        col_dict = my_col_basis.colorize(myamplitude['diagrams'][3],
                                         self.mymodel)

        goal_dict = {
            (0, 0):
            color.ColorString([
                color.T(-1000, 2, 1),
                color.f(-1001, 3, 5),
                color.f(-1001, 4, -1000)
            ]),
            (0, 1):
            color.ColorString([
                color.T(-1000, 2, 1),
                color.f(-1002, 3, -1000),
                color.f(-1002, 4, 5)
            ]),
            (0, 2):
            color.ColorString([
                color.T(-1000, 2, 1),
                color.f(-1003, 3, 4),
                color.f(-1003, 5, -1000)
            ])
        }

        self.assertEqual(col_dict, goal_dict)

    def test_colorize_funny_model(self):
        """Test the colorize function for uu~ > ggg"""

        mypartlist = base_objects.ParticleList()
        myinterlist = base_objects.InteractionList()
        mymodel = base_objects.Model()

        # A gluon
        mypartlist.append(
            base_objects.Particle({
                'name': 'g',
                'antiname': 'g',
                'spin': 3,
                'color': 8,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'g',
                'antitexname': 'g',
                'line': 'curly',
                'charge': 0.,
                'pdg_code': 21,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))

        # 3 gluon vertiex
        myinterlist.append(base_objects.Interaction({
                      'id': 1,
                      'particles': base_objects.ParticleList(\
                                            [mypartlist[0]] * 3),
                      'color': [color.ColorString([color.f(0, 1, 2)]),
                                color.ColorString([color.f(0, 2, 1)]),
                                color.ColorString([color.f(1, 2, 0)])],
                      'lorentz':['L1'],
            'couplings':{(0, 0):'G', (2,0):'G'},
                      'orders':{'QCD':1}}))

        mymodel.set('particles', mypartlist)
        mymodel.set('interactions', myinterlist)

        myleglist = base_objects.LegList()

        myleglist.append(base_objects.Leg({'id': 21, 'state': False}))
        myleglist.append(base_objects.Leg({'id': 21, 'state': False}))

        myleglist.extend([base_objects.Leg({'id': 21, 'state': True})] * 2)

        myprocess = base_objects.Process({'legs': myleglist, 'model': mymodel})

        myamplitude = diagram_generation.Amplitude()

        myamplitude.set('process', myprocess)

        myamplitude.generate_diagrams()

        my_col_basis = color_amp.ColorBasis()

        # Check that only the color structures that are actually used are included
        col_dict = my_col_basis.colorize(myamplitude['diagrams'][0], mymodel)
        goal_dict = {
            (2, 0):
            color.ColorString([color.f(1, 2, -1000),
                               color.f(3, 4, -1000)]),
            (0, 0):
            color.ColorString([color.f(-1000, 1, 2),
                               color.f(3, 4, -1000)]),
            (0, 2):
            color.ColorString([color.f(-1000, 1, 2),
                               color.f(4, -1000, 3)]),
            (2, 2):
            color.ColorString([color.f(1, 2, -1000),
                               color.f(4, -1000, 3)])
        }

        self.assertEqual(col_dict, goal_dict)

    def test_color_basis_uux_aggg(self):
        """Test the color basis building for uu~ > aggg (3! elements)"""

        myleglist = base_objects.LegList()

        myleglist.append(base_objects.Leg({'id': -2, 'state': False}))
        myleglist.append(base_objects.Leg({'id': 2, 'state': False}))

        myleglist.append(base_objects.Leg({'id': 22, 'state': True}))
        myleglist.extend([base_objects.Leg({'id': 21, 'state': True})] * 3)

        myprocess = base_objects.Process({
            'legs': myleglist,
            'model': self.mymodel
        })

        myamplitude = diagram_generation.Amplitude()

        myamplitude.set('process', myprocess)

        myamplitude.generate_diagrams()

        new_col_basis = color_amp.ColorBasis(myamplitude)

        self.assertEqual(len(new_col_basis), 6)

        # Test the color flow decomposition
        self.assertEqual(
            new_col_basis.color_flow_decomposition(
                {
                    1: 3,
                    2: -3,
                    3: 1,
                    4: 8,
                    5: 8,
                    6: -8
                }, 2), [{
                    1: [0, 501],
                    2: [504, 0],
                    3: [0, 0],
                    4: [502, 501],
                    5: [503, 502],
                    6: [504, 503]
                }, {
                    1: [0, 501],
                    2: [503, 0],
                    3: [0, 0],
                    4: [502, 501],
                    5: [503, 504],
                    6: [504, 502]
                }, {
                    1: [0, 501],
                    2: [504, 0],
                    3: [0, 0],
                    4: [502, 503],
                    5: [503, 501],
                    6: [504, 502]
                }, {
                    1: [0, 501],
                    2: [502, 0],
                    3: [0, 0],
                    4: [502, 504],
                    5: [503, 501],
                    6: [504, 503]
                }, {
                    1: [0, 501],
                    2: [503, 0],
                    3: [0, 0],
                    4: [502, 504],
                    5: [503, 502],
                    6: [504, 501]
                }, {
                    1: [0, 501],
                    2: [502, 0],
                    3: [0, 0],
                    4: [502, 503],
                    5: [503, 504],
                    6: [504, 501]
                }])

    def test_color_flow_string(self):
        """Test the color flow decomposition of various color strings"""

        # qq~>qq~
        my_cs = color.ColorString([color.T(-1000, 1, 2), color.T(-1000, 3, 4)])

        goal_cs = color.ColorString([color.T(1, 4), color.T(3, 2)])
        goal_cs.coeff = fractions.Fraction(1, 2)

        self.assertEqual(color_amp.ColorBasis.get_color_flow_string(my_cs, []),
                         goal_cs)

        # qq~>qq~g
        my_cs = color.ColorString(
            [color.T(-1000, 1, 2),
             color.T(-1000, 3, 4),
             color.T(5, 4, 6)])
        goal_cs = color.ColorString(
            [color.T(1, 2005),
             color.T(3, 2), color.T(1005, 6)])
        goal_cs.coeff = fractions.Fraction(1, 4)
        self.assertEqual(
            color_amp.ColorBasis.get_color_flow_string(my_cs,
                                                       [(8, 5, 1005, 2005)]),
            goal_cs)

        # gg>gg
        my_cs = color.ColorString(
            [color.Tr(-1000, 1, 2),
             color.Tr(-1000, 3, 4)])

        goal_cs = color.ColorString([
            color.T(1001, 2002),
            color.T(1002, 2003),
            color.T(1003, 2004),
            color.T(1004, 2001)
        ])
        goal_cs.coeff = fractions.Fraction(1, 32)

        self.assertEqual(
            color_amp.ColorBasis.get_color_flow_string(my_cs,
                                                       [(8, 1, 1001, 2001),
                                                        (8, 2, 1002, 2002),
                                                        (8, 3, 1003, 2003),
                                                        (8, 4, 1004, 2004)]),
            goal_cs)

    def test_color_flow_string_epsilon(self):
        """Test the color flow decomposition of strings including Epsilon
        and color sextets"""

        # g q > trip q
        my_cs = color.ColorString(
            [color.Epsilon(-1000, 2, 3),
             color.T(1, 4, -1000)])

        goal_cs = color.ColorString(
            [color.T(4, 2001), color.Epsilon(2, 3, 1001)])
        goal_cs.coeff = fractions.Fraction(1, 2)
        self.assertEqual(
            color_amp.ColorBasis.get_color_flow_string(my_cs,
                                                       [(8, 1, 1001, 2001)]),
            goal_cs)

        # g q > six q~
        my_cs = color.ColorString(
            [color.K6(3, -1000, 4),
             color.T(1, -1000, 2)])
        goal_cs = color.ColorString(
            [color.T(1001, 2),
             color.T(1003, 4),
             color.T(3003, 2001)])
        goal_cs.coeff = fractions.Fraction(1, 4)
        self.assertEqual(
            color_amp.ColorBasis.get_color_flow_string(my_cs,
                                                       [(8, 1, 1001, 2001),
                                                        (6, 3, 1003, 3003)]),
            goal_cs)

        # g q~ > trip > q~ q q~

        my_cs = color.ColorString([
            color.Epsilon(-1000, 2, 4),
            color.EpsilonBar(-1001, 3, 5),
            color.T(1, -1001, -1000)
        ])

        goal_cs = color.ColorString(
            [color.Epsilon(2, 4, 1001),
             color.EpsilonBar(3, 5, 2001)])
        goal_cs.coeff = fractions.Fraction(1, 2)
        self.assertEqual(
            color_amp.ColorBasis.get_color_flow_string(my_cs,
                                                       [(8, 1, 1001, 2001)]),
            goal_cs)

    def perform_color_connections_test(self, combinations):
        """ Handy wrapping function to test selected cases for the color correlated matrices."""
        def compare_result(color_matrices, color_connections, reference):

            important_info_CM = [(c[0], dict(c[1][1].col_matrix_fixed_Nc))
                                 for c in color_matrices
                                 if c[1][1].col_matrix_fixed_Nc is not None]

            important_info_CC = {
                k: [(i, cc['tuple_representation']) for i, cc in enumerate(v)]
                for k, v in color_connections.items()
            }

            #pprint(important_info_CM)
            #pprint(important_info_CC)
            self.assertListEqual(important_info_CM,
                                 reference['color_matrices'])
            self.assertDictEqual(important_info_CC,
                                 reference['color_connections'])

        for n_gluons, initial_states, max_order, reference in combinations:
            myleglist = base_objects.LegList()
            myleglist.append(
                base_objects.Leg({
                    'id': initial_states[0],
                    'state': False
                }))
            myleglist.append(
                base_objects.Leg({
                    'id': initial_states[1],
                    'state': False
                }))
            myleglist.append(base_objects.Leg({'id': 1, 'state': True}))
            myleglist.append(base_objects.Leg({'id': -1, 'state': True}))
            myleglist.extend([base_objects.Leg({
                'id': 21,
                'state': True
            })] * (n_gluons))

            myprocess = base_objects.Process({
                'legs': myleglist,
                'model': self.mymodel
            })

            myamplitude = diagram_generation.Amplitude()
            myamplitude.set('process', myprocess)
            myamplitude.generate_diagrams()

            col_basis = color_amp.ColorBasis(myamplitude)

            col_matrix = color_amp.ColorMatrix(col_basis, Nc=3)

            all_color_matrices, color_connections = col_matrix.build_color_correlated_matrices(
                myleglist,
                self.mymodel,
                order=max_order,
                Nc=3,
                Nc_power_min=None,
                Nc_power_max=None)

            # Now aggregate all color_connections in a single list:
            all_color_connections = []
            for o in range(max_order.count('N') + 1):
                all_color_connections.extend(color_connections['N' * o + 'LO'])

            sorted_color_matrices = sorted(all_color_matrices.items(),
                                           key=lambda el: el[0])

            compare_result(sorted_color_matrices, color_connections, reference)

            if False:
                pprint([(
                    '(%d,%d) -> [ %s | %s ]' %
                    (k[0], k[1],
                     str(all_color_connections[k[0]]['tuple_representation']),
                     str(all_color_connections[k[1]]['tuple_representation'])),
                    (v[0], v[1].col_matrix_fixed_Nc))
                        for k, v in sorted_color_matrices])

            if False:
                print(
                    "\nA total of %d correlator matrices have been computed (among which %d are zero)"
                    % (len(sorted_color_matrices),
                       len([
                           _ for _ in sorted_color_matrices
                           if _[1][1].col_matrix_fixed_Nc is None
                       ])))

    def test_color_correlators_computation_slow(self):
        """ Test the functions in ColorMatrix related to the generation of color-correlated 
        Matrix elements, for cases that are fast to evaluate."""

        from color_connections_reference_results import \
            epem_ddxgg_NNLO_res, \
            ddx_ddx_NNLO_res, \
            epem_ddx_NNNLO_res

        # The tests that are commented are not because they're not working but just
        # because they are too slow for unit tests
        combinations = [

            ## NNLO
            #######
            ## e+ e- > d d~ g g @NNLO
            (2, [11, -11], 'NNLO', epem_ddxgg_NNLO_res),
            ## d d~ > d d~ @NNLO
            (0, [1, -1], 'NNLO', ddx_ddx_NNLO_res),

            ## NNNLO
            #######
            ## e+ e- > d d~ @NNNLO
            (0, [11, -11], 'NNNLO', epem_ddx_NNNLO_res),
            ## g g > d d~ @NNNLO
            # These results are too large to be stored in the hardcoded reference
            # results and we skip them. (also too slow for unit tests.
            # But it's working in principle.
            #(0,[21,21],'NNNLO',gg_ddx_NNNLO_res),
            ## d d~ > d d~ @NNNLO
            #(0,[1,-1],'NNNLO',ddx_ddx_NNNLO_res),
        ]

        self.perform_color_connections_test(combinations)

    def test_color_correlators_computation_fast(self):
        """ Test the functions in ColorMatrix related to the generation of color-correlated 
        Matrix elements, for cases that are slow to evaluate."""

        from color_connections_reference_results import \
            epem_ddx_NLO_res, \
            epem_ddx_NNLO_res, \
            epem_ddxg_NLO_res, \
            epem_ddxg_NNLO_res, \
            epem_ddxgg_NLO_res, \
            gg_ddx_NLO_res, \
            gg_ddxg_NLO_res, \
            ddx_ddx_NLO_res

        # The tests that are commented are not because they're not working but just
        # because they are too slow for unit tests
        combinations = [

            ## NLO
            ######
            ## e+ e- > d d~ @NLO
            (0, [11, -11], 'NLO', epem_ddx_NLO_res),
            ## e+ e- > d d~ g @NLO
            (1, [11, -11], 'NLO', epem_ddxg_NLO_res),
            ## e+ e- > d d~ g g @NLO
            (2, [11, -11], 'NLO', epem_ddxgg_NLO_res),
            ## g g > d d~ @NLO
            (0, [21, 21], 'NLO', gg_ddx_NLO_res),
            ## g g > d d~ g @NLO
            (1, [21, 21], 'NLO', gg_ddxg_NLO_res),
            ## d d~ > d d~ @NLO
            (0, [1, -1], 'NLO', ddx_ddx_NLO_res),

            ## NNLO
            #######
            ## e+ e- > d d~ @NNLO
            (0, [11, -11], 'NNLO', epem_ddx_NNLO_res),
            ## e+ e- > d d~ g @NNLO
            (1, [11, -11], 'NNLO', epem_ddxg_NNLO_res),
        ]

        self.perform_color_connections_test(combinations)
def loadLoopToyModel():
    """Setup the Loop Toy QCD model with d,dx,u,ux and the gluon"""
    
    mypartlist = base_objects.ParticleList()
    myinterlist = base_objects.InteractionList()
    myloopmodel = loop_base_objects.LoopModel()

    # A gluon
    mypartlist.append(base_objects.Particle({'name':'g',
                  'antiname':'g',
                  'spin':3,
                  'color':8,
                  'mass':'ZERO',
                  'width':'ZERO',
                  'texname':'G',
                  'antitexname':'G',
                  'line':'curly',
                  'charge':0.,
                  'pdg_code':21,
                  'propagating':True,
                  'propagator':0,                                             
                  'is_part':True,
                  'counterterm':{('QCD', ((6,),)): {0: 'UVWfct_G_1', -1: 'UVWfct_G_1_1eps'}, ('QCD', ((5,),)): {0: 'UVWfct_G_0', -1: 'UVWfct_G_0_1eps'}},
                  'self_antipart':True}))
    
    # A quark U and its antiparticle
    mypartlist.append(base_objects.Particle({'name':'u',
                  'antiname':'u~',
                  'spin':2,
                  'color':3,
                  'mass':'ZERO',
                  'width':'ZERO',
                  'texname':'u',
                  'antitexname':'u',
                  'line':'straight',
                  'charge':2. / 3.,
                  'pdg_code':2,
                  'propagating':True,
                  'propagator':'',                                             
                  'is_part':True,
                  'self_antipart':False}))
    antiu = copy.copy(mypartlist[1])
    antiu.set('is_part', False)

    # A quark D and its antiparticle
    mypartlist.append(base_objects.Particle({'name':'d',
                  'antiname':'d~',
                  'spin':2,
                  'color':3,
                  'mass':'ZERO',
                  'width':'ZERO',
                  'texname':'d',
                  'antitexname':'d',
                  'line':'straight',
                  'charge':-1. / 3.,
                  'pdg_code':1,
                  'propagating':True,
                  'propagator':'',
                  'is_part':True,
                  'self_antipart':False}))
    antid = copy.copy(mypartlist[2])
    antid.set('is_part', False)

    myloopmodel.set('particles', mypartlist)
    myloopmodel.set('couplings', ['QCD'])        
    myloopmodel.set('interactions', myinterlist)
    myloopmodel.set('perturbation_couplings', ['QCD'])
    myloopmodel.set('order_hierarchy', {'QCD':1})

    return myloopmodel
Example #12
0
    def test_read_interactions(self):
        """Test the output of import interactions.dat file"""

        particles_dat_str = """ve ve~ F S ZERO ZERO S ve 12
                                vm vm~ F S ZERO ZERO S vm 14
                                vt vt~ F S ZERO ZERO S vt 16
                                e- e+ F S ZERO ZERO S e 11
                                m- m+ F S ZERO ZERO S m 13
                                tt- tt+ F S MTA ZERO S tt 15
                                u u~ F S ZERO ZERO T u 2
                                c c~ F S MC ZERO T c 4
                                t t~ F S MT WT T t 6
                                d d~ F S ZERO ZERO T d 1
                                s s~ F S ZERO ZERO T s 3
                                b b~ F S MB ZERO T b 5
                                a a V W ZERO ZERO S a 22
                                z z V W MZ WZ S Z 23
                                w+ w- V W MW WW S W 24
                                g g V C ZERO ZERO O G 21
                                h h S D MH WH S H 25
                                T1 T1 T D ZERO ZERO O T1 8000002"""

        interactions_dat_str = """# Interactions associated with Standard_Model
                                    w+   w-   a MGVX3   QED
                                    g   g   T1 MGVX2   QCD a
                                    w+   w-   w+   w- MGVX6   DUM0   QED QED n
                                    e-   ve   w- MGVX24   QED
                                    e+   ve~   w+ MGVX25   QED
                                    u   u   g MGVX1   QCD
                                    u   u   a MGVX4   QED
                                    # And now some bad format entries
                                    # which should be ignored with a warning
                                    k+ k- a test QED
                                    g g test QCD"""

        fsock_part = StringIO.StringIO(particles_dat_str)
        fsock_inter = StringIO.StringIO(interactions_dat_str)

        myparts = import_v4.read_particles_v4(fsock_part)

        wplus = copy.copy(myparts[14])
        wmin = copy.copy(myparts[14])
        wmin.set('is_part', False)
        eminus = copy.copy(myparts[3])
        eplus = copy.copy(myparts[3])
        eplus.set('is_part', False)
        enu = copy.copy(myparts[0])
        enubar = copy.copy(myparts[0])
        enubar.set('is_part', False)
        photon = copy.copy(myparts[12])
        gluon = copy.copy(myparts[15])
        t1 = copy.copy(myparts[17])
        u = myparts[6]
        ubar = copy.copy(myparts[6])
        ubar.set('is_part', False)

        my_i_f = color.ColorString([color.f(0, 1, 2)])
        my_i_f.is_imaginary = True

        goal_inter_list = base_objects.InteractionList([ \
                    base_objects.Interaction(
                                    {'id':1,
                                     'particles':base_objects.ParticleList([
                                                                wplus,
                                                                wmin,
                                                                photon]),

                                     'color':[],
                                     'lorentz':[''],
                                     'couplings':{(0, 0):'MGVX3'},
                                     'orders':{'QED':1}}),
                     base_objects.Interaction(
                                    {'id':2,
                                     'particles':base_objects.ParticleList([
                                                                gluon,
                                                                gluon,
                                                                t1]),

                                     'color':[my_i_f],
                                     'lorentz':['A'],
                                     'couplings':{(0, 0):'MGVX2'},
                                     'orders':{'QCD':1}}),
                     base_objects.Interaction(
                                    {'id':3,
                                     'particles':base_objects.ParticleList([
                                                                wplus,
                                                                wmin,
                                                                wplus,
                                                                wmin]),
                                     'color':[],
                                     'lorentz':['WWVVN'],
                                     'couplings':{(0, 0):'MGVX6'},
                                     'orders':{'QED':2}}),

                     base_objects.Interaction(
                                    {'id':4,
                                     'particles':base_objects.ParticleList([
                                                                eplus,
                                                                enu,
                                                                wmin]),
                                     'color':[],
                                     'lorentz':[''],
                                     'couplings':{(0, 0):'MGVX24'},
                                     'orders':{'QED':1}}),

                     base_objects.Interaction(
                                    {'id':5,
                                     'particles':base_objects.ParticleList([
                                                                eminus,
                                                                enubar,
                                                                wplus]),
                                     'color':[],
                                     'lorentz':[''],
                                     'couplings':{(0, 0):'MGVX25'},
                                     'orders':{'QED':1}}),

                     base_objects.Interaction(
                                    {'id':6,
                                     'particles':base_objects.ParticleList([
                                                                ubar,
                                                                u,
                                                                gluon]),
                                     'color':[color.ColorString(\
                                              [color.T(2, 1, 0)])],
                                     'lorentz':[''],
                                     'couplings':{(0, 0):'MGVX1'},
                                     'orders':{'QCD':1}}),

                     base_objects.Interaction(
                                    {'id':7,
                                     'particles':base_objects.ParticleList([
                                                                ubar,
                                                                u,
                                                                photon]),
                                     'color':[color.ColorString(\
                                              [color.T(1, 0)])],
                                     'lorentz':[''],
                                     'couplings':{(0, 0):'MGVX4'},
                                     'orders':{'QED':1}})])
        
        result = import_v4.read_interactions_v4(fsock_inter, myparts)
        self.assertEqual(len(result), len(goal_inter_list))
        for i in range(len(result)):
            self.assertEqual(result[i], goal_inter_list[i])
Example #13
0
    def test_ungrouping_lepton(self):
        """check that the routines which ungroup the process for the muon/electron processes
        works as expected"""

        # Setup A simple model
        mypartlist = base_objects.ParticleList()
        myinterlist = base_objects.InteractionList()

        # A quark U and its antiparticle
        mypartlist.append(
            base_objects.Particle({
                'name': 'u',
                'antiname': 'u~',
                'spin': 2,
                'color': 3,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'u',
                'antitexname': '\bar u',
                'line': 'straight',
                'charge': 2. / 3.,
                'pdg_code': 2,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        u = mypartlist[-1]
        antiu = copy.copy(u)
        antiu.set('is_part', False)

        # A quark D and its antiparticle
        mypartlist.append(
            base_objects.Particle({
                'name': 'd',
                'antiname': 'd~',
                'spin': 2,
                'color': 3,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'd',
                'antitexname': '\bar d',
                'line': 'straight',
                'charge': -1. / 3.,
                'pdg_code': 1,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        d = mypartlist[-1]
        antid = copy.copy(d)
        antid.set('is_part', False)

        # A quark C and its antiparticle
        mypartlist.append(
            base_objects.Particle({
                'name': 'c',
                'antiname': 'c~',
                'spin': 2,
                'color': 3,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'u',
                'antitexname': '\bar u',
                'line': 'straight',
                'charge': 2. / 3.,
                'pdg_code': 4,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        c = mypartlist[-1]
        antic = copy.copy(c)
        antic.set('is_part', False)

        # electron/positront
        mypartlist.append(
            base_objects.Particle({
                'name': 'e-',
                'antiname': 'e+',
                'spin': 2,
                'color': 1,
                'mass': 'zero',
                'width': 'zero',
                'texname': 'u',
                'antitexname': '\bar u',
                'line': 'straight',
                'charge': 1,
                'pdg_code': 11,
                'propagating': True,
                'is_part': True,
                'self_antipart': False
            }))
        e = mypartlist[-1]
        antie = copy.copy(e)
        mu = copy.copy(e)
        antie.set('is_part', False)
        mu.set('name', 'mu-')
        mu.set('antiname', 'mu+')
        mu.set('pdg_code', 13)
        mypartlist.append(mu)
        antimu = copy.copy(mu)
        antimu.set('is_part', False)

        # A Z
        mypartlist.append(
            base_objects.Particle({
                'name': 'z',
                'antiname': 'z',
                'spin': 3,
                'color': 1,
                'mass': 'MZ',
                'width': 'WZ',
                'texname': 'Z',
                'antitexname': 'Z',
                'line': 'wavy',
                'charge': 0.,
                'pdg_code': 23,
                'propagating': True,
                'is_part': True,
                'self_antipart': True
            }))
        z = mypartlist[-1]

        # Coupling of Z to quarks

        myinterlist.append(base_objects.Interaction({
                      'id': 1,
                      'particles': base_objects.ParticleList(\
                                            [antiu, \
                                             u, \
                                             z]),
                      'color': [color.ColorString([color.T(1,0)])],
                      'lorentz':['FFV1', 'FFV2'],
                      'couplings':{(0, 0):'GUZ1', (0, 1):'GUZ2'},
                      'orders':{'QED':1}}))

        myinterlist.append(base_objects.Interaction({
                      'id': 2,
                      'particles': base_objects.ParticleList(\
                                            [antid, \
                                             d, \
                                             z]),
                      'color': [color.ColorString([color.T(1,0)])],
                      'lorentz':['FFV1', 'FFV2'],
                      'couplings':{(0, 0):'GDZ1', (0, 1):'GDZ2'},
                      'orders':{'QED':1}}))

        myinterlist.append(base_objects.Interaction({
                      'id': 3,
                      'particles': base_objects.ParticleList(\
                                            [antic, \
                                             c, \
                                             z]),
                      'color': [color.ColorString([color.T(1,0)])],
                      'lorentz':['FFV1', 'FFV2'],
                      'couplings':{(0, 0):'GUZ1', (0, 1):'GUZ2'},
                      'orders':{'QED':1}}))

        # Coupling of Z to leptons
        myinterlist.append(base_objects.Interaction({
                      'id': 4,
                      'particles': base_objects.ParticleList(\
                                            [antie, \
                                             e, \
                                             z]),
                      'color': [color.ColorString()],
                      'lorentz':['FFV1', 'FFV2'],
                      'couplings':{(0, 0):'GLZ1', (0, 1):'GLZ2'},
                      'orders':{'QED':1}}))

        # Coupling of Z to leptons
        myinterlist.append(base_objects.Interaction({
                      'id': 5,
                      'particles': base_objects.ParticleList(\
                                            [antimu, \
                                             mu, \
                                             z]),
                      'color': [color.ColorString()],
                      'lorentz':['FFV1', 'FFV2'],
                      'couplings':{(0, 0):'GLZ1', (0, 1):'GLZ2'},
                      'orders':{'QED':1}}))

        mymodel = base_objects.Model()
        mymodel.set('particles', mypartlist)
        mymodel.set('interactions', myinterlist)
        mymodel.set('name', 'sm')

        procs = [[1, -1, 23], [2, -2, 23], [4, -4, 23]]
        decays = [[23, 11, -11], [23, 13, -13]]
        coreamplitudes = diagram_generation.AmplitudeList()
        decayamplitudes = diagram_generation.AmplitudeList()
        decayprocs = base_objects.ProcessList()

        #Building the basic objects
        for proc in procs:
            # Define the multiprocess
            my_leglist = base_objects.LegList([\
                base_objects.Leg({'id': id, 'state': True}) for id in proc])

            my_leglist[0].set('state', False)
            my_leglist[1].set('state', False)

            my_process = base_objects.Process({
                'legs': my_leglist,
                'model': mymodel
            })
            my_amplitude = diagram_generation.Amplitude(my_process)
            my_amplitude.set('has_mirror_process', True)
            coreamplitudes.append(my_amplitude)

        for proc in decays:
            # Define the multiprocess
            my_leglist = base_objects.LegList([\
                base_objects.Leg({'id': id, 'state': True}) for id in proc])

            my_leglist[0].set('state', False)

            my_process = base_objects.Process({
                'legs': my_leglist,
                'model': mymodel,
                'is_decay_chain': True
            })
            my_amplitude = diagram_generation.Amplitude(my_process)
            decayamplitudes.append(my_amplitude)
            decayprocs.append(my_process)

        decays = diagram_generation.DecayChainAmplitudeList([\
                         diagram_generation.DecayChainAmplitude({\
                                            'amplitudes': decayamplitudes})])

        decay_chains = diagram_generation.DecayChainAmplitude({\
            'amplitudes': coreamplitudes,
            'decay_chains': decays})

        dc_subproc_group = group_subprocs.DecayChainSubProcessGroup.\
              group_amplitudes(\
                 diagram_generation.DecayChainAmplitudeList([decay_chains]))

        subproc_groups = \
                       dc_subproc_group.generate_helas_decay_chain_subproc_groups()

        ######################
        ##  Make the test!! ##
        ######################
        self.assertEqual(len(subproc_groups), 1)
        subproc_groups = subproc_groups.split_lepton_grouping()
        self.assertEqual(len(subproc_groups), 2)

        # check that indeed
        for group in subproc_groups:
            self.assertEqual(len(group['matrix_elements']), 2)
            has_muon = any(
                abs(l['id']) == 13 for l in group['matrix_elements'][0]
                ['processes'][0]['decay_chains'][0]['legs'])
            for me in group['matrix_elements']:
                if abs(me['processes'][0]['legs'][0]['id']) == 1:
                    self.assertEqual(len(me['processes']), 1)
                else:
                    self.assertEqual(len(me['processes']), 2)
                for proc in me['processes']:
                    for dec in proc['decay_chains']:
                        if has_muon:
                            self.assertFalse(
                                any(abs(l['id']) == 11 for l in dec['legs']))
                        else:
                            self.assertFalse(
                                any(abs(l['id']) == 13 for l in dec['legs']))

            self.assertNotEqual(group['name'], 'qq_z_z_ll')
            if has_muon:
                self.assertEqual(group['name'], 'qq_z_z_mummup')
            else:
                self.assertEqual(group['name'], 'qq_z_z_emep')
            group['name']
    def setUp(self):
        """ Setup a toy-model with gluon and down-quark only """

        # A gluon
        self.mypartlist.append(base_objects.Particle({'name':'g',
                      'antiname':'g',
                      'spin':3,
                      'color':8,
                      'mass':'zero',
                      'width':'zero',
                      'texname':'g',
                      'antitexname':'g',
                      'line':'curly',
                      'charge':0.,
                      'pdg_code':21,
                      'propagating':True,
                      'is_part':True,
                      'self_antipart':True}))

        # A quark D and its antiparticle
        self.mypartlist.append(base_objects.Particle({'name':'d',
                      'antiname':'d~',
                      'spin':2,
                      'color':3,
                      'mass':'dmass',
                      'width':'zero',
                      'texname':'d',
                      'antitexname':'\bar d',
                      'line':'straight',
                      'charge':-1. / 3.,
                      'pdg_code':1,
                      'propagating':True,
                      'is_part':True,
                      'self_antipart':False}))
        antid = copy.copy(self.mypartlist[1])
        antid.set('is_part', False)

        # 3 gluon vertex
        self.myinterlist.append(base_objects.Interaction({
                      'id': 1,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[0]] * 3),
                      'color': [],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'G'},
                      'orders':{'QCD':1}}))

        # 4 gluon vertex
        self.myinterlist.append(base_objects.Interaction({
                      'id': 2,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[0]] * 4),
                      'color': [],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'G^2'},
                      'orders':{'QCD':2}}))

        # Gluon coupling to the down-quark
        self.myinterlist.append(base_objects.Interaction({
                      'id': 3,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[1], \
                                             antid, \
                                             self.mypartlist[0]]),
                      'color': [],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQQ'},
                      'orders':{'QCD':1}}))

        self.mymodel.set('particles', self.mypartlist)
        self.mymodel.set('interactions', self.myinterlist)
        self.myproc.set('model',self.mymodel)
        
        self.myloopmodel = save_load_object.load_from_file(os.path.join(_input_file_path,\
                                                            'test_toyLoopModel.pkl'))

        box_diagram, box_struct = self.def_box()
        pent_diagram, pent_struct = self.def_pent()
       
        self.box_drawing = draw_lib.LoopFeynmanDiagram(
                                box_diagram, box_struct, self.myloopmodel)
class LoopDiagramDrawerTest(unittest.TestCase):
    """Test class for all functions related to the LoopDiagramDrawer
        diagram made by hand
    """
    
    myloopmodel = loop_base_objects.LoopModel()
    mypartlist = base_objects.ParticleList()
    myinterlist = base_objects.InteractionList()
    mymodel = base_objects.Model()
    myproc = base_objects.Process()

    def setUp(self):
        """ Setup a toy-model with gluon and down-quark only """

        # A gluon
        self.mypartlist.append(base_objects.Particle({'name':'g',
                      'antiname':'g',
                      'spin':3,
                      'color':8,
                      'mass':'zero',
                      'width':'zero',
                      'texname':'g',
                      'antitexname':'g',
                      'line':'curly',
                      'charge':0.,
                      'pdg_code':21,
                      'propagating':True,
                      'is_part':True,
                      'self_antipart':True}))

        # A quark D and its antiparticle
        self.mypartlist.append(base_objects.Particle({'name':'d',
                      'antiname':'d~',
                      'spin':2,
                      'color':3,
                      'mass':'dmass',
                      'width':'zero',
                      'texname':'d',
                      'antitexname':'\bar d',
                      'line':'straight',
                      'charge':-1. / 3.,
                      'pdg_code':1,
                      'propagating':True,
                      'is_part':True,
                      'self_antipart':False}))
        antid = copy.copy(self.mypartlist[1])
        antid.set('is_part', False)

        # 3 gluon vertex
        self.myinterlist.append(base_objects.Interaction({
                      'id': 1,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[0]] * 3),
                      'color': [],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'G'},
                      'orders':{'QCD':1}}))

        # 4 gluon vertex
        self.myinterlist.append(base_objects.Interaction({
                      'id': 2,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[0]] * 4),
                      'color': [],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'G^2'},
                      'orders':{'QCD':2}}))

        # Gluon coupling to the down-quark
        self.myinterlist.append(base_objects.Interaction({
                      'id': 3,
                      'particles': base_objects.ParticleList(\
                                            [self.mypartlist[1], \
                                             antid, \
                                             self.mypartlist[0]]),
                      'color': [],
                      'lorentz':['L1'],
                      'couplings':{(0, 0):'GQQ'},
                      'orders':{'QCD':1}}))

        self.mymodel.set('particles', self.mypartlist)
        self.mymodel.set('interactions', self.myinterlist)
        self.myproc.set('model',self.mymodel)
        
        self.myloopmodel = save_load_object.load_from_file(os.path.join(_input_file_path,\
                                                            'test_toyLoopModel.pkl'))

        box_diagram, box_struct = self.def_box()
        pent_diagram, pent_struct = self.def_pent()
       
        self.box_drawing = draw_lib.LoopFeynmanDiagram(
                                box_diagram, box_struct, self.myloopmodel)

    def test_fuse_line(self):
        """ check that we fuse line correctly """
        
        self.box_drawing.load_diagram()
        #avoid that element are erase from memory
        line1 = self.box_drawing.lineList[0]
        line2 = self.box_drawing.lineList[1]
        vertex1 = line1.begin
        vertex2 = line1.end
        vertex3 = line2.begin
        vertex4 = line2.end
        
        # fuse line1 and line2
        self.box_drawing.fuse_line(line1, line2)
        
        # check that all link to line1 are ok
        self.assertEqual(line1.begin, vertex1)
        self.assertEqual(line1.end, vertex3)
        self.assertTrue(line1 in vertex1.lines)
        self.assertTrue(line1 in vertex3.lines)
        #self.assertTrue(vertex1 in self.box_drawing.vertexList)
        #self.assertTrue(vertex4 in self.box_drawing.vertexList)

        
        #check that all info to line2 are deleted
        self.assertFalse(line2 in self.box_drawing.lineList)
        self.assertFalse(line2 in vertex1.lines)
        self.assertFalse(line2 in vertex3.lines)
        self.assertFalse(vertex2 in self.box_drawing.vertexList)
        self.assertFalse(vertex3 in self.box_drawing.vertexList)
        

    def def_box(self):
        """ Test the drawing of a simple loop box """
        
        myleglist = base_objects.LegList([base_objects.Leg({'id':21,
                                              'number':num, 'state':True,
                                              'loop_line':False}) \
                                              for num in range(1, 5)])
        myleglist.append(base_objects.Leg({'id':1,'number':5,'loop_line':True}))
        myleglist.append(base_objects.Leg({'id':-1,'number':6,'loop_line':True}))                         
        l1=myleglist[0]
        l1.set('state',False)
        l2=myleglist[1]
        l2.set('state',False)
        l3=myleglist[2]
        l4=myleglist[3]
        l5=myleglist[4]
        l6=myleglist[5]

        
        # One way of constructing this diagram, with a three-point amplitude
        l15 = base_objects.Leg({'id':1,'number':1,'loop_line':True, 'state':False})
        l12 = base_objects.Leg({'id':1,'number':1,'loop_line':True})
        l13 = base_objects.Leg({'id':1,'number':1,'loop_line':True}) 
        lfake = base_objects.Leg({'id':1,'number':1,'loop_line':True})         

        vx15 = base_objects.Vertex({'legs':base_objects.LegList([l1, l5, l15]), 'id': 3})
        vx12 = base_objects.Vertex({'legs':base_objects.LegList([l15, l2, l12]), 'id': 3})
        vx13 = base_objects.Vertex({'legs':base_objects.LegList([l12, l3, l13]), 'id': 3})
        vx164 = base_objects.Vertex({'legs':base_objects.LegList([l13, l6, l4]), 'id': 3})
        fakevx = base_objects.Vertex({'legs':base_objects.LegList([l13, lfake]), 'id': 0})
        ctvx = base_objects.Vertex({'legs':base_objects.LegList([l1, l2, l3, l4]), 'id': 666})

        myVertexList1=base_objects.VertexList([vx15,vx12,vx13,vx164])
        myCTVertexList=base_objects.VertexList([ctvx,])
        myPentaDiag1=loop_base_objects.LoopDiagram({'vertices':myVertexList1,'type':1,\
                                                    'CT_vertices':myCTVertexList})
        
        return myPentaDiag1, []

    def def_pent(self):       
        """ Test the gg>gggg d*dx* tagging of a quark pentagon which is tagged"""

        # Five gluon legs with two initial states
        myleglist = base_objects.LegList([base_objects.Leg({'id':21,
                                              'number':num,
                                              'loop_line':False}) \
                                              for num in range(1, 7)])
        myleglist.append(base_objects.Leg({'id':1,'number':7,'loop_line':True}))
        myleglist.append(base_objects.Leg({'id':-1,'number':8,'loop_line':True}))                         
        l1=myleglist[0]
        l2=myleglist[1]
        l3=myleglist[2]
        l4=myleglist[3]
        l5=myleglist[4]
        l6=myleglist[5]
        l7=myleglist[6]
        l8=myleglist[7]

        # One way of constructing this diagram, with a three-point amplitude
        l17 = base_objects.Leg({'id':1,'number':1,'loop_line':True})
        l12 = base_objects.Leg({'id':1,'number':1,'loop_line':True})
        l68 = base_objects.Leg({'id':-1,'number':6,'loop_line':True}) 
        l56 = base_objects.Leg({'id':-1,'number':5,'loop_line':True})
        l34 = base_objects.Leg({'id':21,'number':3,'loop_line':False})

        self.myproc.set('legs',myleglist)

        vx17 = base_objects.Vertex({'legs':base_objects.LegList([l1, l7, l17]), 'id': 3})
        vx12 = base_objects.Vertex({'legs':base_objects.LegList([l17, l2, l12]), 'id': 3})
        vx68 = base_objects.Vertex({'legs':base_objects.LegList([l6, l8, l68]), 'id': 3})
        vx56 = base_objects.Vertex({'legs':base_objects.LegList([l5, l68, l56]), 'id': 3})
        vx34 = base_objects.Vertex({'legs':base_objects.LegList([l3, l4, l34]), 'id': 1})
        vx135 = base_objects.Vertex({'legs':base_objects.LegList([l12, l56, l34]), 'id': 3})

        myVertexList1=base_objects.VertexList([vx17,vx12,vx68,vx56,vx34,vx135])

        myPentaDiag1=loop_base_objects.LoopDiagram({'vertices':myVertexList1,'type':1})

        myStructRep=loop_base_objects.FDStructureList()
        
        myPentaDiag1.tag(myStructRep,7,8,self.myproc)
        
        return myPentaDiag1,myStructRep
        # test the drawing of myPentaDiag with its loop vertices and those in the 
        # structures of myStructRep
        

    def def_diagrams_epemddx(self):
        """ Test the drawing of diagrams from the loop process e+e- > dd~ """
    
        myleglist = base_objects.LegList()
        myleglist.append(base_objects.Leg({'id':-11,
                                         'state':False}))
        myleglist.append(base_objects.Leg({'id':11,
                                         'state':False}))
        myleglist.append(base_objects.Leg({'id':1,
                                         'state':True}))
        myleglist.append(base_objects.Leg({'id':-1,
                                         'state':True}))
        
        myproc = base_objects.Process({'legs':myleglist,
                                        'model':self.myloopmodel,
                                        'orders':{},
                                        'perturbation_couplings':['QCD',],
                                        'squared_orders':{}})
    
        myloopamplitude = loop_diagram_generation.LoopAmplitude()
        myloopamplitude.set('process', myproc)
        myloopamplitude.generate_diagrams()
        
        # Now the drawing test on myloopamplitude['loop_diagrams']
        return myloopamplitude['loop_diagrams']