Source code for hepi.madgraph.run

"""Runs MadGraph"""
import shutil
from typing import List
import subprocess
from string import Template
import warnings
from hepi.input import Order
from hepi.run import RunParam, Runner
import numpy as np
import pkgutil
from .. import Input, Result, LD2DL, get_output_dir, get_input_dir
import re
import os.path
from pathlib import Path
from .result import MadgraphResult, is_valid, parse_single
import time
import difflib
from smpl.parallel import par
import hashlib

class MadGraphRunParams(RunParam):    
    """
    Parameters for MadGraph.
    """
    def __init__(self, dic,  skip=False,madstr=True):
        super().__init__(skip)
        self.dic = dic
        # self.skip = skip
        self.madstr = madstr

class MadGraphRunner(Runner):
    def _check_input(self,p: Input)-> bool:
        """Checks input parameter for compatibility with Prospino"""
        if p.order != Order.LO and p.order != Order.NLO:
            warnings.warn("MadGraph does only support LO/NLO.")
            return False
        return True
    def _is_valid(self, file: str, p: Input, d) -> bool:
        return is_valid(file,p,d)
    def _parse_file(self, file: str) -> Result:
        return parse_single(file)
    def _run(self,rps: List[RunParam], wait=True,parallel=True,sleep=0,**kwargs):
        # TODO clean up on exit emergency
        global madgraph_path
        skipall = True
        for rp in rps:
            if not rp.skip:
                skipall = False
        template = ""
        if not skipall:
            lo = "&& nice -n 5 {dir}/bin/calculate_xsect LO -f >> {out} " if rps[0].dic["order"] == Order.NLO else ""
            template =  \
                'rm -rf {dir} && cp -r ' + rps[0].dic["bdir"] + \
                ' {dir}  && cp {slha} {dir}/Cards/param_card.dat && cp {run} {dir}/Cards/run_card.dat && echo "nb_core = 1" >> {dir}/Cards/amcatnlo_configuration.txt ' + lo+ '&& nice -n 5 {dir}/bin/calculate_xsect -f >> {out}  && rm -rf {dir}'
            print(rps[0].dic["out"])
        if not rps[0].skip:
            mgcom = 'bin/mg5_aMC'
            if rps[0].madstr:
                mgcom = 'bin/mg5_aMC --mode="MadSTR"'
            com = madgraph_path + mgcom + \
                ' --file {in} >> {out} && cp {slha} {bdir}/Cards/param_card.dat && cp {run} {bdir}/Cards/run_card.dat && sed -i \'s/.*= req_acc_FO/ 1 = req_acc_FO/g\' {bdir}/Cards/run_card.dat && echo "automatic_html_opening = False" >> {bdir}/Cards/amcatnlo_configuration.txt && nice -n 5 {bdir}/bin/calculate_xsect -f'
            pp = subprocess.Popen(com.format(**rps[0].dic), shell=True)
            pp.wait()
        # Run commands in parallel
        processes = []

        for rp in rps:
            if not rp.skip:
                command = template.format(**rp.dic)
                # print(command)
                process = subprocess.Popen(command, shell=True)
                processes.append(process)
                if not parallel:
                    process.wait()
                # Forced delay to prevent overloading clusters when registering jobs
                time.sleep(sleep)
        if wait:
            # Collect statuses
            output = [p.wait() for p in processes]
            return output
        return []

    def _prepare(self, p: Input, **kwargs) -> RunParam:
        rp = super()._prepare(p,**kwargs)
        name = rp.name
        if not rp.skip:
            d= p.__dict__
            d["dir"] = get_output_dir() + name + ".dir"
            d["bdir"] = get_output_dir() + name + ".bdir"

            infile = ""
            if p.order is Order.LO:
                infile ="lo.mg"
            elif p.order is Order.NLO:
                infile ="nlo.mg"
            else:
                raise ValueError("Madgraph only supported for LO/NLO.")

            data = pkgutil.get_data(__name__, infile).decode( 'utf-8')
            src = Template(data)
            result = src.substitute(d)
            open(get_output_dir() + name + ".mg", "w").write(result)
            if not rp.skip:
                open(get_output_dir() + name + ".out", "w").write(result + "\n\n")

            if p.order == Order.LO:
                mgfile = "run_card_no_madstr.dat"
            elif p.order == Order.NLO:
                mgfile = "run_card_with_madstr.dat"
            else:
                raise ValueError("Order must be one of LO/NLO in MadGraph.")


            if "madstr" in kwargs and not kwargs["madstr"]:
                mgfile = "run_card_no_madstr.dat"
                rp.madstr = False
            else:
                rp.madstr = True
            data = pkgutil.get_data(__name__, mgfile).decode( 'utf-8')

            src = Template(data)
            result = src.substitute(d)
            open(get_output_dir() + name + ".dat", "w").write(result)
            if not rp.skip:
                open(get_output_dir() + name + ".out", "a").write(result + "\n\n")

            sname = d['slha']
            with open(get_output_dir() + sname, 'r') as f:
                #src = Template(f.read())
                #result = src.substitute(d)
                #open(get_input_dir() + sname + ".in", "w").write(result)
                if not rp.skip:
                    open(get_output_dir() + name + ".out",
                         "a").write(f.read() + "\n\n")
            rp.dic = {'in': get_output_dir()+name + ".mg",
                              'dir': d["dir"],
                              'bdir': d["bdir"],
                              'run': get_output_dir() + name + ".dat",
                              'slha': get_output_dir() + sname,
                              'out': get_output_dir()+name + ".out",
                              'order' : p.order}
            rp.out_file = rp.dic['out']
        return rp


# Legacy 
default_madgraph_runner = MadGraphRunner("/opt/MG5_aMC_v2_7_0/")
"""Default MadGraph Runner to provide backward compatibility"""
run = default_madgraph_runner.run
set_path = default_madgraph_runner.set_path
get_path = default_madgraph_runner.get_path

# print(sp)
# names = queue(sp)
# run(names)