Add MadtRex tests to CI

.github/workflows/generate_madtrex_assets.py

@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+"""Script to (re)generate the MadtRex CI test assets for one process:
+ - an LHE event sample (test/MadtRex_baseline/<process>_events.lhe.gz)
+ - the corresponding reweight baseline (test/MadtRex_baseline/<process>_rwgt.csv)
+
+This is NOT run in CI. The idea is to run it locally every now and then
+(whenever the generated *.mad code or something crucial changes) and commit the
+resulting files, the same way epochX/cudacpp/CODEGEN/allGenerateAndCompare.sh
+is used to regenerate reference outputs.
+
+Usage (from epochX/cudacpp):
+  ../../.github/workflows/generate_madtrex_assets.py gg_tt.mad
+"""
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+import madtrex_common as common
+
+MADGRAPH_CLI = Path.cwd() / ".." / ".." / "MG5aMC" / "mg5amcnlo" / "bin" / "mg5_aMC"
+
+def generate_dat_content(process_dir: Path, rwgt_card_path: Path, process: str) -> str:
+    dat = common.model_import_lines(process)
+    dat += f"launch {process_dir}\n"
+    dat += "reweight=madtrex\n"
+    dat += "set nevents 10000\n"
+    dat += f"set iseed {common.ISEED}\n"
+    dat += f"{rwgt_card_path}\n"
+    return dat
+
+def main() -> int:
+    process_dir = Path(sys.argv[1])
+    process = process_dir.name.replace(".mad", "")
+
+    HOME = Path.cwd()
+    process_path = (HOME / process_dir).resolve()
+    if not process_path.exists():
+        print(f"ERROR: Process {process} not found at: {process_path}", file=sys.stderr)
+        return 1
+
+    if process not in common.ALLOWED_PROCESSES:
+        print(
+            f"ERROR: PROCESS '{process}' is not in the allowed list.\n"
+            f"Allowed: {sorted(common.ALLOWED_PROCESSES)}",
+            file=sys.stderr,
+        )
+        return 1
+
+    rwgt_card_path = HOME / "rwgt_card.dat"
+    common.write_rwgt_card(rwgt_card_path)
+
+    dat_path = HOME / f"{process}.dat"
+    dat_path.write_text(generate_dat_content(process_dir, rwgt_card_path, process), encoding="utf-8")
+
+    # Check that the CUDACPP_OUTPUT plugin is present: required for MadtRex reweighting
+    error = common.check_cudacpp_plugin_present(HOME)
+    if error:
+        print(error, file=sys.stderr)
+        return 1
+
+    LOGS = HOME / "logs"
+    LOGS.mkdir(exist_ok=True)
+    stdout_log = LOGS / f"mg5_{process}.stdout.log"
+    stderr_log = LOGS / f"mg5_{process}.stderr.log"
+    print(f"Launching: {MADGRAPH_CLI} {dat_path}")
+    with stdout_log.open("wb") as out, stderr_log.open("wb") as err:
+        result = subprocess.run(
+            [str(MADGRAPH_CLI), str(dat_path)],
+            cwd=str(HOME),
+            stdout=out,
+            stderr=err,
+            check=False,
+        )
+    if result.returncode != 0:
+        print(f"ERROR: mg5_aMC exited with code {result.returncode}.", file=sys.stderr)
+        common.dump_logs(stdout_log, stderr_log)
+        return result.returncode
+    print(f"mg5_aMC finished. Logs:\n  stdout: {stdout_log}\n  stderr: {stderr_log}")
+
+    dat_path.unlink(missing_ok=True)
+
+    # Locate the generated run's LHE file (the only run in a freshly generated process dir)
+    run_dirs = sorted((process_path / "Events").glob("run_*"))
+    if not run_dirs:
+        print(f"ERROR: No Events/run_* directory found under {process_path}", file=sys.stderr)
+        common.dump_logs(stdout_log, stderr_log)
+        return 1
+    if len(run_dirs) > 1:
+        print(f"WARNING: Multiple runs found, using the last one: {[d.name for d in run_dirs]}")
+    run_dir = run_dirs[-1]
+    lhe_src = run_dir / "unweighted_events.lhe.gz"
+    if not lhe_src.exists():
+        print(f"ERROR: Expected LHE file not found at: {lhe_src}", file=sys.stderr)
+        common.dump_logs(stdout_log, stderr_log)
+        return 1
+
+    madtrex_csv = process_path / "rw_me" / "SubProcesses" / "rwgt_results.csv"
+    if not madtrex_csv.exists():
+        print(f"ERROR: Expected results not found at: {madtrex_csv}", file=sys.stderr)
+        common.dump_logs(stdout_log, stderr_log)
+        return 1
+
+    lhe_dst = common.baseline_lhe_path(HOME, process)
+    csv_dst = common.baseline_csv_path(HOME, process)
+    lhe_dst.parent.mkdir(parents=True, exist_ok=True)
+    shutil.copyfile(lhe_src, lhe_dst)
+    shutil.copyfile(madtrex_csv, csv_dst)
+    print(f"Updated asset: {lhe_dst}")
+    print(f"Updated asset: {csv_dst}")
+    print("Please review and commit these files.")
+
+    return 0
+
+if __name__ == "__main__":
+    sys.exit(main())

.github/workflows/madtrex_common.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+"""Shared constants/helpers for the MadtRex CI test (run_madtrex.py) and the
+dev-only asset regeneration script (generate_madtrex_assets.py)."""
+import csv
+import os
+import re
+from pathlib import Path
+
+ALLOWED_PROCESSES = [ "ee_mumu", "gg_tt", "gg_tt01g", "gg_ttg", "gg_ttgg", "gg_ttggg", "gq_ttq", "heft_gg_bb", "nobm_pp_ttW", "pp_tt012j", "smeft_gg_tttt", "susy_gg_t1t1", "susy_gg_tt" ]
+
+# Run name used for the (real or restored) event sample: must match MadGraph's
+# default first-run naming, since the asset is restored without a results database.
+RUN_NAME = "run_01"
+
+ISEED = 489
+
+BASELINE_DIR = Path("CODEGEN") / "PLUGIN" / "CUDACPP_SA_OUTPUT" / "test" / "MadtRex_baseline"
+
+def baseline_csv_path(home: Path, process: str) -> Path:
+    return home / BASELINE_DIR / f"{process}_rwgt.csv"
+
+def baseline_lhe_path(home: Path, process: str) -> Path:
+    return home / BASELINE_DIR / f"{process}_events.lhe.gz"
+
+def is_executable(path: Path) -> bool:
+    return path.is_file() and os.access(path, os.X_OK)
+
+def mg5amcnlo_dir(home: Path) -> Path:
+    """Location of the MG5aMC/mg5amcnlo checkout relative to HOME (epochX/cudacpp)."""
+    return home / ".." / ".." / "MG5aMC" / "mg5amcnlo"
+
+def check_cudacpp_plugin_present(home: Path) -> str:
+    """Check that PLUGIN/CUDACPP_OUTPUT exists (directory or symlink to one) inside
+    the mg5amcnlo checkout: it is required for MadtRex reweighting to work.
+    Returns an error message if missing, or an empty string if the check passes."""
+    plugin_dir = mg5amcnlo_dir(home) / "PLUGIN" / "CUDACPP_OUTPUT"
+    if not plugin_dir.is_dir():
+        return (
+            f"ERROR: CUDACPP_OUTPUT plugin not found at:\n  {plugin_dir}\n"
+            f"It is required for MadtRex reweighting. Create it, e.g. with:\n"
+            f"  cd {plugin_dir.parent} && ln -s ../../MG5aMC_PLUGIN/CUDACPP_OUTPUT ./"
+        )
+    return ""
+
+def set_mg5_path(me5_configuration_path: Path, mg5_path: Path) -> None:
+    """Set (uncomment/overwrite) the mg5_path entry in me5_configuration.txt so that
+    bin/madevent can find the mg5amcnlo checkout needed for MadtRex reweighting,
+    without going through bin/mg5_aMC's 'launch' command."""
+    text = me5_configuration_path.read_text(encoding="utf-8")
+    text = re.sub(r"^#?\s*mg5_path\s*=.*$", f"mg5_path = {mg5_path}", text, flags=re.MULTILINE)
+    me5_configuration_path.write_text(text, encoding="utf-8")
+
+def write_rwgt_card(path: Path) -> None:
+    if path.exists():
+        return
+    content = """launch\nset sminputs 1 scan:[j for j in range(100,200,10)]\n"""
+    path.write_text(content, encoding="utf-8")
+
+def load_csv(path):
+    with open(path, newline="") as f:
+        reader = csv.DictReader(f, fieldnames=["RWGT", "VALUE", "ERROR"])
+        for row in reader:
+            yield float(row["VALUE"]), float(row["ERROR"])
+
+def dump_logs(stdout_log, stderr_log):
+    print("Dumping run logs...")
+    print("==== STDOUT ====")
+    with open(stdout_log, "r") as file:
+        print(file.read())
+    print("\n\n==== STDERR ====")
+    with open(stderr_log, "r") as file:
+        print(file.read())
+    print("================")
+
+def compare_csv(baseline_csv: Path, madtrex_csv: Path) -> bool:
+    all_ok = True
+    for i, ((v_base, _), (v_mad, _)) in enumerate(zip(load_csv(baseline_csv), load_csv(madtrex_csv)), start=1):
+        diff = abs(v_base - v_mad)
+        tol = 0.05 * v_mad
+        if diff >= tol:
+            print(f"Error: Row {i}: |{v_base} - {v_mad}| = {diff} >= {tol}")
+            all_ok = False
+    return all_ok

.github/workflows/run_madtrex.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+import madtrex_common as common
+
+def generate_dat_content(process: str) -> str:
+    dat = f"reweight {common.RUN_NAME} -from_cards --plugin=madtrex\n"
+    return dat
+
+def restore_lhe(lhe_asset: Path, process_path: Path) -> Path:
+    """Decompress the committed LHE asset into Events/<run_name>/unweighted_events.lhe.gz,
+    the location MadGraph expects for an existing run (no run database needed: the banner
+    is read straight from the LHE file itself)."""
+    events_dir = process_path / "Events" / common.RUN_NAME
+    events_dir.mkdir(parents=True, exist_ok=True)
+    lhe_dst = events_dir / "unweighted_events.lhe.gz"
+    shutil.copyfile(lhe_asset, lhe_dst)
+    return lhe_dst
+
+def main() -> int:
+    # Name of the directory of the process to test
+    process_dir = Path(sys.argv[1])
+    # Label for the process (must be in the allowed list)
+    process = process_dir.name.replace(".mad", "")
+
+    # Treat current working directory as HOME
+    HOME = Path.cwd()
+
+    process_path = (HOME / process_dir).resolve()
+    if not process_path.exists():
+        print(f"ERROR: Process {process} not found at: {process_path}", file=sys.stderr)
+        return 1
+
+    if process not in common.ALLOWED_PROCESSES:
+        print(
+            f"ERROR: PROCESS '{process}' is not in the allowed list.\n"
+            f"Allowed: {sorted(common.ALLOWED_PROCESSES)}",
+            file=sys.stderr,
+        )
+        return 1
+
+    # Check that baseline rwgt.csv exists
+    baseline_csv = common.baseline_csv_path(HOME, process)
+    if not baseline_csv.exists():
+        print(
+            f"ERROR: Baseline rwgt.csv not found at:\n  {baseline_csv}\n"
+            f"Ensure the baseline file exists before running.",
+            file=sys.stderr,
+        )
+        return 1
+
+    # Check that the pre-generated LHE asset exists, and restore it into the run directory
+    lhe_asset = common.baseline_lhe_path(HOME, process)
+    if not lhe_asset.exists():
+        print(
+            f"ERROR: LHE asset not found at:\n  {lhe_asset}\n"
+            f"Generate it with generate_madtrex_assets.py and commit it before running.",
+            file=sys.stderr,
+        )
+        return 1
+    restore_lhe(lhe_asset, process_path)
+
+    # Write reweight_card.dat directly into the process Cards directory (consumed via -from_cards)
+    rwgt_card_path = process_path / "Cards" / "reweight_card.dat"
+    common.write_rwgt_card(rwgt_card_path)
+
+    # Write PROCESS.dat to HOME
+    dat_path = HOME / f"{process}.dat"
+    dat_path.write_text(generate_dat_content(process), encoding="utf-8")
+
+    # Check that the CUDACPP_OUTPUT plugin is present: required for MadtRex reweighting
+    error = common.check_cudacpp_plugin_present(HOME)
+    if error:
+        print(error, file=sys.stderr)
+        return 1
+
+    # Point me5_configuration.txt at this repo's mg5amcnlo checkout: required for
+    # MadtRex reweighting, which uses it directly instead of going through bin/mg5_aMC
+    common.set_mg5_path(
+        process_path / "Cards" / "me5_configuration.txt",
+        common.mg5amcnlo_dir(HOME).resolve(),
+    )
+
+    # Run bin/madevent with PROCESS.dat as argument, wait for completion
+    madevent_bin = process_path / "bin" / "madevent"
+    LOGS = HOME / "logs"
+    LOGS.mkdir(exist_ok=True)
+    stdout_log = LOGS / f"mg5_{process}.stdout.log"
+    stderr_log = LOGS / f"mg5_{process}.stderr.log"
+    print(f"Launching: {madevent_bin} {dat_path}")
+    try:
+        with stdout_log.open("wb") as out, stderr_log.open("wb") as err:
+            result = subprocess.run(
+                [str(madevent_bin), str(dat_path)],
+                cwd=str(process_path),
+                stdout=out,
+                stderr=err,
+                check=False,
+            )
+        if result.returncode != 0:
+            print(
+                f"ERROR: bin/madevent exited with code {result.returncode}. "
+                f"See logs:\n  stdout: {stdout_log}\n  stderr: {stderr_log}",
+                file=sys.stderr,
+            )
+            common.dump_logs(stdout_log, stderr_log)
+            return result.returncode
+        else:
+            print(f"bin/madevent finished. Logs:\n  stdout: {stdout_log}\n  stderr: {stderr_log}")
+    except FileNotFoundError:
+        print(f"ERROR: Failed to launch {madevent_bin}", file=sys.stderr)
+        return 1
+    except Exception as e:
+        print(f"ERROR: bin/madevent run failed: {e}", file=sys.stderr)
+        return 1
+
+    # Remove process.dat
+    dat_path.unlink(missing_ok=True)
+
+    # Get rwgt_results.csv results file
+    madtrex_csv = process_path / "rw_me" / "SubProcesses" / "rwgt_results.csv"
+    if not madtrex_csv.exists():
+        print(
+            f"ERROR: Expected results not found at:\n  {madtrex_csv}\n"
+            f"Ensure the run produced rwgt_results.csv.",
+            file=sys.stderr,
+        )
+        common.dump_logs(stdout_log, stderr_log)
+        return 1
+
+    if not common.compare_csv(baseline_csv, madtrex_csv):
+        print(f"Some checks failed for process {process}.", file=sys.stderr)
+        sys.exit(1)
+    print(f"All checks passed for process {process}.")
+
+    return 0
+
+if __name__ == "__main__":
+    sys.exit(main())