patchMEInt.py

#!/usr/bin/env python

"""Patches computation of squared ME to keep interference terms only.

This script patches Fortran code generated for each subprocess.  When
summing over helicities, it computes the squared matrix element as is,
then inverts the sign of the given coupling, recomputes the squared
matrix element, and finally returns the half-difference of the two. This
way only terms with odd degrees in the coupling are kept.

Definitions of couplings can be found in file couplings.py in the
directory with the model.

This trick was proposed by Sebastien Wertz.
"""

from __future__ import print_function
import argparse
import os
import re
import shutil
import tempfile


if __name__ == '__main__':
    
    argParser = argparse.ArgumentParser(
        epilog=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    argParser.add_argument('dir', help='Directory generated by MadGraph for a given process')
    argParser.add_argument(
        'coupling', type=int,
        help='Index of the coupling whose sign will be inversed'
    )
    args = argParser.parse_args()
    
    couplingName = 'GC_{}'.format(args.coupling)
    
    
    subProcDir = os.path.join(args.dir, 'SubProcesses')
    
    if not os.path.exists(subProcDir):
        raise RuntimeError(
            'Directory "{}" does not contain required subdirectory "SubProcesses".'.format(
                args.dir
            )
        )
    
    regexInclude = re.compile(r'(\s*)INCLUDE\s+\'maxamps\.inc\'\s*')
    regexME = re.compile(r'(\s*)T=(MATRIX1\(.+\))\s*')
    
    
    # Loop over all subprocesses
    for entry in os.listdir(subProcDir):
        if not os.path.isdir(os.path.join(subProcDir, entry)):
            continue
        
        srcFileName = os.path.join(subProcDir, entry, 'matrix1.f')
        print('Patching file "{}".'.format(srcFileName))
        
        
        srcFile = open(srcFileName)
        outFile = tempfile.NamedTemporaryFile(mode='w', delete=False)
        outFileName = outFile.name
        
        
        # Copy content of the source file with up to the end of the
        # include block.  Add inclusion of file with couplings.
        while True:
            line = srcFile.readline()
            
            if not line:
                raise RuntimeError(
                    'Failed to locate the end of the include block in file "{}".'.format(
                        srcFileName
                    )
                )
            
            outFile.write(line)
            match = regexInclude.match(line)
            
            if match:
                outFile.write(match.group(1) + 'INCLUDE \'coupl.inc\'\n')
                break
        
        
        # Copy the remaining content of the source file.  When the
        # computation of the squared ME is encountered, apply the trick
        # with flipping of the sign of selected coupling.
        nMEPatches = 0
        
        while True:
            line = srcFile.readline()
            
            if not line:
                break
            
            outFile.write(line)
            match = regexME.match(line)
            
            if match:
                spacing = match.group(1)
                meCall = match.group(2)
                
                outFile.write('{s}{c}=-{c}\n'.format(s=spacing, c=couplingName))
                outFile.write('{}T=5D-1*(T-{})\n'.format(spacing, meCall))
                outFile.write('{s}{c}=-{c}\n'.format(s=spacing, c=couplingName))
                
                nMEPatches += 1
        
        if nMEPatches != 2:
            raise RuntimeError(
                'In file "{}" found {} places in which the squared matrix element is evaluated '
                'while 2 places expected.'.format(srcFileName, nMEPatches)
            )
        
        
        srcFile.close()
        outFile.close()
        
        shutil.move(srcFileName, srcFileName + '.orig')
        shutil.move(outFileName, srcFileName)