← Back to branch summary

~madteam/mg5amcnlo/series2.0

« back to all changes in this revision

Viewing changes to madgraph/core/diagram_generation.py

mwrge with 2.7.3

Show diffs side-by-side

added added

removed removed

Lines of Context:
madgraph/core/diagram_generation.py
20
20
"""
21
21
 
22
22
from __future__ import absolute_import
 
23
from six.moves import filter
23
24
#force filter to be a generator # like in py3
24
 
try:
25
 
    from  itertools import ifilter as filter
26
 
except:
27
 
    pass
 
25
 
28
26
 
29
27
import array
30
28
import copy
813
811
                    nexttolastvertex = copy.copy(vertices.pop())
814
812
                    legs = copy.copy(nexttolastvertex.get('legs'))
815
813
                    ntlnumber = legs[-1].get('number')
816
 
                    lastleg = next(filter(lambda leg: leg.get('number') != ntlnumber,
817
 
                                     lastvx.get('legs')))
 
814
                    lastleg = [leg for leg in lastvx.get('legs') if leg.get('number') != ntlnumber][0]
818
815
                    # Reset onshell in case we have forbidden s-channels
819
816
                    if lastleg.get('onshell') == False:
820
817
                        lastleg.set('onshell', None)
901
898
                                             fcts=['remove_diag'])
902
899
        else:
903
900
            #example and simple tests
904
 
            def remove_diag(diag):
 
901
            def remove_diag(diag, model=None):
905
902
                for vertex in diag['vertices']: #last 
906
903
                    if vertex['id'] == 0: #special final vertex
907
904
                        continue 
912
909
 
913
910
        res = diag_list.__class__()                
914
911
        nb_removed = 0 
 
912
        model = self['process']['model'] 
915
913
        for diag in diag_list:
916
 
            if remove_diag(diag):
 
914
            if remove_diag(diag, model):
917
915
                nb_removed +=1
918
916
            else:
919
917
                res.append(diag)