← Back to branch summary

~madteam/mg5amcnlo/series2.0

« back to all changes in this revision

Viewing changes to madgraph/iolibs/helas_call_writers.py

  • Committer: olivier-mattelaer
  • Date: 2020-08-21 09:16:56 UTC
  • mfrom: (284.2.24 python3)
  • Revision ID: olivier-mattelaer-20200821091656-iizux2mj94tkssuo
pass to 2.8.0 (and move to python3 branch)

Show diffs side-by-side

added added

removed removed

Lines of Context:
madgraph/iolibs/helas_call_writers.py
13
13
#
14
14
################################################################################
15
15
"""Classes for writing Helas calls. HelasCallWriter is the base class."""
16
 
 
 
16
from __future__ import absolute_import
 
17
 
 
18
 
 
19
import re
17
20
import madgraph.core.base_objects as base_objects
18
21
import madgraph.core.helas_objects as helas_objects
19
22
import madgraph.loop.loop_helas_objects as loop_helas_objects
22
25
import aloha
23
26
from madgraph import MadGraph5Error
24
27
import madgraph.various.misc as misc
 
28
from six.moves import range
25
29
 
26
30
class HelasWriterError(Exception):
27
31
    """Class for the error of this module """
62
66
        self['model'] = base_objects.Model()
63
67
        self['wavefunctions'] = {}
64
68
        self['amplitudes'] = {}
 
69
        self.width_tchannel_set_tozero = False
65
70
 
66
71
    def filter(self, name, value):
67
72
        """Filter for model property values"""
68
73
 
69
74
        if name == 'model':
70
75
            if not isinstance(value, base_objects.Model):
71
 
                raise self.PhysicsObjectError, \
72
 
                    "Object of type %s is not a model" % type(value)
 
76
                raise self.PhysicsObjectError("Object of type %s is not a model" % type(value))
73
77
 
74
78
        if name == 'wavefunctions':
75
79
            # Should be a dictionary of functions returning strings, 
76
80
            # with keys (spins, flow state)
77
81
            if not isinstance(value, dict):
78
 
                raise self.PhysicsObjectError, \
79
 
                        "%s is not a valid dictionary for wavefunction" % \
80
 
                                                                str(value)
 
82
                raise self.PhysicsObjectError("%s is not a valid dictionary for wavefunction" % \
 
83
                                                                str(value))
81
84
 
82
85
            for key in value.keys():
83
86
                self.add_wavefunction(key, value[key])
86
89
            # Should be a dictionary of functions returning strings, 
87
90
            # with keys (spins, flow state)
88
91
            if not isinstance(value, dict):
89
 
                raise self.PhysicsObjectError, \
90
 
                        "%s is not a valid dictionary for amplitude" % \
91
 
                                                                str(value)
 
92
                raise self.PhysicsObjectError("%s is not a valid dictionary for amplitude" % \
 
93
                                                                str(value))
92
94
 
93
95
            for key in value.keys():
94
96
                self.add_amplitude(key, value[key])
231
233
 
232
234
        res = []
233
235
        for diagram in matrix_element.get('diagrams'):
 
236
            
 
237
            
234
238
            res.extend([ self.get_wavefunction_call(wf) for \
235
239
                         wf in diagram.get('wavefunctions') ])
236
240
            res.append("# Amplitude(s) for diagram number %d" % \
274
278
        corresponding to the key"""
275
279
 
276
280
        try:
 
281
            #misc.sprint(wavefunction['number_external'])
277
282
            call = self["wavefunctions"][wavefunction.get_call_key()](\
278
 
                wavefunction)
279
 
            return call
280
 
        except KeyError:
 
283
                                                                   wavefunction)
 
284
        except KeyError as error:
281
285
            return ""
 
286
        
 
287
        if  self.options['zerowidth_tchannel'] and wavefunction.is_t_channel():
 
288
            call, n = re.subn(',\s*fk_(?!ZERO)\w*\s*,', ', ZERO,', str(call), flags=re.I)
 
289
            if n:
 
290
                self.width_tchannel_set_tozero = True
 
291
        return call
 
292
        
282
293
 
283
294
    def get_amplitude_call(self, amplitude):
284
295
        """Return the function for writing the amplitude
286
297
 
287
298
        try:
288
299
            call = self["amplitudes"][amplitude.get_call_key()]
289
 
        except KeyError, error:
 
300
        except KeyError as error:
290
301
            return ""
291
302
        else:
292
303
            return call(amplitude)
324
335
 
325
336
    # Customized constructor
326
337
 
327
 
    def __init__(self, argument={}):
 
338
    def __init__(self, argument={}, options={}):
328
339
        """Allow generating a HelasCallWriter from a Model
329
340
        """
330
341
 
 
342
        default_options = {'zerowidth_tchannel': True}
 
343
        
 
344
        self.options = dict(default_options)
 
345
        self.options.update(options)
 
346
        
331
347
        if isinstance(argument, base_objects.Model):
332
348
            super(HelasCallWriter, self).__init__()
333
349
            self.set('model', argument)
594
610
        # If function not already existing, try to generate it.
595
611
 
596
612
        if len(wavefunction.get('mothers')) > 3:
597
 
            raise self.PhysicsObjectError, \
598
 
                  """Automatic generation of Fortran wavefunctions not
599
 
                  implemented for > 3 mothers"""
 
613
            raise self.PhysicsObjectError("""Automatic generation of Fortran wavefunctions not
 
614
                  implemented for > 3 mothers""")
600
615
 
601
616
        self.generate_helas_call(wavefunction)
602
617
        return super(FortranHelasCallWriter, self).get_wavefunction_call(\
615
630
        # If function not already existing, try to generate it.
616
631
 
617
632
        if len(amplitude.get('mothers')) > 4:
618
 
            raise self.PhysicsObjectError, \
619
 
                  """Automatic generation of Fortran amplitudes not
620
 
                  implemented for > 4 mothers"""
 
633
            raise self.PhysicsObjectError("""Automatic generation of Fortran amplitudes not
 
634
                  implemented for > 4 mothers""")
621
635
 
622
636
        self.generate_helas_call(amplitude)
623
637
        return super(FortranHelasCallWriter, self).get_amplitude_call(amplitude)
645
659
        
646
660
        if not isinstance(argument, helas_objects.HelasWavefunction) and \
647
661
           not isinstance(argument, helas_objects.HelasAmplitude):
648
 
            raise self.PhysicsObjectError, \
649
 
                  "get_helas_call must be called with wavefunction or amplitude"
 
662
            raise self.PhysicsObjectError("get_helas_call must be called with wavefunction or amplitude")
650
663
 
651
664
        call = "CALL "
652
665
 
837
850
        if isinstance(arg, helas_objects.HelasWavefunction):
838
851
            return "".join(sorted([HelasCallWriter.mother_dict[\
839
852
            wf.get_spin_state_number()] for wf in arg.get('mothers')],
840
 
                          lambda l1, l2: \
841
 
                          FortranHelasCallWriter.sort_wf[l2] - \
842
 
                          FortranHelasCallWriter.sort_wf[l1]))
 
853
                                  key= lambda l: FortranHelasCallWriter.sort_wf[l],
 
854
                                  reverse=True))
843
855
 
844
856
        if isinstance(arg, helas_objects.HelasAmplitude):
845
857
            return "".join(sorted([HelasCallWriter.mother_dict[\
846
858
            wf.get_spin_state_number()] for wf in arg.get('mothers')],
847
 
                          lambda l1, l2: \
848
 
                          FortranHelasCallWriter.sort_amp[l2] - \
849
 
                          FortranHelasCallWriter.sort_amp[l1]))
 
859
                                  key= lambda l: FortranHelasCallWriter.sort_amp[l],
 
860
                                  reverse=True))
850
861
 
851
862
    @staticmethod
852
863
    def sorted_mothers(arg):
1003
1014
 
1004
1015
    mp_prefix = check_param_card.ParamCard.mp_prefix
1005
1016
 
1006
 
    def __init__(self, argument={}, hel_sum = False):
 
1017
    def __init__(self, argument={}, hel_sum = False, options={}):
1007
1018
        """Allow generating a HelasCallWriter from a Model.The hel_sum argument
1008
1019
        specifies if amplitude and wavefunctions must be stored specifying the
1009
1020
        helicity, i.e. W(1,i) vs W(1,i,H).
1010
1021
        """
1011
1022
        self.hel_sum = hel_sum
1012
 
        super(FortranUFOHelasCallWriter, self).__init__(argument)
 
1023
        super(FortranUFOHelasCallWriter, self).__init__(argument, options=options)
1013
1024
 
1014
1025
    def format_helas_object(self, prefix, number):
1015
1026
        """ Returns the string for accessing the wavefunction with number in
1085
1096
 
1086
1097
        if not isinstance(argument, helas_objects.HelasWavefunction) and \
1087
1098
           not isinstance(argument, helas_objects.HelasAmplitude):
1088
 
            raise self.PhysicsObjectError, \
1089
 
                  "generate_helas_call must be called with wavefunction or amplitude"
 
1099
            raise self.PhysicsObjectError("generate_helas_call must be called with wavefunction or amplitude")
1090
1100
        
1091
1101
        call = "CALL "
1092
1102
 
1256
1266
                elif lcutpart.get('spin')==2 or lcutpart.get('spin')==3:
1257
1267
                    res.append("DO I=1,4")
1258
1268
                else:
1259
 
                    raise self.PhysicsObjectError, \
1260
 
                  "The L-cut particle type is not supported"
 
1269
                    raise self.PhysicsObjectError("The L-cut particle type is not supported")
1261
1270
                # Temporarily relabel the 'me_id' attribute of the external wfs
1262
1271
                # in this wavefunction's mothers so to have them matching the
1263
1272
                # convention in the loop helas calls.
1496
1505
            
1497
1506
        if isinstance(argument, helas_objects.HelasAmplitude) and \
1498
1507
                                                  argument.get('type')=='loop':
1499
 
           raise MadGraph5Error, 'There should not be any helas call '+\
1500
 
                                'associated with helas amplitudes of type loop.'
 
1508
           raise MadGraph5Error('There should not be any helas call '+\
 
1509
                                'associated with helas amplitudes of type loop.')
1501
1510
 
1502
1511
        # Check if we need to append a charge conjugation flag
1503
1512
        l = [str(l) for l in argument.get('lorentz')]
1612
1621
 
1613
1622
        if not isinstance(argument, helas_objects.HelasWavefunction) and \
1614
1623
           not isinstance(argument, helas_objects.HelasAmplitude):
1615
 
            raise self.PhysicsObjectError, \
1616
 
                  "get_helas_call must be called with wavefunction or amplitude"
 
1624
            raise self.PhysicsObjectError("get_helas_call must be called with wavefunction or amplitude")
1617
1625
        
1618
1626
        call = ""
1619
1627
 
1742
1750
                                   and wf.get('spin') == 3 \
1743
1751
                                   and wf.get('mass').lower() == 'zero']
1744
1752
                if not gauge_check_wfs:
1745
 
                    raise HelasWriterError, \
1746
 
                          'no massless spin one particle for gauge check'
 
1753
                    raise HelasWriterError('no massless spin one particle for gauge check')
1747
1754
                gauge_check_wf = wfs.pop(wfs.index(gauge_check_wfs[0]))
1748
1755
                res.append(self.generate_helas_call(gauge_check_wf, True)(\
1749
1756
                                                    gauge_check_wf))
1764
1771
 
1765
1772
        if not isinstance(argument, helas_objects.HelasWavefunction) and \
1766
1773
           not isinstance(argument, helas_objects.HelasAmplitude):
1767
 
            raise self.PhysicsObjectError, \
1768
 
                  "get_helas_call must be called with wavefunction or amplitude"
 
1774
            raise self.PhysicsObjectError("get_helas_call must be called with wavefunction or amplitude")
1769
1775
        
1770
1776
        call_function = None
1771
1777