MadGraphConfig.py

Go to the documentation of this file.

# Copyright (C) 2002-2026 CERN for the benefit of the ATLAS collaboration
 
from AthenaConfiguration.ComponentAccumulator import ComponentAccumulator
from GeneratorConfig.Sequences import EvgenSequence, EvgenSequenceFactory
from GeneratorConfig.GeneratorInfoSvcConfig import GeneratorInfoSvcCfg
 
import os
 
from MadGraphControl.MadGraphPDFSettings import MadGraphPDFSets, get_pdf_set
 
 
def _get_nevents(flags, safety):
    """Helper function to determing number of events to be generated
    in MadGraph, based on MaxEvents or nEventsPerJob and a user-provided 
    safety factor (the latter defaults to 1.1 to allow for failures 
    in showering stage)."""
    try:
        sf = float(safety)
    except (TypeError, ValueError) as exc:
        raise RuntimeError(f"safety must be numeric, got {safety}.") from exc
    if sf <= 0:
        raise RuntimeError(f"safety must be > 0, got {safety}.")
 
    base_events = (
        flags.Exec.MaxEvents
        if flags.Exec.MaxEvents > 0
        else flags.Generator.nEventsPerJob
    )
 
    return int(base_events * sf)
 
 
def _prepare_lhe_for_shower(produced_output, lhe_file):
    primary_output = None
    if produced_output:
        if produced_output.endswith(".tar.gz"):
            root = produced_output[:-7]
        elif produced_output.endswith(".tgz"):
            root = produced_output[:-4]
        elif produced_output.endswith(".gz"):
            root = produced_output[:-3]
        else:
            root, _ = os.path.splitext(produced_output)
        primary_output = f"{root}.events"
 
    # If the transform requested a specific TXT output name, symlink the 
    # produced output to the filename that the transform expects
    candidates = [candidate 
                  for candidate in (primary_output, 
                                    "tmp_LHE_events.events", 
                                    "events.events") 
                  if candidate]
    if _symlink_first_existing(lhe_file, candidates, overwrite=True):
        return
 
    raise RuntimeError(
        "Could not prepare LHE file for showering. "
        f"Expected one of: {', '.join(candidates)}"
    )
 
 
def _symlink_first_existing(link_name, candidates, overwrite=False):
    """
    Helper function to symlink the first existing file in candidates to link_name.
    """
    if os.path.exists(link_name) and not overwrite:
        return True
 
    for candidate in candidates:
        if not candidate or not os.path.exists(candidate):
            continue
        if os.path.abspath(candidate) == os.path.abspath(link_name):
            return True
        if os.path.lexists(link_name):
            os.remove(link_name)
        os.symlink(os.path.abspath(candidate), link_name)
        return True
 
    return False
 
 
def MadGraphBaseCfg(flags, **kwargs):
    """Base MadGraph CA fragment. It returns a CA object 
    that contains the generator metadata and registers 
    default values for steering the MGC object 
    (to be created by the top-level config)."""
    from MadGraphControl.MGC import (
        MADGRAPH_CATCH_ERRORS,
        MADGRAPH_DEVICES,
        MADGRAPH_PDFSETTING,
    )
 
    # Default values for MGC. Use MGC defaults for now, but these
    # can be declared here in the future.
    defaults = {
        "safety": 1.1,
        "pdf_setting": MADGRAPH_PDFSETTING,
        "devices": MADGRAPH_DEVICES,
        "catch_errors": MADGRAPH_CATCH_ERRORS,
        "lhe_version": 3,
        "saveProcDir": False,
        "keepJpegs": False,
        "usePMGSettings": False,
    }
 
    # Create a dictionary with default settings.
    # This can be used in top-level configs.
    cfg = {**defaults, **{k: v for k, v in kwargs.items() if v is not None}}
 
    # Create the CA object adding the generator metadata
    ca = ComponentAccumulator(EvgenSequenceFactory(EvgenSequence.Generator))
    ca.merge(
        GeneratorInfoSvcCfg(flags, Generators=["MadGraph"]),
        sequenceName=EvgenSequence.Generator.value,
    )
    
    return ca, cfg
 
 
def MadGraphCfg(
    flags,
    process_definition,
    *,
    safety=None,
    settings=None,
    pdf_setting=None,
    devices=None,
    catch_errors=None,
    lhe_version=None,
    saveProcDir=None,
    plugin=None,
    keepJpegs=None,
    usePMGSettings=None,
    prepare_lhe_for_shower=False,
    lhe_file="events.lhe",
):
    """
    Fragment for configuring a LHE generation step.
 
    This starts from MadGraphBaseCfg and creates a MGC instance
    that is later used to call the MadGraphUtils functions that steer
    the event generation.
 
    All arguments after * are keyword-only to avoid confusion 
    between MadGraphControl settings and CA configuration options.
    Set prepare_lhe_for_shower=True when the same job should feed the
    produced LHE file into a shower generator.
 
    process_definition is required, the rest is optional.
 
    If prepare_lhe_for_shower is True, the produced LHE file will be 
    symlinked to lhe_file (default: events.lhe) 
    for later use in the showering step.
    """
 
    from MadGraphControl.MGC import MGControl
    import MadGraphControl.MadGraphUtils as MadGraphUtils
 
    if isinstance(pdf_setting, MadGraphPDFSets):
        pdf_setting = get_pdf_set(pdf_setting)
 
    # TODO: implement deduplication of settings as done in Pythia8Config
    ca, cfg = MadGraphBaseCfg(
        flags,
        safety=safety,
        pdf_setting=pdf_setting,
        devices=devices,
        catch_errors=catch_errors,
        lhe_version=lhe_version,
        saveProcDir=saveProcDir,
        keepJpegs=keepJpegs,
        usePMGSettings=usePMGSettings,
    )
 
    run_card_settings = {} if settings is None else dict(settings)
 
    # Get nEvents
    run_card_settings["nevents"] = _get_nevents(flags, cfg["safety"])
 
    # Create the MGC instance
    mgc = MGControl(
        process=process_definition,
        plugin=plugin,
        keepJpegs=cfg["keepJpegs"],
        usePMGSettings=cfg["usePMGSettings"],
        pdf_setting=cfg["pdf_setting"],
        devices=cfg["devices"],
        catch_errors=cfg["catch_errors"],
    )
 
    # Bind the MGC instance to the MadGraphUtils module, 
    # so that it can be used in the calls to the MadGraphUtils functions 
    # This is not following the CA logic completely, but it avoids 
    # having to pass the MGC instance through multiple function calls.
    MadGraphUtils.my_MGC_instance = mgc
 
    # Create the process directory
    process_dir = mgc.process_dir
    MadGraphUtils.modify_run_card(
        process_dir=process_dir,
        flags=flags,
        settings=run_card_settings,
        pdf_setting=cfg["pdf_setting"],
    )
 
    # Generate events
    MadGraphUtils.generate(process_dir=process_dir, flags=flags, pdf_setting=cfg["pdf_setting"])
    produced_output = MadGraphUtils.arrange_output(
        process_dir=process_dir,
        flags=flags,
        lhe_version=cfg["lhe_version"],
        saveProcDir=cfg["saveProcDir"],
        pdf_setting=cfg["pdf_setting"],
    )
 
    # If requested, prepare the produced LHE file for showering
    # by symlinking it to the filename that pythia expects, 
    # by default "events.lhe".
    if prepare_lhe_for_shower:
        _prepare_lhe_for_shower(produced_output, lhe_file)
 
    # If the transform requested a specific TXT output name, symlink the 
    # produced output to the filename that the transform expects
    # (only if the file does not exist).
    requested_output = flags.Output.TXTFileName
    if requested_output and not os.path.exists(requested_output):
        root, _ = os.path.splitext(requested_output)
        candidates = [candidate for candidate in (produced_output, f"{root}.events", "events.events") if candidate]
        _symlink_first_existing(requested_output, candidates, overwrite=True)
 
    return ca