run_model

#!/usr/bin/env python3

import os, sys, fileinput, subprocess, shlex, shutil, gzip
from pathlib import Path
from contextlib import nullcontext
from copy import deepcopy
from magiconfig import ArgumentParser, MagiConfig, MagiConfigOptions, ArgumentDefaultsRawHelpFormatter
from svjHelper import svjHelper, extHelper
import Histogram

sys.path.append(os.getcwd()+"/configs")

def separate_args(args,groups):
    separated = {group.title : MagiConfig(**{a.dest:getattr(args,a.dest,None) for a in group._group_actions}) for group in groups}
    return separated

def run_cmd(cmd,log,verbose=False,shell=False):
    if verbose:
        print(cmd)
    with open(log,'w') as logfile:
        if not shell:
            cmd = shlex.split(cmd)
        subprocess.check_call(cmd,stdout=logfile,stderr=logfile,shell=shell)

if __name__=="__main__":
    _parser_common = ArgumentParser(add_help=False)

    allowed_steps = ['pythia','delphes','hist']

    common = _parser_common.add_argument_group("common")
    common.add_argument("--steps", type=str, nargs='*', default=[], choices=allowed_steps+['all'], help="run these steps")
    common.add_argument("--events", type=int, default=0, help="generate this many events in Pythia")
    print_group = common.add_mutually_exclusive_group()
    print_group.add_argument("--verbose", default=False, action="store_true", help="increase verbosity of printouts")
    print_group.add_argument("--quiet", default=False, action="store_true", help="suppress all printouts")
    common.add_argument("--dir", type=str, default="models", help="output directory")
    common.add_argument("--pythia", type=str, nargs='*', default=["cards/CMS_Common.txt","cards/CMS_Tune_CP5.txt"], help="additional settings for Pythia")
    common.add_argument("--delphes", type=str, default="cards/delphes_card_CMS.tcl", help="template card for Delphes")

    parser = ArgumentParser(
        formatter_class=ArgumentDefaultsRawHelpFormatter
    )
    subparsers = parser.add_subparsers(dest="model_type")
    helpers = {
        "helper": svjHelper,
        "external": extHelper,
    }
    parsers = {}
    model_groups = {}
    for model_type in helpers.keys():
        parsers[model_type] = subparsers.add_parser(model_type, config_options=MagiConfigOptions(), parents=[_parser_common])
        model_groups[model_type] = parsers[model_type].add_argument_group("model")
        helpers[model_type].add_arguments(model_groups[model_type])
        parsers[model_type].set_defaults(model_type=model_type)

    args_orig = parser.parse_args()
    model_type = args_orig.model_type
    helper_class = helpers[model_type]
    args = separate_args(args_orig,[common,model_groups[model_type]])

    if 'all' in args["common"].steps:
        args["common"].steps = allowed_steps
    if 'pythia' in args["common"].steps and not args["common"].events>0:
        raise RuntimeError("No events requested to generate")
    if args["common"].pythia==['']: # filled by default, turn this into empty
        args["common"].pythia = []

    # absolute paths for files relative to current dir
    args["common"].pythia = [Path(txt).absolute() for txt in args["common"].pythia]
    args["common"].delphes = Path(args["common"].delphes).absolute()

    helper = helper_class(args["model"])
    name = helper.name()
    outdir = os.path.join(args["common"].dir,name)
    os.makedirs(outdir, exist_ok=True)

    # change to output dir
    os.chdir(outdir)
    if args["common"].verbose: print(f'cd {outdir}')

    config_fname = "config.py"
    parser.write_config(args_orig,config_fname)
    if args["common"].verbose: print(f'wrote {config_fname}')

    pythia_lines = helper.getPythiaSettings()
    if 'pythia' in args["common"].steps:
        pythia_lines.append('Main:numberOfEvents = {}'.format(args["common"].events))
    pythia_fname = "pythia_card.txt"
    with open(pythia_fname,'w') as file, (fileinput.input(args["common"].pythia) if len(args["common"].pythia)>0 else nullcontext()) as inputs:
        file.write('\n'.join(pythia_lines))
        file.write('\n')
        if inputs is not None:
            for line in inputs:
                file.write(line)
    if args["common"].verbose: print(f'wrote {pythia_fname}')

    delphes_lines = helper.getDelphesSettings(args["common"].delphes)
    delphes_fname = "delphes_card.txt"
    with open(delphes_fname,'w') as file:
        file.write(delphes_lines)
    if args["common"].verbose: print(f'wrote {delphes_fname}')

    # intermediate output filenames
    hepmc_fname = "events.hepmc"
    hepmc_fname_gz = hepmc_fname+".gz"
    root_fname = "events.root"

    if 'pythia' in args["common"].steps:
        # step 1: pythia
        log_fname = "log_pythia8.log"
        pythia_exe = os.path.expandvars("$PYTHIA8/examples/main42")
        pythia_cmd = f'{pythia_exe} {pythia_fname} {hepmc_fname}'
        if not args["common"].quiet: print(f'Running Pythia ({log_fname})')
        run_cmd(pythia_cmd, log_fname, args["common"].verbose)

        # step 1.5: compress events
        if not args["common"].quiet: print("Compressing Pythia output")
        with open(hepmc_fname,'rb') as ifile, gzip.open(hepmc_fname_gz,'wb') as ofile:
            shutil.copyfileobj(ifile, ofile)
        os.remove(hepmc_fname)
        if args["common"].verbose: print(f'wrote {hepmc_fname_gz}')

    if 'delphes' in args["common"].steps:
        # check for input
        if not os.path.exists(hepmc_fname_gz):
            raise RuntimeError(f'Could not find Delphes input {hepmc_fname_gz}')

        # step 2: delphes
        log_fname = "log_delphes.log"
        delphes_exe = "DelphesHepMC2"
        
        if os.path.exists(root_fname):
            if args["common"].verbose: print(f'removing old {root_fname}')
            os.remove(root_fname)
        delphes_cmd = f'gunzip -c {hepmc_fname_gz} | {delphes_exe} {delphes_fname} {root_fname}'
        if not args["common"].quiet: print(f'Running Delphes ({log_fname})')
        run_cmd(delphes_cmd, log_fname, args["common"].verbose, shell=True)
        if args["common"].verbose: print(f'wrote {root_fname}')
        
    if 'hist' in args["common"].steps:
        
        Histogram.histogram(root_fname, helper)
        if args["common"].verbose: print(f'wrote histograms of {root_fname}')